/**
 * Cursor ACP Extension for pi
 * 
 * Enables pi to use Cursor's AI agent via the Agent Client Protocol.
 * Sessions are created fresh for each request and cleaned up after.
 */

import { spawn, ChildProcess } from "node:child_process";
import readline from "node:readline";
import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
import { Type } from "@sinclair/typebox";

// ============================================================================
// Types
// ============================================================================

interface JsonRpcMessage {
  jsonrpc: "2.0";
  id?: number;
  method?: string;
  params?: unknown;
  result?: unknown;
  error?: { code: number; message: string; data?: unknown };
}

interface SessionUpdate {
  sessionUpdate: string;
  content?: { text?: string; type?: string };
  task?: { id: string; status: string };
}

interface PermissionRequest {
  toolName: string;
  toolArgs: Record<string, unknown>;
  justification?: string;
}

// ============================================================================
// ACP Session (created per-request)
// ============================================================================

interface ACPResponse {
  stopReason: string;
}

class ACPSession {
  private agent: ChildProcess | null = null;
  private pending = new Map<number, { resolve: (v: unknown) => void; reject: (e: Error) => void }>();
  private nextId = 1;
  private sessionId: string | null = null;
  private streamedContent = "";
  private stopReason = "unknown";
  private updateCallback?: (content: string) => void;
  private permissionHandler?: (request: PermissionRequest) => Promise<string>;

  /**
   * Run a complete request: spawn, init, auth, session, prompt, cleanup
   */
  async run(text: string, mode: "agent" | "plan" | "ask" = "agent", cwd?: string): Promise<string> {
    this.streamedContent = "";
    this.stopReason = "unknown";
    
    await this.start();
    
    try {
      await this.createSession(cwd);
      
      // This waits for the response, while notifications populate streamedContent
      const result = (await this.send("session/prompt", {
        sessionId: this.sessionId!,
        prompt: [{ type: "text", text }],
        mode,
      })) as ACPResponse;
      
      this.stopReason = result.stopReason;
    } finally {
      this.stop();
    }
    
    return this.streamedContent || `[Completed: ${this.stopReason}]`;
  }

  private sessionId: string | null = null;

  private async start(): Promise<void> {
    return new Promise((resolve, reject) => {
      this.agent = spawn("agent", ["acp"], {
        stdio: ["pipe", "pipe", "pipe"],
        env: { ...process.env },
      });

      this.agent.on("error", (err) => reject(err));

      const rl = readline.createInterface({
        input: this.agent.stdout!,
        crlfDelay: Infinity,
      });

      rl.on("line", (line) => {
        try {
          this.handleMessage(JSON.parse(line));
        } catch (e) {
          // Ignore non-JSON lines
        }
      });

      const timeout = setTimeout(() => reject(new Error("Init timeout (60s)")), 60000);

      this.initialize()
        .then(() => this.authenticate())
        .then(() => {
          clearTimeout(timeout);
          resolve();
        })
        .catch((err) => {
          clearTimeout(timeout);
          reject(err);
        });
    });
  }

  private send(method: string, params?: unknown): Promise<unknown> {
    return new Promise((resolve, reject) => {
      if (!this.agent || !this.agent.stdin) {
        reject(new Error("Agent not running"));
        return;
      }

      const id = this.nextId++;
      this.pending.set(id, { resolve, reject });

      this.agent.stdin.write(JSON.stringify({ jsonrpc: "2.0", id, method, params }) + "\n");

      setTimeout(() => {
        if (this.pending.has(id)) {
          this.pending.delete(id);
          reject(new Error(`Request ${method} timed out (120s)`));
        }
      }, 120000);
    });
  }

  private handleMessage(msg: JsonRpcMessage): void {
    if (msg.id !== undefined && this.pending.has(msg.id)) {
      const waiter = this.pending.get(msg.id)!;
      this.pending.delete(msg.id);
      if (msg.error) waiter.reject(new Error(msg.error.message));
      else waiter.resolve(msg.result);
      return;
    }

    if (msg.method) {
      this.handleNotification(msg.method, msg.params);
    }
  }

  private handleNotification(method: string, params?: unknown): void {
    if (method === "session/update") {
      const update = (params as { update: SessionUpdate })?.update;
      if (update) {
        if (update.sessionUpdate === "agent_message_chunk" && update.content?.text) {
          this.streamedContent += update.content.text;
          this.updateCallback?.(this.streamedContent);
        }
        if (update.sessionUpdate === "agent_thought_chunk" && update.content?.text) {
          // Optionally include thoughts
        }
      }
    }
    
    if (method === "session/request_permission" && this.permissionHandler) {
      const req = params as { id: number; params: PermissionRequest };
      if (req?.id) {
        this.permissionHandler(req.params).then((outcome) => {
          this.agent?.stdin?.write(JSON.stringify({
            jsonrpc: "2.0",
            id: req.id,
            result: { outcome: "selected", optionId: outcome }
          }) + "\n");
        });
      }
    }
  }

