/**
 * Level 2: Shadow + Shadow Manager + Tool Registry
 * 
 * This adds:
 * 1. Shadow class with context isolation
 * 2. Shadow Manager for spawning/terminating
 * 3. Tool registry (read, write, edit, bash, grep, find, ls)
 * 4. Basic concurrency control
 */

import { Agent, type AgentTool, type AgentMessage, type AgentEvent } from "@mariozechner/pi-agent-core";
import { registerBuiltInApiProviders } from "@mariozechner/pi-ai";
import type { Model } from "@mariozechner/pi-ai";
import * as fs from "fs";
import * as path from "path";
import { exec, spawn } from "child_process";

// ============== CONFIG ==============
const OPENROUTER_API_KEY = process.env.OPENROUTER_API_KEY || "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";
process.env.OPENROUTER_API_KEY = OPENROUTER_API_KEY;

// Model config (using free stepfun model)
const model: Model<"openai-responses"> = {
  id: "stepfun/step-3.5-flash:free",
  name: "Step-3.5 Flash (Free)",
  api: "openai-responses",
  provider: "openrouter",
  baseUrl: "https://openrouter.ai/api/v1",
  reasoning: false,
  input: ["text"],
  output: ["text"],
  cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
  contextWindow: 128000,
  maxTokens: 8192,
};

// ============== TOOL REGISTRY ==============
function createTools(cwd: string = process.cwd()): AgentTool[] {
  return [
    {
      name: "read",
      label: "Read File",
      description: "Read the contents of a file",
      parameters: {
        type: "object",
        properties: {
          path: { type: "string", description: "Path to the file to read" },
        },
        required: ["path"],
      } as const,
      execute: async (toolCallId: string, params: { path: string }) => {
        const fullPath = path.resolve(cwd, params.path);
        try {
          if (!fs.existsSync(fullPath)) {
            throw new Error(`File not found: ${fullPath}`);
          }
          const content = fs.readFileSync(fullPath, "utf-8");
          return {
            content: [{ type: "text", text: content }],
            details: { path: fullPath, lines: content.split("\n").length },
          };
        } catch (error: any) {
          throw new Error(`Failed to read file: ${error.message}`);
        }
      },
    },
    {
      name: "write",
      label: "Write File",
      description: "Write content to a file (creates or overwrites)",
      parameters: {
        type: "object",
        properties: {
          path: { type: "string", description: "Path to the file to write" },
          content: { type: "string", description: "Content to write" },
        },
        required: ["path", "content"],
      } as const,
      execute: async (toolCallId: string, params: { path: string; content: string }) => {
        const fullPath = path.resolve(cwd, params.path);
        try {
          // Ensure directory exists
          const dir = path.dirname(fullPath);
          if (!fs.existsSync(dir)) {
            fs.mkdirSync(dir, { recursive: true });
          }
          fs.writeFileSync(fullPath, params.content, "utf-8");
          return {
            content: [{ type: "text", text: `Written ${params.content.length} bytes to ${fullPath}` }],
            details: { path: fullPath, bytes: params.content.length },
          };
        } catch (error: any) {
          throw new Error(`Failed to write file: ${error.message}`);
        }
      },
    },
    {
      name: "edit",
      label: "Edit File",
      description: "Edit a file by replacing specific text",
      parameters: {
        type: "object",
        properties: {
          path: { type: "string", description: "Path to the file to edit" },
          find: { type: "string", description: "Text to find" },
          replace: { type: "string", description: "Text to replace with" },
        },
        required: ["path", "find"],
      } as const,
      execute: async (toolCallId: string, params: { path: string; find: string; replace?: string }) => {
        const fullPath = path.resolve(cwd, params.path);
        try {
          if (!fs.existsSync(fullPath)) {
            throw new Error(`File not found: ${fullPath}`);
          }
          let content = fs.readFileSync(fullPath, "utf-8");
          const newContent = params.replace !== undefined 
            ? content.replace(params.find, params.replace)
            : content.replace(params.find, "");
          
          if (content === newContent) {
            throw new Error(`Text not found: "${params.find}"`);
          }
          
          fs.writeFileSync(fullPath, newContent, "utf-8");
          return {
            content: [{ type: "text", text: `Edited ${fullPath}` }],
            details: { path: fullPath },
          };
        } catch (error: any) {
          throw new Error(`Failed to edit file: ${error.message}`);
        }
      },
    },
    {
      name: "bash",
      label: "Run Command",
      description: "Run a shell command",
      parameters: {
        type: "object",
        properties: {
          command: { type: "string", description: "Command to run" },
        },
        required: ["command"],
      } as const,
      execute: async (toolCallId: string, params: { command: string }) => {
        return new Promise((resolve) => {
          exec(params.command, { cwd }, (error, stdout, stderr) => {
            if (error) {
              resolve({
                content: [{ type: "text", text: stderr || error.message }],
                details: { command: params.command, exitCode: error.code },
                isError: true,
              });
            } else {
              resolve({
                content: [{ type: "text", text: stdout }],
                details: { command: params.command, exitCode: 0 },
              });
            }
          });
        });
      },
    },
    {
      name: "grep",
      label: "Search Text",
      description: "Search for text in files",
      parameters: {
        type: "object",
        properties: {
          pattern: { type: "string", description: "Pattern to search for" },
          path: { type: "string", description: "Path to search in (file or directory)" },
        },
        required: ["pattern"],
      } as const,
      execute: async (toolCallId: string, params: { pattern: string; path?: string }) => {
        const searchPath = params.path || cwd;
        return new Promise((resolve) => {
          exec(`grep -r "${params.pattern}" ${searchPath} --line-number 2>/dev/null || true`, { cwd }, (error, stdout) => {
            resolve({
              content: [{ type: "text", text: stdout || `No matches found for "${params.pattern}"` }],
              details: { pattern: params.pattern, path: searchPath },
            });
          });
        });
      },
    },
    {
      name: "ls",
      label: "List Files",
      description: "List files in a directory",
      parameters: {
        type: "object",
        properties: {
          path: { type: "string", description: "Directory to list" },
        },
      } as const,
      execute: async (toolCallId: string, params: { path?: string }) => {
        const listPath = params.path ? path.resolve(cwd, params.path) : cwd;
        try {
          const files = fs.readdirSync(listPath);
          return {
            content: [{ type: "text", text: files.join("\n") }],
            details: { path: listPath, count: files.length },
          };
        } catch (error: any) {
          throw new Error(`Failed to list: ${error.message}`);
        }
      },
    },
  ];
}

