import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
import { StringEnum } from "@mariozechner/pi-ai";
import { Type } from "@sinclair/typebox";
import { getServiceKey } from "../shared/auth";
import { isAbortError } from "../shared/errors";
import { extractWithExa, extractWithJina, extractWithMetaso, extractWithTavilyEngine, type ReadContent, type ReadFailure, stringifyReadError } from "../clients/read";
import { shortenUrl } from "../shared/render";
import { READ_ENGINES, type ReadEngine, type WebReadDetails } from "../types";
import { buildExpandedParams, kylinFrameForCall, kylinFrameForResult } from "../../../vera-theme/src/public";

function buildWebReadParams(args: any, theme: any): string {
  const urls = args.urls ?? [];
  const engines = normalizeEngines(args.engines);
  const target = urls.length === 1 ? shortenUrl(urls[0]) : `${urls.length} URLs`;
  const enginesPart = engines.length > 0 ? formatEngineList(engines) : "—";
  return (
    theme.fg("text", target) +
    theme.fg("borderMuted", " · ") +
    theme.fg("dim", enginesPart) +
    (args.format && args.format !== "markdown"
      ? theme.fg("borderMuted", " · ") + theme.fg("muted", args.format)
      : "")
  );
}

function extractTextPreview(content: any, maxLines: number): { lines: string[]; totalLines: number } {
  if (content?.type !== "text" || typeof content.text !== "string") {
    return { lines: [], totalLines: 0 };
  }
  const allLines = content.text.split("\n");
  return {
    lines: allLines.slice(0, maxLines),
    totalLines: allLines.length,
  };
}

const ENGINE_LABELS: Record<ReadEngine, string> = {
  tavily: "Tavily",
  jina: "Jina",
  exa: "Exa",
  metaso: "Metaso",
};

const ENGINE_ENV_VARS: Record<ReadEngine, string> = {
  tavily: "TAVILY_API_KEY",
  jina: "JINA_API_KEY",
  exa: "EXA_API_KEY",
  metaso: "METASO_API_KEY",
};

const READ_CLIENTS: Record<ReadEngine, (urls: string[], format: "markdown" | "text", apiKey: string, signal?: AbortSignal) => Promise<{ results: ReadContent[]; failed: ReadFailure[] }>> = {
  tavily: extractWithTavilyEngine,
  jina: extractWithJina,
  exa: extractWithExa,
  metaso: extractWithMetaso,
};

function normalizeEngines(value: unknown): ReadEngine[] {
  if (!Array.isArray(value)) return [];
  const valid = value.filter((engine): engine is ReadEngine => READ_ENGINES.includes(engine as ReadEngine));
  return Array.from(new Set(valid));
}

function formatEngineList(engines: ReadEngine[]): string {
  return engines.map((engine) => ENGINE_LABELS[engine]).join(" + ");
}

