import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { Type } from "typebox";
import * as fs from "node:fs";
import * as path from "node:path";
import * as os from "node:os";
import * as crypto from "node:crypto";
import { fileURLToPath } from "node:url";
import { spawn, type ChildProcess } from "node:child_process";

// ── Config (no hardcoded paths — everything env-overridable) ──────────────
// PIMAGE_DIR: where inbox/outbox/signal live (default ~/.pimage)
// PIMAGE_DAEMON: explicit daemon path override (default: package daemons/)
const PIDIR = process.env.PIMAGE_DIR ?? path.join(os.homedir(), ".pimage");
const INBOX = path.join(PIDIR, "in");
const OUTBOX = path.join(PIDIR, "out");
const SIGNAL = path.join(PIDIR, "current");
function resolveDaemonScript(): string {
  // 1. explicit override
  if (process.env.PIMAGE_DAEMON) return process.env.PIMAGE_DAEMON;
  // 2. legacy location ~/.pimage/pimage-daemon.py (kept for existing installs)
  const legacy = path.join(PIDIR, "pimage-daemon.py");
  if (fs.existsSync(legacy)) return legacy;
  // 3. package daemons dir
  return path.join(path.dirname(fileURLToPath(import.meta.url)), "..", "daemons", "pimage-daemon.py");
}

const DAEMON_SCRIPT = resolveDaemonScript();

// Timeout: how long to wait for pimage to respond (ms)
const POLL_TIMEOUT = Number(process.env.PIMAGE_TIMEOUT_MS ?? 120_000);
const POLL_INTERVAL = 200;
const DAEMON_READY_TIMEOUT = Number(process.env.PIMAGE_DAEMON_READY_MS ?? 8_000);

// ── Lazy daemon lifecycle ──────────────────────────────────────
let daemonProcess: ChildProcess | null = null;
let daemonStarting: Promise<boolean> | null = null;

/**
 * Ensure the pimage daemon is running. Spawns it on demand.
 * Returns true if ready, false if it couldn't be started.
 * Thread-safe for concurrent calls via daemonStarting promise gate.
 */
async function ensureDaemon(signal: AbortSignal): Promise<boolean> {
  if (daemonProcess && daemonProcess.exitCode === null) return true;
  if (daemonStarting) return daemonStarting;

  daemonStarting = startDaemon(signal);
  try {
    return await daemonStarting;
  } finally {
    daemonStarting = null;
  }
}

async function startDaemon(signal: AbortSignal): Promise<boolean> {
  if (!fs.existsSync(DAEMON_SCRIPT)) {
    console.error(`[pimage] daemon script not found at ${DAEMON_SCRIPT}`);
    return false;
  }

  return new Promise((resolve) => {
    const proc = spawn("python3", [DAEMON_SCRIPT], {
      stdio: ["ignore", "pipe", "pipe"],
      env: { ...process.env, PIMAGE_DIR: PIDIR },
    });

    daemonProcess = proc;

    let settled = false;
    const finish = (ok: boolean) => {
      if (settled) return;
      settled = true;
      resolve(ok);
    };

    proc.stdout?.on("data", (data: Buffer) => {
      const text = data.toString();
      // Daemon prints "Watching ..." once it's in the polling loop
      if (text.includes("Watching")) {
        finish(true);
      }
    });

    proc.on("error", () => finish(false));

    proc.on("exit", (code) => {
      daemonProcess = null;
      finish(code === 0);
    });

    // Safety timeout — proceed anyway after DAEMON_READY_TIMEOUT
    // so a slow startup doesn't block forever
    const timeout = setTimeout(() => finish(true), DAEMON_READY_TIMEOUT);

    // If the tool gets aborted during daemon startup, clean up
    const onAbort = () => {
      clearTimeout(timeout);
      proc.kill();
      daemonProcess = null;
      finish(false);
    };
    if (signal.aborted) {
      onAbort();
    } else {
      signal.addEventListener("abort", onAbort, { once: true });
    }
  });
}

