import { spawn } from "node:child_process";
import { createHash } from "node:crypto";
import { readdirSync, statSync } from "node:fs";
import { createConnection } from "node:net";
import { join } from "node:path";
import { StringEnum, Type, type ImageContent, type TextContent } from "@earendil-works/pi-ai";
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { Text } from "@earendil-works/pi-tui";

const runtimeDir = process.env.XDG_RUNTIME_DIR ?? `/run/user/${process.getuid?.()}`;

interface ComputerUseDetails {
  screenshotIds: string[];
}

function waylandEnv() {
  const configured = process.env.WAYLAND_DISPLAY;
  if (configured) {
    try {
      if (statSync(join(runtimeDir, configured)).isSocket()) {
        return { ...process.env, XDG_RUNTIME_DIR: runtimeDir, WAYLAND_DISPLAY: configured };
      }
    } catch {}
  }

  const display = readdirSync(runtimeDir)
    .filter((name) => /^wayland-\d+$/.test(name))
    .map((name) => ({ name, modified: statSync(join(runtimeDir, name)).mtimeMs }))
    .sort((a, b) => b.modified - a.modified)[0]?.name;

  if (!display) throw new Error(`No Wayland compositor found in ${runtimeDir}`);
  return { ...process.env, XDG_RUNTIME_DIR: runtimeDir, WAYLAND_DISPLAY: display };
}

function wait(ms: number, signal?: AbortSignal): Promise<void> {
  return new Promise((resolve, reject) => {
    if (signal?.aborted) {
      reject(new Error("Operation aborted"));
      return;
    }

    const timer = setTimeout(() => {
      signal?.removeEventListener("abort", abort);
      resolve();
    }, ms);
    const abort = () => {
      clearTimeout(timer);
      reject(new Error("Operation aborted"));
    };
    signal?.addEventListener("abort", abort, { once: true });
  });
}

function run(command: string, args: string[], signal?: AbortSignal): Promise<Buffer> {
  return new Promise((resolve, reject) => {
    if (signal?.aborted) {
      reject(new Error("Operation aborted"));
      return;
    }

    const child = spawn(command, args, {
      env: waylandEnv(),
      stdio: ["ignore", "pipe", "pipe"],
    });
    const stdout: Buffer[] = [];
    const stderr: Buffer[] = [];
    const abort = () => child.kill();

    child.stdout.on("data", (chunk: Buffer) => stdout.push(chunk));
    child.stderr.on("data", (chunk: Buffer) => stderr.push(chunk));
    child.on("error", reject);
    child.on("close", (code) => {
      signal?.removeEventListener("abort", abort);
      if (signal?.aborted) reject(new Error("Operation aborted"));
      else if (code === 0) resolve(Buffer.concat(stdout));
      else {
        reject(
          new Error(
            Buffer.concat(stderr).toString().trim() || `${command} exited with code ${code}`,
          ),
        );
      }
    });

    signal?.addEventListener("abort", abort, { once: true });
  });
}

async function connectVnc(signal?: AbortSignal) {
  const socket = createConnection({ host: "127.0.0.1", port: 5900, signal });
  socket.on("error", () => {});
  socket.setTimeout(5_000, () => socket.destroy(new Error("WayVNC handshake timed out")));
  const chunks = socket.iterator({ destroyOnReturn: false });
  let buffered = Buffer.alloc(0);

  async function read(length: number): Promise<Buffer> {
    while (buffered.length < length) {
      const chunk = await chunks.next();
      if (chunk.done) throw new Error("WayVNC disconnected during handshake");
      buffered = Buffer.concat([buffered, chunk.value]);
    }
    const result = buffered.subarray(0, length);
    buffered = buffered.subarray(length);
    return result;
  }

  try {
    if ((await read(12)).toString() !== "RFB 003.008\n") {
      throw new Error("Expected WayVNC to support RFB 3.8");
    }
    socket.write("RFB 003.008\n");
    const securityTypes = await read((await read(1))[0]);
    if (!securityTypes.includes(1)) {
      throw new Error("Local WayVNC must allow unauthenticated connections");
    }
    socket.write(Buffer.from([1]));
    if ((await read(4)).readUInt32BE() !== 0) {
      throw new Error("WayVNC rejected the connection");
    }
    socket.write(Buffer.from([1])); // Share the desktop with existing VNC clients.
    await read(24);
    await chunks.return?.();
    socket.setTimeout(0);
    socket.resume();

    // Give Firefox time to bind the seat's newly advertised input devices.
    await wait(100, signal);
    if (socket.destroyed) throw socket.errored ?? new Error("WayVNC disconnected");
    return socket;
  } catch (error) {
    socket.destroy();
    throw error;
  }
}