export function registerWebReadTool(pi: ExtensionAPI): void {
  pi.registerTool({
    name: "web_read",
    label: "Web Read",
    description: "Extract clean, readable content from one or more web URLs using selected fetch engines (Tavily, Jina, Exa, Metaso). Returns markdown or plain text. Use this when you need to read a web page's content.",
    promptSnippet: "Use web_read to fetch and extract content from web URLs with selected engines.",
    parameters: Type.Object({
      urls: Type.Array(Type.String({ description: "URL to extract content from" }), {
        minItems: 1,
        maxItems: 5,
        description: "URLs to extract (1-5)",
      }),
      engines: Type.Array(StringEnum(READ_ENGINES, { description: "Fetch engine to use" }), {
        minItems: 1,
        uniqueItems: true,
        description: "Fetch engines to query, e.g. [\"jina\", \"tavily\"]",
      }),
      format: Type.Optional(StringEnum(["markdown", "text"] as const, {
        description: "Output format (default: markdown)",
      })),
    }),
    renderShell: "self" as const,
    renderCall(args, theme, ctx) {
      const state = ctx.state as { startedAt?: number };
      if (state.startedAt === undefined) state.startedAt = Date.now();
      return kylinFrameForCall(ctx, {
        name: "web_read",
        theme,
        status: "pending",
        paramSummary: buildWebReadParams(args, theme),
        startedAt: state.startedAt,
      });
    },
    renderResult(result, { expanded, isPartial }, theme, context) {
      const details = result.details as WebReadDetails | undefined;
      const state = context.state as { startedAt?: number };
      const startedAt = state.startedAt;
      const paramSummary = buildWebReadParams(context.args, theme);
      const expandedParams = expanded ? buildExpandedParams([
        { label: "urls", value: context.args?.urls },
        { label: "engines", value: context.args?.engines },
        { label: "format", value: context.args?.format },
      ], theme) : undefined;

      if (isPartial) {
        return kylinFrameForResult(context, {
          name: "web_read",
          theme,
          status: "pending",
          paramSummary,
          expanded,
          expandedParams,
          startedAt,
        });
      }

      if (details?.error) {
        return kylinFrameForResult(context, {
          name: "web_read",
          theme,
          status: "error",
          paramSummary,
          errorMessage: details.error,
          expanded,
          expandedParams,
          startedAt,
        });
      }
      if (!details) {
        return kylinFrameForResult(context, {
          name: "web_read",
          theme,
          status: "error",
          paramSummary,
          errorMessage: "No result",
          expanded,
          expandedParams,
          startedAt,
        });
      }

      const counts: string[] = [];
      if (details.succeeded > 0) counts.push(theme.fg("success", `✓ ${details.succeeded} extracted`));
      if (details.failed > 0) counts.push(theme.fg("error", `✗ ${details.failed} failed`));
      const resultSummary = counts.length > 0 ? counts.join(theme.fg("borderMuted", " · ")) : undefined;

      const expandedBody: string[] = [];
      for (const meta of details.results ?? []) {
        expandedBody.push(
          theme.fg("success", "✓ ") +
            theme.fg("accent", `[${meta.engine}] `) +
            theme.fg("text", shortenUrl(meta.url)) +
            theme.fg("dim", ` · ${meta.lines} lines`),
        );
      }
      for (const failure of details.failedUrls ?? []) {
        expandedBody.push(
          theme.fg("error", "✗ ") +
            theme.fg("accent", `[${failure.engine}] `) +
            theme.fg("text", shortenUrl(failure.url)) +
            theme.fg("dim", ` · ${failure.error}`),
        );
      }

      let totalExpandedLines = expandedBody.length;
      if (expanded) {
        const preview = extractTextPreview(result.content[0], 20);
        if (preview.lines.length > 0) {
          if (expandedBody.length > 0) expandedBody.push(theme.fg("borderMuted", "─".repeat(40)));
          for (const line of preview.lines) expandedBody.push(theme.fg("dim", line));
          totalExpandedLines = expandedBody.length + Math.max(0, preview.totalLines - preview.lines.length);
        }
      }

      return kylinFrameForResult(context, {
        name: "web_read",
        theme,
        status: details.failed > 0 && details.succeeded === 0 ? "error" : "success",
        paramSummary,
        resultSummary,
        expanded,
        expandedParams,
        expandedBody: expanded ? expandedBody : undefined,
        expandedTotalLines: totalExpandedLines,
        startedAt,
      });
    },
    async execute(_toolCallId, params, signal, onUpdate) {
      const urls = params.urls;
      const engines = normalizeEngines(params.engines);
      const format = params.format ?? "markdown";
      const serviceKeys: Record<ReadEngine, string | null> = {
        tavily: getServiceKey("tavily", "TAVILY_API_KEY"),
        jina: getServiceKey("jina", "JINA_API_KEY"),
        exa: getServiceKey("exa", "EXA_API_KEY"),
        metaso: getServiceKey("metaso", "METASO_API_KEY"),
      };

      if (engines.length === 0) {
        return {
          content: [{ type: "text" as const, text: "Error: No fetch engines selected. Choose at least one of tavily, jina, exa, metaso." }],
          details: { urls, succeeded: 0, failed: 0, error: "No fetch engines selected" } as WebReadDetails,
        };
      }

      const configuredEngines = engines.filter((engine) => serviceKeys[engine]);
      if (configuredEngines.length === 0) {
        const envVars = engines.map((engine) => ENGINE_ENV_VARS[engine]).join(" / ");
        return {
          content: [{ type: "text" as const, text: `Error: No API keys configured for selected engines. Set ${envVars} or add them to agent/auth.json.` }],
          details: { urls, engines, succeeded: 0, failed: 0, error: "No API keys configured for selected engines" } as WebReadDetails,
        };
      }

      onUpdate?.({
        content: [{ type: "text" as const, text: "Fetching..." }],
        details: { urls, engines, succeeded: 0, failed: 0 } as WebReadDetails,
      });

      try {
        const engineResults = await Promise.all(engines.map(async (engine) => {
          const apiKey = serviceKeys[engine];
          if (!apiKey) {
            return {
              engine,
              results: [] as ReadContent[],
              failed: urls.map((url) => ({ url, error: "No API key" })),
            };
          }
          try {
            const data = await READ_CLIENTS[engine](urls, format, apiKey, signal);
            return { engine, ...data };
          } catch (error) {
            if (isAbortError(error)) throw error;
            return {
              engine,
              results: [] as ReadContent[],
              failed: urls.map((url) => ({ url, error: stringifyReadError(error) })),
            };
          }
        }));

        const parts: string[] = [];
        const results: { url: string; engine: ReadEngine; lines: number }[] = [];
        const failedUrls: { url: string; engine: ReadEngine; error: string }[] = [];

        for (const engineResult of engineResults) {
          for (const result of engineResult.results) {
            parts.push(`## [${engineResult.engine}] ${result.url}\n\n${result.content}`);
            results.push({ url: result.url, engine: engineResult.engine, lines: result.content.split("\n").length });
          }
          for (const failure of engineResult.failed) {
            parts.push(`## [${engineResult.engine}] ${failure.url}\n\n[Failed: ${failure.error}]`);
            failedUrls.push({ url: failure.url, engine: engineResult.engine, error: failure.error });
          }
        }

        return {
          content: [{ type: "text" as const, text: parts.join("\n\n---\n\n") }],
          details: {
            urls,
            engines,
            succeeded: results.length,
            failed: failedUrls.length,
            results,
            failedUrls,
          } as WebReadDetails,
        };
      } catch (error) {
        if (isAbortError(error)) {
          return {
            content: [{ type: "text" as const, text: "Request cancelled." }],
            details: { urls, engines, succeeded: 0, failed: 0, error: "Cancelled" } as WebReadDetails,
          };
        }

        const message = stringifyReadError(error);
        return {
          content: [{ type: "text" as const, text: message }],
          details: { urls, engines, succeeded: 0, failed: 0, error: message } as WebReadDetails,
        };
      }
    },
  });
}