function stopDaemon(): void {
  if (daemonProcess && daemonProcess.exitCode === null) {
    daemonProcess.kill();
    daemonProcess = null;
  }
}

// ── Helpers ─────────────────────────────────────────────────────

function uuid(): string {
  return crypto.randomUUID();
}

function ensureDir(dir: string) {
  if (!fs.existsSync(dir)) fs.mkdirSync(dir, { recursive: true });
}

function sleep(ms: number): Promise<void> {
  return new Promise((r) => setTimeout(r, ms));
}

async function waitForResult(
  resultDir: string,
  signal: AbortSignal,
  timeout: number
): Promise<{ success: boolean; content: string; error: string | null } | null> {
  const deadline = Date.now() + timeout;
  while (Date.now() < deadline) {
    if (signal.aborted) return null;
    const resultFile = path.join(resultDir, "result.json");
    try {
      const raw = fs.readFileSync(resultFile, "utf-8");
      return JSON.parse(raw);
    } catch {
      // not ready yet
    }
    await sleep(POLL_INTERVAL);
  }
  return null; // timeout
}

// ── Extension ────────────────────────────────────────────────────
export default function (pi: ExtensionAPI) {
  // ── Ensure pimage directories exist ──────────────────────────
  ensureDir(INBOX);
  ensureDir(OUTBOX);

  // ── Tool: forward_to_pimage ──────────────────────────────────
  pi.registerTool({
    name: "forward_to_pimage",
    label: "Forward to Pimage",
    description:
      "Forward an image to the pimage vision agent for processing. " +
      "Use this when you receive an image but cannot process it yourself " +
      "(your model does not support vision).",
    promptSnippet: "Forward an image to the vision-specialist agent (pimage) for description, OCR, or analysis",
    promptGuidelines: [
      "Use forward_to_pimage when the user pastes or references an image and your model doesn't support vision.",
      "Pass the EXACT image path from the context — do not guess or hallucinate paths.",
      "Provide a clear instruction: 'Describe this image', 'Extract text from this image', 'What objects are in this image?', etc.",
      "The tool returns the vision model's response — relay it to the user as-is or summarize it.",
    ],
    parameters: Type.Object({
      image_path: Type.String({
        description:
          "Absolute path to the image file on disk. You MUST get this from the context — the user's message or the system injected the path.",
      }),
      instruction: Type.String({
        description:
          "What to do with the image. Be specific. Examples: 'Describe this image in detail', 'Extract all text', 'What is the woman wearing?', 'Read the chart values'.",
      }),
    }),

    async execute(toolCallId, params, signal, onUpdate, ctx) {
      const id = uuid();
      const reqDir = path.join(INBOX, id);
      const outDir = path.join(OUTBOX, id);
      ensureDir(reqDir);

      const imagePath = params.image_path;
      const instruction = params.instruction;

      // ── Lazy-start pimage daemon if not running ──────────────
      if (!(await ensureDaemon(signal))) {
        return {
          content: [
            {
              type: "text",
              text: `❌ Failed to start pimage daemon (script not found at ${DAEMON_SCRIPT}). Set PIMAGE_DAEMON to point at pimage-daemon.py if your install layout differs.`,
            },
          ],
          isError: true,
          details: {},
        };
      }

      // ── Stage the image ──────────────────────────────────────
      let stagedImage = false;
      if (imagePath && fs.existsSync(imagePath)) {
        try {
          const dest = path.join(reqDir, "image");
          fs.copyFileSync(imagePath, dest);
          stagedImage = true;
        } catch (e) {
          // non-fatal — daemon will warn
        }
      }

      // ── Write request ────────────────────────────────────────
      const request = {
        id,
        image_path: imagePath,
        instruction,
        prompt: instruction,
      };
      fs.writeFileSync(
        path.join(reqDir, "request.json"),
        JSON.stringify(request, null, 2)
      );

      // ── Signal pimage daemon ─────────────────────────────────
      fs.writeFileSync(SIGNAL, id);

      // ── Wait for result ──────────────────────────────────────
      await onUpdate({
        content: [
          {
            type: "text",
            text: stagedImage
              ? `⏳ Forwarded image to pimage (${path.basename(imagePath)}) — waiting for vision model...`
              : `⚠️  Image file not found at ${imagePath}, but still sending request to pimage...`,
          },
        ],
        details: {},
      });

      const result = await waitForResult(outDir, signal, POLL_TIMEOUT);

      // Cleanup request files
      try {
        fs.rmSync(reqDir, { recursive: true, force: true });
      } catch {
        // best effort
      }

      if (!result) {
        return {
          content: [
            {
              type: "text",
              text: `❌ pimage did not respond within ${POLL_TIMEOUT / 1000}s. The vision model may still be starting up — try again.`,
            },
          ],
          isError: true,
          details: {},
        };
      }

      if (!result.success) {
        return {
          content: [
            {
              type: "text",
              text: `❌ pimage failed: ${result.error ?? "Unknown error"}`,
            },
          ],
          isError: true,
          details: {},
        };
      }

      return {
        content: [
          {
            type: "text",
            text: `🧩 **pimage response:**\n\n${result.content}`,
          },
        ],
        details: {},
      };
    },
  });

  // ── Clean up daemon on session end ─────────────────────────
  pi.on("session_shutdown", async () => {
    stopDaemon();
  });

  // ── Inject system prompt context about the bridge ────────────
  pi.on("before_agent_start", async (event, ctx) => {
    // Always add a brief note about the bridge so the LLM knows about the tool
    let systemPromptAdditions = `
## Vision Bridge
` +
      `- Your model does not support images. Use the \`forward_to_pimage\` tool to process images.\n` +
      `- pimage is a vision agent that can describe, OCR, analyze, or answer questions about images.\n` +
      `- Pass the exact image path and a clear instruction about what to do with it.\n`;

    let messageToInject: string | null = null;

    // If the model doesn't support images but the user attached an image,
    // inject a message with the image path so the LLM can use forward_to_pimage
    if (event.images && event.images.length > 0) {
      const model = ctx.model;
      // Check model support: ctx.model may be an object {input, ...} or a string id
      const modelId = typeof model === "string" ? model : (model as any)?.id ?? "";
      const modelInput = typeof model === "object" ? (model as any)?.input : undefined;
      const supportsImages =
        modelInput?.includes("image") ??
        /mimo|vision|gpt-4o|claude-3\.5-sonnet|gemini|image/i.test(modelId);

      if (!supportsImages) {
        // Save images to a known location and inject the paths
        const savedPaths: string[] = [];
        for (let i = 0; i < event.images.length; i++) {
          const img = event.images[i];
          const imgPath = typeof img === "string" ? img : (img as any).path;
          if (imgPath && fs.existsSync(imgPath)) {
            const imgDir = path.join(os.tmpdir(), "pimage-attachments");
            ensureDir(imgDir);
            const stablePath = path.join(
              imgDir,
              `img-${Date.now()}-${i}${path.extname(imgPath) || ".png"}`
            );
            fs.copyFileSync(imgPath, stablePath);
            savedPaths.push(stablePath);
          }
        }

        if (savedPaths.length > 0) {
          const pathList = savedPaths.map((p) => `- \`${p}\``).join("\n");
          messageToInject =
            `## 📸 Image Attached, But Not Visible to Me\n\n` +
            `The user attached ${
              savedPaths.length === 1
                ? "an image"
                : `${savedPaths.length} images`
            }. My model does not support vision, so I cannot see it directly.\n\n` +
            `**Image path${savedPaths.length > 1 ? "s" : ""}:**\n${pathList}\n\n` +
            `**To process this image, use the \`forward_to_pimage\` tool.**\n` +
            `Pass the image path and your instruction (describe, OCR, analyze, etc.).\n` +
            `pimage will process it and return the result.`;
        }
      }
    }

    return {
      ...(messageToInject
        ? {
            message: {
              customType: "pimage-bridge",
              content: messageToInject,
              display: true,
            },
          }
        : {}),
      systemPrompt: event.systemPrompt + systemPromptAdditions,
    };
  });
}