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 { errorMessage, isAbortError } from "../shared/errors";
import { dedupSearchResults, searchBrave, searchExa, searchJina, searchMetaso } from "../clients/search";
import { searchWithTavily } from "../clients/tavily";
import { DEFAULT_SEARCH_RESULTS, SEARCH_ENGINES, type SearchDetails, type SearchEngine, type SearchResult } from "../types";
import { buildExpandedParams, kylinFrameForCall, kylinFrameForResult } from "../../../vera-theme/src/public";

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

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

const SEARCH_CLIENTS: Record<SearchEngine, (query: string, apiKey: string, count: number, signal?: AbortSignal) => Promise<{ results: SearchResult[]; error?: string }>> = {
  exa: searchExa,
  brave: searchBrave,
  metaso: searchMetaso,
  jina: searchJina,
  tavily: searchWithTavily,
};

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

function normalizeCount(value: unknown): number {
  const parsed = Number(value ?? DEFAULT_SEARCH_RESULTS);
  if (!Number.isFinite(parsed)) return DEFAULT_SEARCH_RESULTS;
  return Math.max(1, Math.floor(parsed));
}

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

function applyEngineOutcome(details: SearchDetails, engine: SearchEngine, outcome: { results: SearchResult[]; error?: string }): void {
  switch (engine) {
    case "exa":
      details.exaCount = outcome.results.length;
      details.exaError = outcome.error;
      return;
    case "brave":
      details.braveCount = outcome.results.length;
      details.braveError = outcome.error;
      return;
    case "metaso":
      details.metasoCount = outcome.results.length;
      details.metasoError = outcome.error;
      return;
    case "jina":
      details.jinaCount = outcome.results.length;
      details.jinaError = outcome.error;
      return;
    case "tavily":
      details.tavilyCount = outcome.results.length;
      details.tavilyError = outcome.error;
      return;
  }
}

function truncate(text: string, max = 70): string {
  return text.length > max ? text.slice(0, max - 1) + "…" : text;
}

function buildParamSummary(args: any, theme: any): string {
  const engines = normalizeEngines(args.engines);
  const count = normalizeCount(args.count);
  const engineText = engines.length > 0 ? formatEngineList(engines) : "—";
  return (
    theme.fg("text", `"${truncate(String(args.query ?? ""), 58)}"`) +
    theme.fg("borderMuted", " · ") +
    theme.fg("dim", engineText) +
    theme.fg("borderMuted", " · ") +
    theme.fg("muted", `${count}/engine`)
  );
}

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

function renderSourceStatus(theme: any, details: SearchDetails): string[] {
  return (details.engines ?? []).map((engine) => {
    const label = ENGINE_LABELS[engine];
    switch (engine) {
      case "exa":
        return details.exaError
          ? theme.fg("error", `${label} ✗`)
          : theme.fg("accent", label) + theme.fg("dim", ` ${details.exaCount ?? 0}`);
      case "brave":
        return details.braveError
          ? theme.fg("error", `${label} ✗`)
          : theme.fg("accent", label) + theme.fg("dim", ` ${details.braveCount ?? 0}`);
      case "metaso":
        return details.metasoError
          ? theme.fg("error", `${label} ✗`)
          : theme.fg("accent", label) + theme.fg("dim", ` ${details.metasoCount ?? 0}`);
      case "jina":
        return details.jinaError
          ? theme.fg("error", `${label} ✗`)
          : theme.fg("accent", label) + theme.fg("dim", ` ${details.jinaCount ?? 0}`);
      case "tavily":
        return details.tavilyError
          ? theme.fg("error", `${label} ✗`)
          : theme.fg("accent", label) + theme.fg("dim", ` ${details.tavilyCount ?? 0}`);
    }
  });
}

