import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { Type } from "typebox";
import { AnySearchClient, AnySearchError, type JsonObject, type SearchRequest } from "../src/client.js";

const MAX_OUTPUT_CHARS = 12_000;
const querySchema = Type.String({ minLength: 1, maxLength: 1_000, description: "A single web search query." });
const commonSearchParameters = {
  max_results: Type.Optional(Type.Integer({ minimum: 1, maximum: 20, description: "Results per query (default: 5)." })),
  tag: Type.Optional(Type.String({ minLength: 1, maxLength: 200, description: "An AnySearch tag or sub-domain returned by get_sub_domains." })),
  zone: Type.Optional(Type.Union([Type.Literal("cn"), Type.Literal("intl")], { description: "Search zone." })),
  language: Type.Optional(Type.String({ minLength: 1, maxLength: 40, description: "Preferred result language." })),
  params: Type.Optional(Type.Object({}, { additionalProperties: true, description: "JSON object of sub_domain_params for the selected tag." })),
  include_content: Type.Optional(Type.Boolean({ description: "Include result content when available (default: false)." })),
  fresh: Type.Optional(Type.Boolean({ description: "Bypass the 60-second local cache (default: false)." })),
};

const webSearchParameters = Type.Union([
  Type.Object({ query: querySchema, ...commonSearchParameters }, { additionalProperties: false }),
  Type.Object({ queries: Type.Array(querySchema, { minItems: 1, maxItems: 4, description: "One to four independent queries." }), ...commonSearchParameters }, { additionalProperties: false }),
]);

export default function registerAnySearch(pi: ExtensionAPI): void {
  const client = new AnySearchClient();

  pi.registerTool({
    name: "web_search",
    label: "AnySearch Web Search",
    description: "Search the web through AnySearch. Provide exactly one of query or queries (1-4). Result snippets are capped at 500 characters; content is included only with include_content and capped at 2,000 characters; total output is capped at 12,000 characters.",
    promptSnippet: "Search the web with AnySearch, including AnySearch vertical tags.",
    promptGuidelines: [
      "Use get_sub_domains before web_search when a vertical or domain-specific search could help, then pass its returned sub-domain in web_search.tag and its documented parameters in web_search.params.",
      "Use web_search with one precise query when possible; use queries only for up to four genuinely independent searches.",
    ],
    parameters: webSearchParameters,
    async execute(_toolCallId, params, signal) {
      const queries = "query" in params ? [params.query] : params.queries;
      const requestBase = {
        maxResults: params.max_results ?? 5,
        tag: params.tag,
        zone: params.zone,
        language: params.language,
        params: params.params as JsonObject | undefined,
        includeContent: params.include_content ?? false,
        fresh: params.fresh ?? false,
      };

      if (queries.length === 1) {
        const response = await client.search({ query: queries[0]!, ...requestBase }, signal);
        return {
          content: [{ type: "text", text: formatSearchResult(queries[0]!, response.items, requestBase.includeContent) }],
          details: {
            cached: response.cached,
            queryCount: 1,
            successfulQueries: 1,
            failedQueries: [],
            resultCount: response.items.length,
            requestIds: response.requestId ? [response.requestId] : [],
          },
        };
      }

      const outcomes = await mapConcurrent(queries, 3, async (query) => {
        try {
          return { query, response: await client.search({ query, ...requestBase }, signal) };
        } catch (error) {
          return { query, error: safeToolError(error) };
        }
      });
      const successes = outcomes.filter((outcome): outcome is { query: string; response: Awaited<ReturnType<AnySearchClient["search"]>> } => "response" in outcome);
      if (successes.length === 0) {
        throw new AnySearchError(`All AnySearch queries failed: ${outcomes.map((outcome) => `${outcome.query}: ${"error" in outcome ? outcome.error : "unknown error"}`).join("; ")}`);
      }
      const failed = outcomes.filter((outcome): outcome is { query: string; error: string } => "error" in outcome);
      const sections = successes.map(({ query, response }) => formatSearchResult(query, response.items, requestBase.includeContent));
      if (failed.length > 0) sections.unshift(`> Warning: ${failed.length} of ${queries.length} searches failed: ${failed.map((item) => item.query).join(", ")}.`);
      return {
        content: [{ type: "text", text: truncateTotal(sections.join("\n\n")) }],
        details: {
          cached: successes.every((item) => item.response.cached),
          queryCount: queries.length,
          successfulQueries: successes.length,
          failedQueries: failed.map((item) => item.query),
          resultCount: successes.reduce((count, item) => count + item.response.items.length, 0),
          requestIds: successes.flatMap((item) => item.response.requestId ? [item.response.requestId] : []),
        },
      };
    },
  });

  pi.registerTool({
    name: "get_sub_domains",
    label: "AnySearch Sub-domains",
    description: "Ask AnySearch which official vertical-search sub-domains and parameters apply to up to five domains. Use this before web_search for vertical searches. Returns an official Markdown capability table, capped at 12,000 characters.",
    promptSnippet: "Discover AnySearch vertical-search sub-domains and their parameters.",
    promptGuidelines: [
      "Use get_sub_domains before web_search for vertical or domain-specific searches; use only the returned capability table rather than inventing tags or parameter names.",
    ],
    parameters: Type.Object({
      domains: Type.Array(Type.String({ minLength: 1, maxLength: 200 }), { minItems: 1, maxItems: 5, description: "Up to five vertical-search domains to look up." }),
    }, { additionalProperties: false }),
    async execute(_toolCallId, params, signal) {
      const response = await client.getSubDomains(params.domains, signal);
      return {
        content: [{ type: "text", text: truncateTotal(response.markdown) }],
        details: {
          cached: response.cached,
          domainCount: params.domains.length,
          requestIds: response.requestId ? [response.requestId] : [],
        },
      };
    },
  });
}