  private async initialize(): Promise<void> {
    const result = (await this.send("initialize", {
      protocolVersion: 1,
      clientCapabilities: {},
      clientInfo: { name: "pi-cursor", version: "1.0.0" },
    })) as { protocolVersion: number };
    
    if (result.protocolVersion !== 1) {
      throw new Error(`Unsupported protocol: ${result.protocolVersion}`);
    }
  }

  private async authenticate(): Promise<void> {
    await this.send("authenticate", { methodId: "cursor_login" });
  }

  private async createSession(cwd?: string): Promise<void> {
    const result = (await this.send("session/new", {
      cwd: cwd || process.cwd(),
      mcpServers: [],
    })) as { sessionId: string };
    
    this.sessionId = result.sessionId;
  }

  setPermissionHandler(handler: (request: PermissionRequest) => Promise<string>): void {
    this.permissionHandler = handler;
  }

  setUpdateCallback(cb: (content: string) => void): void {
    this.updateCallback = cb;
  }

  stop(): void {
    this.pending.forEach(({ reject }) => reject(new Error("Stopped")));
    this.pending.clear();
    if (this.agent) {
      this.agent.kill();
      this.agent = null;
    }
    this.sessionId = null;
  }
}

// ============================================================================
// Extension
// ============================================================================

export default function (pi: ExtensionAPI) {
  console.log("[pi-cursor] Extension loaded");

  // ==========================================================================
  // Tool: Run Cursor Agent
  // ==========================================================================
  pi.registerTool({
    name: "cursor_agent",
    label: "Cursor Agent",
    description: "Use Cursor's AI agent to help with coding tasks",
    parameters: Type.Object({
      prompt: Type.String({ description: "What to ask Cursor's agent to do" }),
      mode: Type.Optional(
        Type.Union([
          Type.Literal("agent", { description: "Full tool access" }),
          Type.Literal("plan", { description: "Read-only planning" }),
          Type.Literal("ask", { description: "Q&A mode" }),
        ])
      ),
      cwd: Type.Optional(Type.String({ description: "Working directory" })),
    }),
    async execute(toolCallId, params, signal, onUpdate, ctx) {
      const session = new ACPSession();
      
      // Set up streaming
      session.setUpdateCallback((content) => {
        onUpdate?.({ content: [{ type: "text", text: content + "..." }] });
      });

      // Set up permission handler
      session.setPermissionHandler(async (request) => {
        const toolName = request.toolName || "unknown";
        const argsStr = JSON.stringify(request.toolArgs || {}, null, 2);
        const confirmed = await ctx.ui.confirm(
          `Cursor wants to run: ${toolName}`,
          `${request.justification || ""}\n\nArgs: ${argsStr}\n\nAllow?`
        );
        return confirmed ? "allow-once" : "reject-once";
      });

      try {
        const result = await session.run(
          params.prompt,
          params.mode || "agent",
          params.cwd || String(ctx.cwd || process.cwd())
        );
        
        return { content: [{ type: "text", text: result }] };
      } catch (error) {
        console.error("[pi-cursor] Tool error:", error);
        return {
          content: [{
            type: "text",
            text: `Error: ${error instanceof Error ? error.message : String(error)}`,
          }],
        };
      }
    },
  });

  // ==========================================================================
  // Tool: Status
  // ==========================================================================
  pi.registerTool({
    name: "cursor_status",
    label: "Cursor Status",
    description: "Check Cursor agent status",
    parameters: Type.Object({}),
    async execute() {
      return {
        content: [{
          type: "text",
          text: "Use @cursor_agent to run Cursor. Each request creates a fresh session.",
        }],
      };
    },
  });

  // ==========================================================================
  // Command: Help
  // ==========================================================================
  pi.registerCommand("cursor:help", {
    description: "Show Cursor extension help",
    handler: async (_args, ctx) => {
      ctx.ui.notify(
        "Cursor Extension\n\n" +
        "Use @cursor_agent(prompt=\"...\") to run Cursor.\n" +
        "Modes: agent (default), plan, ask\n" +
        "Example: @cursor_agent(prompt=\"Fix this bug\", mode=\"agent\")",
        "info"
      );
    },
  });
}