// ============== SHADOW CLASS ==============
interface ShadowConfig {
  id: string;
  systemPrompt: string;
  worktreePath: string;
  modelId?: string;
}

interface ShadowState {
  id: string;
  status: "idle" | "running" | "completed" | "error";
  createdAt: Date;
  worktreePath: string;
}

class Shadow {
  public readonly id: string;
  public readonly agent: Agent;
  public readonly worktreePath: string;
  public state: ShadowState;
  
  private eventCallback?: (event: AgentEvent) => void;

  constructor(config: ShadowConfig) {
    this.id = config.id;
    this.worktreePath = config.worktreePath;
    this.state = {
      id: config.id,
      status: "idle",
      createdAt: new Date(),
      worktreePath: config.worktreePath,
    };

    // Create Pi Agent with isolated context
    this.agent = new Agent({
      initialState: {
        systemPrompt: config.systemPrompt,
        model: model,
        tools: createTools(config.worktreePath) as any,
        messages: [],
      },
      convertToLlm: (messages: AgentMessage[]) => {
        // ISOLATION: Filter to only this shadow's messages
        // We add a special role suffix to identify messages from this shadow
        return messages
          .filter((m) => {
            // Keep messages that either:
            // 1. Have no shadowId (legacy) OR
            // 2. Have matching shadowId
            const msg = m as any;
            return !msg._shadowId || msg._shadowId === this.id;
          })
          .filter((m) => m.role === "user" || m.role === "assistant" || m.role === "toolResult")
          .map((m) => ({
            role: m.role,
            content: m.content,
          }));
      },
    });

    // Subscribe to events
    this.agent.subscribe((event) => {
      // Track state changes
      if (event.type === "agent_start") {
        this.state.status = "running";
      } else if (event.type === "agent_end") {
        this.state.status = "completed";
      } else if (event.type === "tool_execution_start") {
        // Tool running
      } else if (event.type === "tool_execution_end" && (event as any).isError) {
        this.state.status = "error";
      }
      
      // Forward events
      this.eventCallback?.(event);
    });
  }

  onEvent(callback: (event: AgentEvent) => void) {
    this.eventCallback = callback;
  }

