import { defineTool, type ExtensionAPI } from "@earendil-works/pi-coding-agent";

import { CuaComputerRuntime, type CuaComputerRuntimeOptions } from "./runtime.ts";
import { computerParameters, validateComputerAction } from "./schema.ts";

export interface ComputerUseExtensionOptions extends CuaComputerRuntimeOptions {
  runtime?: CuaComputerRuntime;
}

export function registerComputerUseExtension(
  pi: ExtensionAPI,
  options: ComputerUseExtensionOptions = {},
): CuaComputerRuntime {
  const runtime = options.runtime ?? new CuaComputerRuntime(options);

  pi.registerTool(
    defineTool({
      name: "computer",
      label: "Computer",
      description:
        "Control the visible primary desktop using coordinates from the latest returned screenshot. " +
        "Supports screenshot, click, double_click, scroll, type, wait, move, keypress, and drag. " +
        "Every supported call returns a fresh screenshot.",
      promptSnippet: "Inspect and control the visible primary desktop",
      promptGuidelines: [
        "Use computer with action.type=\"screenshot\" before the first coordinate action and ground every coordinate in the latest returned screenshot.",
        "After computer reports an unknown action outcome, inspect its fresh screenshot before deciding whether a retry is safe.",
      ],
      parameters: computerParameters,
      executionMode: "sequential",
      async execute(_toolCallId, params, signal) {
        return runtime.execute(validateComputerAction(params.action), signal);
      },
    }),
  );

  pi.registerCommand("computer-permissions", {
    description: "Set up macOS Accessibility and Screen Recording permissions",
    handler: async (_args, ctx) => {
      try {
        const result = await runtime.setupPermissions();
        if (!result.supported) {
          ctx.ui.notify("Computer Use permission setup is only needed on macOS.", "info");
          return;
        }
        if (result.after?.accessibility && result.after.screenRecording) {
          ctx.ui.notify("Computer Use permissions are ready.", "info");
          return;
        }
        const missing = [
          !result.after?.accessibility ? "Accessibility" : undefined,
          !result.after?.screenRecording ? "Screen Recording" : undefined,
        ].filter(Boolean);
        ctx.ui.notify(
          `Grant ${missing.join(" and ")} to the app hosting Pi, then restart Pi.`,
          "warning",
        );
      } catch (error) {
        ctx.ui.notify(`Computer permission setup failed: ${String(error)}`, "error");
      }
    },
  });

  pi.on("session_shutdown", async () => {
    await runtime.close();
  });

  return runtime;
}

export default function computerUseExtension(pi: ExtensionAPI): void {
  registerComputerUseExtension(pi);
}

export { CuaComputerRuntime } from "./runtime.ts";
export { computerParameters, validateComputerAction } from "./schema.ts";
export type { ComputerAction, ComputerParameters } from "./schema.ts";