export default function (pi: ExtensionAPI) {
  if (!process.env.WAYLAND_DISPLAY) return;

  const visibleScreenshotResults = new Set<string>();
  pi.on("agent_end", () => visibleScreenshotResults.clear());
  pi.on("context", (event) => ({
    messages: event.messages.map((message) => {
      if (
        message.role !== "toolResult" ||
        message.toolName !== "computer_use" ||
        visibleScreenshotResults.has(message.toolCallId)
      ) {
        return message;
      }

      return { ...message, content: message.content.filter((content) => content.type !== "image") };
    }),
  }));

  pi.registerTool({
    name: "computer_use",
    label: "Computer Use",
    description:
      "Interact with visible desktop applications. Prefer web_search and web_fetch for ordinary web research and page reading because they return text directly; use the GUI browser when visual or interactive access is useful. Coordinates are absolute across the current 1920×1080 screen, from (0, 0) at top-left to (1919, 1079) at bottom-right. Take a screenshot before and after interacting; use separate calls when an intermediate screenshot is needed.",
    promptSnippet:
      "Interact with visible desktop applications; prefer textual web tools for ordinary research",
    promptGuidelines: [
      "Prefer web_search/web_fetch over opening Firefox and reading screenshots when either textual tool can handle the task adequately.",
      "All coordinates are absolute screen positions on a 1920×1080 screen, not positions relative to a window, element, screenshot crop, or the current pointer.",
      "Screenshots remain available throughout the current turn. Take another after interacting with the computer, or recall one from an earlier turn by its ID.",
      "Recalled screenshots are not visible until the entire action sequence completes.",
    ],
    parameters: Type.Object(
      {
        actions: Type.Array(
          Type.Union([
            Type.Object(
              {
                action: StringEnum(["move"] as const),
                x: Type.Integer({
                  minimum: 0,
                  maximum: 1919,
                  description: "Absolute screen x-coordinate: 0 is left, 1919 is right",
                }),
                y: Type.Integer({
                  minimum: 0,
                  maximum: 1079,
                  description: "Absolute screen y-coordinate: 0 is top, 1079 is bottom",
                }),
              },
              { additionalProperties: false },
            ),
            Type.Object(
              {
                action: StringEnum(["click"] as const),
                x: Type.Optional(
                  Type.Integer({
                    minimum: 0,
                    maximum: 1919,
                    description:
                      "Absolute screen x-coordinate (0 left to 1919 right); omit x and y to click in place",
                  }),
                ),
                y: Type.Optional(
                  Type.Integer({
                    minimum: 0,
                    maximum: 1079,
                    description:
                      "Absolute screen y-coordinate (0 top to 1079 bottom); omit x and y to click in place",
                  }),
                ),
                button: Type.Optional(
                  StringEnum(["left", "middle", "right"] as const, { default: "left" }),
                ),
              },
              { additionalProperties: false },
            ),
            Type.Object(
              {
                action: StringEnum(["scroll"] as const),
                x: Type.Optional(
                  Type.Integer({
                    minimum: 0,
                    maximum: 1919,
                    description:
                      "Absolute screen x-coordinate (0 left to 1919 right); x and y must be supplied together",
                  }),
                ),
                y: Type.Optional(
                  Type.Integer({
                    minimum: 0,
                    maximum: 1079,
                    description:
                      "Absolute screen y-coordinate (0 top to 1079 bottom); x and y must be supplied together",
                  }),
                ),
                deltaX: Type.Optional(Type.Integer({ description: "Horizontal scroll amount" })),
                deltaY: Type.Optional(
                  Type.Integer({ description: "Vertical scroll amount; positive scrolls down" }),
                ),
              },
              { additionalProperties: false },
            ),
            Type.Object(
              {
                action: StringEnum(["type"] as const),
                text: Type.String({
                  description: "Literal text to type; use a key action for named keys and shortcuts",
                }),
              },
              { additionalProperties: false },
            ),
            Type.Object(
              {
                action: StringEnum(["key"] as const),
                key: Type.String({
                  description:
                    "One XKB key name, such as Return, Escape, Tab, BackSpace, Delete, Left, Page_Down, F5, or a. Put shortcut modifiers in modifiers; do not put them in key",
                  pattern: "^[A-Za-z0-9_]+$",
                }),
                modifiers: Type.Optional(
                  Type.Array(StringEnum(["shift", "ctrl", "alt", "logo"] as const), {
                    description: "Modifiers held while pressing the key",
                    uniqueItems: true,
                  }),
                ),
              },
              { additionalProperties: false },
            ),
            Type.Object(
              {
                action: StringEnum(["new_browser_tab"] as const),
                url: Type.String({
                  description:
                    "HTTP or HTTPS URL to open in a new Firefox tab every time; prefer web_fetch for ordinary page reading",
                  pattern: "^https?://",
                }),
              },
              { additionalProperties: false },
            ),
            Type.Object(
              {
                action: StringEnum(["recall"] as const),
                id: Type.String({
                  description: "ID from an earlier screenshot result",
                  pattern: "^sc_[0-9a-f]{16}$",
                }),
              },
              { additionalProperties: false },
            ),
            Type.Object(
              { action: StringEnum(["screenshot"] as const) },
              { additionalProperties: false },
            ),
            Type.Object(
              {
                action: StringEnum(["sleep"] as const),
                ms: Type.Integer({
                  minimum: 0,
                  maximum: 30_000,
                  description: "Time to wait, in milliseconds",
                }),
              },
              { additionalProperties: false },
            ),
          ]),
          { minItems: 1, maxItems: 20 },
        ),
      },
      { additionalProperties: false },
    ),
    renderCall(args, theme) {
      const actions = args.actions.map((action) => {
        if (action.action === "move") return `move (${action.x ?? "…"}, ${action.y ?? "…"})`;
        if (action.action === "click") {
          const position =
            action.x === undefined && action.y === undefined
              ? ""
              : ` (${action.x ?? "…"}, ${action.y ?? "…"})`;
          return `click ${action.button ?? "left"}${position}`;
        }
        if (action.action === "scroll") return `scroll (${action.deltaX ?? 0}, ${action.deltaY ?? 0})`;
        if (action.action === "type") return `type ${JSON.stringify(action.text ?? "")}`;
        if (action.action === "key") return `press ${[...(action.modifiers ?? []), action.key ?? "…"].join("+")}`;
        if (action.action === "new_browser_tab") return `new tab ${action.url ?? "…"}`;
        if (action.action === "recall") return `recall ${action.id ?? "…"}`;
        if (action.action === "screenshot") return "screenshot";
        return `wait ${action.ms ?? 0}ms`;
      });
      return new Text(
        theme.fg("toolTitle", theme.bold("computer_use ")) +
          theme.fg("accent", actions.join(" → ")),
        0,
        0,
      );
    },
    async execute(toolCallId, params, signal, _onUpdate, ctx) {
      const screenshots: { id: string; image: ImageContent; recalled: boolean }[] = [];

      // WayVNC keeps seat capabilities stable while wlrctl/wtype come and go.
      const vnc = params.actions.some((action) =>
        ["move", "click", "scroll", "type", "key"].includes(action.action),
      )
        ? await connectVnc(signal)
        : undefined;
      try {
        for (const action of params.actions) {
          if (vnc?.destroyed) throw vnc.errored ?? new Error("WayVNC disconnected");
          if (action.action === "move") {
            await run("wlrctl", ["pointer", "move", "-100000", "-100000"], signal);
            await run("wlrctl", ["pointer", "move", String(action.x), String(action.y)], signal);
          } else if (action.action === "click") {
            const hasPosition = action.x !== undefined || action.y !== undefined;
            if (hasPosition && (action.x === undefined || action.y === undefined)) {
              throw new Error("x and y must be supplied together");
            }
            if (hasPosition) {
              await run("wlrctl", ["pointer", "move", "-100000", "-100000"], signal);
              await run("wlrctl", ["pointer", "move", String(action.x), String(action.y)], signal);
            }
            await run("wlrctl", ["pointer", "click", action.button ?? "left"], signal);
          } else if (action.action === "scroll") {
            const hasPosition = action.x !== undefined || action.y !== undefined;
            if (hasPosition && (action.x === undefined || action.y === undefined)) {
              throw new Error("x and y must be supplied together");
            }
            if (hasPosition) {
              await run("wlrctl", ["pointer", "move", "-100000", "-100000"], signal);
              await run("wlrctl", ["pointer", "move", String(action.x), String(action.y)], signal);
            }
            await run(
              "wlrctl",
              ["pointer", "scroll", String(action.deltaY ?? 0), String(action.deltaX ?? 0)],
              signal,
            );
          } else if (action.action === "type") {
            await run("wtype", ["--", action.text], signal);
          } else if (action.action === "key") {
            await run(
              "wtype",
              [...(action.modifiers ?? []).flatMap((modifier) => ["-M", modifier]), "-k", action.key],
              signal,
            );
          } else if (action.action === "new_browser_tab") {
            const url = new URL(action.url);
            if (url.protocol !== "http:" && url.protocol !== "https:") {
              throw new Error("Only HTTP and HTTPS URLs are supported");
            }

            const child = spawn("firefox", ["--new-tab", url.href], {
              env: { ...waylandEnv(), MOZ_ENABLE_WAYLAND: "1" },
              detached: true,
              stdio: "ignore",
            });
            await new Promise<void>((resolve, reject) => {
              child.once("spawn", resolve);
              child.once("error", reject);
            });
            child.unref();
          } else if (action.action === "recall") {
            const branch = ctx.sessionManager.getBranch();
            let image: ImageContent | undefined;

            for (let index = branch.length - 1; index >= 0 && !image; index--) {
              const entry = branch[index];
              if (
                entry.type !== "message" ||
                entry.message.role !== "toolResult" ||
                entry.message.toolName !== "computer_use"
              ) {
                continue;
              }

              const screenshotIds = (entry.message.details as ComputerUseDetails | undefined)
                ?.screenshotIds;
              if (!Array.isArray(screenshotIds)) continue;

              const imageIndex = screenshotIds.indexOf(action.id);
              if (imageIndex === -1) continue;
              image = entry.message.content.filter((content) => content.type === "image")[imageIndex];
            }

            if (!image) throw new Error(`Screenshot ${action.id} was not found in this session branch`);
            screenshots.push({ id: action.id, image, recalled: true });
          } else if (action.action === "screenshot") {
            const png = await run("grim", ["-c", "-"], signal);
            screenshots.push({
              id: `sc_${createHash("sha256").update(png).digest("hex").slice(0, 16)}`,
              image: {
                type: "image",
                data: png.toString("base64"),
                mimeType: "image/png",
              },
              recalled: false,
            });
          } else {
            await wait(action.ms ?? 0, signal);
          }
        }

        if (vnc?.destroyed) throw vnc.errored ?? new Error("WayVNC disconnected");
      } finally {
        if (vnc) {
          // Let clients process the last input before removing the seat's devices.
          await wait(100);
          vnc.destroy();
        }
      }

      const content: (TextContent | ImageContent)[] = [
        { type: "text", text: `${params.actions.length} actions completed.` },
      ];
      for (const screenshot of screenshots) {
        content.push(
          {
            type: "text",
            text: `Screenshot ${screenshot.id} (${screenshot.recalled ? "recalled" : "captured"}):`,
          },
          screenshot.image,
        );
      }

      if (screenshots.length) visibleScreenshotResults.add(toolCallId);
      return {
        content,
        details: { screenshotIds: screenshots.map(({ id }) => id) } satisfies ComputerUseDetails,
      };
    },
  });
}