  async prompt(message: string): Promise<AgentEvent[]> {
    const events: AgentEvent[] = [];
    
    // Tag message with shadow ID for isolation
    const shadowMessage: AgentMessage = {
      role: "user",
      content: [{ type: "text", text: message }],
      timestamp: Date.now(),
      _shadowId: this.id,  // Our custom field for isolation
    };
    
    return this.agent.prompt(shadowMessage);
  }

  abort() {
    this.agent.abort();
  }

  reset() {
    this.agent.reset();
    this.state.status = "idle";
  }
}

// ============== SHADOW MANAGER ==============
interface ShadowManagerConfig {
  maxConcurrent?: number;
  defaultSystemPrompt?: string;
}

class ShadowManager {
  private shadows: Map<string, Shadow> = new Map();
  private maxConcurrent: number;
  private defaultSystemPrompt: string;
  private activeCount = 0;

  constructor(config: ShadowManagerConfig = {}) {
    this.maxConcurrent = config.maxConcurrent || 5;
    this.defaultSystemPrompt = config.defaultSystemPrompt || `You are a helpful coding assistant. You have access to tools: read, write, edit, bash, grep, ls. Use them to help the user. Be concise and practical.`;
  }

  async createShadow(worktreePath: string, customPrompt?: string): Promise<Shadow> {
    // Check concurrency limit
    if (this.activeCount >= this.maxConcurrent) {
      throw new Error(`Max concurrent shadows reached (${this.maxConcurrent})`);
    }

    const id = `shadow-${Date.now()}-${Math.random().toString(36).substr(2, 9)}`;
    
    const shadow = new Shadow({
      id,
      systemPrompt: customPrompt || this.defaultSystemPrompt,
      worktreePath,
    });

    this.shadows.set(id, shadow);
    this.activeCount++;
    
    console.log(`📦 Created shadow ${id} (active: ${this.activeCount}/${this.maxConcurrent})`);
    
    return shadow;
  }

  getShadow(id: string): Shadow | undefined {
    return this.shadows.get(id);
  }

  listShadows(): ShadowState[] {
    return Array.from(this.shadows.values()).map((s) => s.state);
  }

  async terminateShadow(id: string): Promise<void> {
    const shadow = this.shadows.get(id);
    if (!shadow) {
      throw new Error(`Shadow ${id} not found`);
    }

    shadow.abort();
    this.shadows.delete(id);
    this.activeCount--;
    
    console.log(`🗑️ Terminated shadow ${id} (active: ${this.activeCount}/${this.maxConcurrent})`);
  }

  getStats() {
    return {
      active: this.activeCount,
      maxConcurrent: this.maxConcurrent,
      totalShadows: this.shadows.size,
      shadows: this.listShadows(),
    };
  }
}

// ============== MAIN ==============
async function main() {
  console.log("🚀 Level 2: Shadow + Shadow Manager\n");

  // Initialize
  registerBuiltInApiProviders();

  // Create manager
  const manager = new ShadowManager({
    maxConcurrent: 3,
  });

  // Create a shadow
  console.log("📦 Creating shadow...");
  const shadow = await manager.createShadow("/home/shoko/repositories/shadows");
  
  // Subscribe to events
  shadow.onEvent((event) => {
    switch (event.type) {
      case "agent_start":
        console.log("🤖 Agent started");
        break;
      case "turn_start":
        console.log("🔄 Turn started");
        break;
      case "message_update":
        const ev = event as any;
        if (ev.assistantMessageEvent?.type === "text_delta") {
          process.stdout.write(ev.assistantMessageEvent.delta || "");
        }
        break;
      case "tool_execution_start":
        console.log(`\n🔧 Tool: ${event.toolName}`);
        break;
      case "tool_execution_end":
        console.log(`   → Done (error: ${(event as any).isError})`);
        break;
      case "turn_end":
        console.log("\n✅ Turn ended");
        break;
      case "agent_end":
        console.log("\n🏁 Agent finished");
        break;
    }
  });

  // Run a task
  console.log("\n📝 Running task: List files and check current directory\n");
  await shadow.prompt("List the files in the current directory, then run 'pwd' to check the current directory.");

  // Show stats
  console.log("\n📊 Manager Stats:", manager.getStats());

  // Cleanup
  await manager.terminateShadow(shadow.id);
  console.log("\n✅ Done!");
}

main().catch(console.error);