function formatSearchResult(query: string, items: Awaited<ReturnType<AnySearchClient["search"]>>["items"], includeContent: boolean): string {
  const lines = [`## AnySearch results: ${query}`];
  if (items.length === 0) return `${lines[0]}\n\nNo results returned.`;
  for (const [index, item] of items.entries()) {
    const url = safeUrl(item.url);
    lines.push(url ? `${index + 1}. [${escapeLabel(item.title)}](${url})` : `${index + 1}. ${escapeLabel(item.title)}`);
    if (item.snippet) lines.push(`   ${truncateField(item.snippet, 500, "snippet")}`);
    if (includeContent && item.content) lines.push(`\n   ${truncateField(item.content, 2_000, "content")}`);
  }
  return truncateTotal(lines.join("\n"));
}

function truncateField(text: string, limit: number, label: string): string {
  return text.length <= limit ? text : `${text.slice(0, limit)}… [${label} truncated at ${limit} characters]`;
}

function truncateTotal(text: string): string {
  if (text.length <= MAX_OUTPUT_CHARS) return text;
  const marker = `\n\n… [output truncated at ${MAX_OUTPUT_CHARS} characters]`;
  return `${text.slice(0, MAX_OUTPUT_CHARS - marker.length)}${marker}`;
}

function escapeLabel(value: string): string {
  return value.replace(/[\[\]]/g, "\\$&");
}

function safeUrl(value: string): string | undefined {
  try {
    const parsed = new URL(value);
    return parsed.protocol === "https:" || parsed.protocol === "http:" ? parsed.href : undefined;
  } catch {
    return undefined;
  }
}

function safeToolError(error: unknown): string {
  if (error instanceof Error) {
    return error.message
      .replace(/Bearer\s+\S+/gi, "Bearer [redacted]")
      .replace(/\bas_(?:sk_)?[A-Za-z0-9_-]{6,}\b/g, "[redacted]");
  }
  return "unknown error";
}

async function mapConcurrent<T, R>(values: T[], limit: number, mapper: (value: T) => Promise<R>): Promise<R[]> {
  const output = new Array<R>(values.length);
  let next = 0;
  const workers = Array.from({ length: Math.min(limit, values.length) }, async () => {
    while (next < values.length) {
      const index = next++;
      output[index] = await mapper(values[index]!);
    }
  });
  await Promise.all(workers);
  return output;
}