import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { Type } from "typebox";
import { chromium, Browser } from "playwright";

import { tmpdir } from "node:os";
import { join } from "node:path";
import { randomUUID } from "node:crypto";

export default function (pi: ExtensionAPI) {
  let browserPromise: Promise<Browser> | undefined;

  pi.on("session_shutdown", async () => {
    if (browserPromise) {
      const browser = await browserPromise;
      await browser.close();
    }
  });

  pi.registerTool({
    name: "screenshot",
    label: "Screenshot",
    description: "Take a screenshot of a webpage and return it as an image for visual analysis",
    parameters: Type.Object({
      url: Type.String({ description: "The URL to screenshot" }),
      fullPage: Type.Optional(Type.Boolean({ description: "Capture the full scrollable page" })),
      width: Type.Optional(Type.Integer({ description: "Viewport width in pixels", default: 1920 })),
      height: Type.Optional(Type.Integer({ description: "Viewport height in pixels", default: 1080 })),
    }),
    async execute(_toolCallId, params, signal) {
      if (!browserPromise) {
        browserPromise = chromium.launch({ headless: true });
      }
      const browser = await browserPromise;

      const context = await browser.newContext({
        viewport: {
          width: params.width ?? 1920,
          height: params.height ?? 1080,
        },
      });

      const page = await context.newPage();

      try {
        if (signal?.aborted) throw new Error("Cancelled");

        await page.goto(params.url, {
          waitUntil: "networkidle",
          timeout: 120_000,
        });

        const path = join(tmpdir(), `pi-screenshot-${randomUUID()}.png`);
        await page.screenshot({ path, fullPage: params.fullPage ?? false });

        return {
          content: [
            {
              type: "text",
              text: `Screenshot saved to ${path}. Use the read tool to view it.`,
            },
          ],
          details: { path },
        };
      } finally {
        await context.close();
      }
    },
  });
}