export function registerWebSearchTool(pi: ExtensionAPI): void {
  pi.registerTool({
    name: "web_search",
    label: "Web Search",
    description: "Search the web using one or more selected search engines (Exa, Brave, Metaso, Jina, Tavily). Returns aggregated, deduplicated results with titles, URLs, and snippets. Use this when you need to find information on the web.",
    promptSnippet: "Use web_search to search with selected engines and return deduplicated web results.",
    parameters: Type.Object({
      query: Type.String({ description: "The search query" }),
      engines: Type.Array(StringEnum(SEARCH_ENGINES, { description: "Search engine to use" }), {
        minItems: 1,
        uniqueItems: true,
        description: "Search engines to query, e.g. [\"jina\", \"exa\"]",
      }),
      count: Type.Optional(Type.Number({ minimum: 1, description: "Maximum results to request from each selected engine (default: 10)." })),
    }),
    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_search",
        theme,
        status: "pending",
        paramSummary: buildParamSummary(args, theme),
        startedAt: state.startedAt,
      });
    },
    renderResult(result, { expanded, isPartial }, theme, context) {
      const details = result.details as SearchDetails | undefined;
      const state = context.state as { startedAt?: number };
      const startedAt = state.startedAt;
      const paramSummary = buildParamSummary(context.args, theme);
      const expandedParams = expanded ? buildExpandedParams([
        { label: "query", value: context.args?.query, multiline: true },
        { label: "engines", value: context.args?.engines },
        { label: "count", value: context.args?.count },
      ], theme) : undefined;

      if (isPartial) {
        const phase = details?.phase ?? "searching";
        return kylinFrameForResult(context, {
          name: "web_search",
          theme,
          status: "pending",
          paramSummary,
          resultSummary: theme.fg("warning", phase),
          expanded,
          expandedParams,
          startedAt,
        });
      }

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

      const sources = renderSourceStatus(theme, details);
      const expandedBody: string[] = [];
      if (sources.length > 0) expandedBody.push(...sources);
      const preview = extractTextPreview(result.content[0], 40);
      if (preview.lines.length > 0) {
        if (expandedBody.length > 0) expandedBody.push(theme.fg("borderMuted", "─".repeat(40)));
        expandedBody.push(...preview.lines.map((line) => theme.fg("dim", line.replace(/^\s+/, ""))));
      }

      return kylinFrameForResult(context, {
        name: "web_search",
        theme,
        status: "success",
        paramSummary,
        resultSummary:
          theme.fg("success", `${details.totalAfterDedup ?? 0} results`) +
          theme.fg("dim", ` (${details.totalBeforeDedup ?? 0} raw)`),
        expanded,
        expandedParams,
        expandedBody: expanded ? expandedBody : undefined,
        expandedTotalLines: expandedBody.length + Math.max(0, preview.totalLines - preview.lines.length),
        startedAt,
      });
    },
    async execute(_toolCallId, params, signal, onUpdate) {
      const engines = normalizeEngines(params.engines);
      const count = normalizeCount(params.count);
      const query = params.query;
      const serviceKeys: Record<SearchEngine, string | null> = {
        exa: getServiceKey("exa", "EXA_API_KEY"),
        brave: getServiceKey("brave", "BRAVE_API_KEY"),
        metaso: getServiceKey("metaso", "METASO_API_KEY"),
        jina: getServiceKey("jina", "JINA_API_KEY"),
        tavily: getServiceKey("tavily", "TAVILY_API_KEY"),
      };

      if (engines.length === 0) {
        return {
          content: [{ type: "text" as const, text: "Error: No search engines selected. Choose at least one of exa, brave, metaso, jina, tavily." }],
          details: { query, phase: "done", error: "No search engines selected" } as SearchDetails,
        };
      }

      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: { query, engines, count, phase: "done", error: "No API keys configured for selected engines" } as SearchDetails,
        };
      }

      onUpdate?.({
        content: [{ type: "text" as const, text: "Searching..." }],
        details: { query, engines, count, phase: "searching" } as SearchDetails,
      });

      try {
        const searchResults = await Promise.all(engines.map(async (engine) => {
          const apiKey = serviceKeys[engine];
          if (!apiKey) {
            return { engine, results: [] as SearchResult[], error: "No API key" };
          }
          const outcome = await SEARCH_CLIENTS[engine](query, apiKey, count, signal);
          return { engine, ...outcome };
        }));

        const allResults = searchResults.flatMap((outcome) => outcome.results);

        onUpdate?.({
          content: [{ type: "text" as const, text: "Merging results..." }],
          details: { query, engines, count, phase: "merging" } as SearchDetails,
        });

        const merged = dedupSearchResults(allResults);
        const lines: string[] = [];
        for (let i = 0; i < merged.length; i++) {
          const result = merged[i];
          lines.push(`[${i + 1}] ${result.title}`);
          lines.push(`    ${result.url}`);
          if (result.snippet) lines.push(`    ${result.snippet}`);
          lines.push(`    [${result.source}]`);
          lines.push("");
        }

        const details: SearchDetails = {
          query,
          engines,
          count,
          phase: "done",
          totalBeforeDedup: allResults.length,
          totalAfterDedup: merged.length,
        };
        for (const outcome of searchResults) {
          applyEngineOutcome(details, outcome.engine, outcome);
        }

        return {
          content: [{ type: "text" as const, text: lines.join("\n") }],
          details,
        };
      } catch (error) {
        if (isAbortError(error)) {
          return {
            content: [{ type: "text" as const, text: "Search cancelled." }],
            details: { query, engines, count, phase: "done", error: "Cancelled" } as SearchDetails,
          };
        }
        return {
          content: [{ type: "text" as const, text: `Search error: ${errorMessage(error)}` }],
          details: { query, engines, count, phase: "done", error: errorMessage(error) } as SearchDetails,
        };
      }
    },
  });
}