import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { Type } from "@sinclair/typebox";
import { truncateHead, DEFAULT_MAX_BYTES, DEFAULT_MAX_LINES } from "@earendil-works/pi-coding-agent";
import { BraveSearchProvider } from "../src/search.js";
import { WebFetchCache } from "../src/cache.js";
import { stripHtml, isInternalUrl } from "../src/utils.js";
import { withRetry } from "../src/retry.js";

// ─── Tool Schemas ────────────────────────────────────────────────────────────

const WebSearchParams = Type.Object({
  query: Type.String({ description: "Search query" }),
  count: Type.Optional(Type.Number({ default: 5, description: "Number of results (max 20)" })),
});

const WebFetchParams = Type.Object({
  url: Type.String({ description: "Full URL to fetch" }),
  prompt: Type.String({ description: "What to look for in this page. The tool extracts only content relevant to this prompt." }),
});

async function fetchAndExtract(url: string, signal?: AbortSignal): Promise<{ text: string; finalUrl: string; contentType: string }> {
  const doFetch = async () => {
    const timeoutSignal = AbortSignal.timeout(30000);
    const combinedSignal = signal ? AbortSignal.any([signal, timeoutSignal]) : timeoutSignal;

    if (isInternalUrl(url)) {
      throw new Error(`Blocked fetch to internal/resource URL: ${url}`);
    }

    const response = await fetch(url, {
      headers: { "User-Agent": "pi-web-research/1.0" },
      redirect: "follow",
      signal: combinedSignal,
    });

    const finalUrl = response.url;
    const contentType = response.headers.get("Content-Type") ?? "";

    if (!response.ok) {
      throw new Error(`HTTP ${response.status} for ${url}`);
    }

    // Reject HTTPS → HTTP cross-origin redirects BEFORE reading body
    if (url.startsWith("https://") && finalUrl.startsWith("http://")) {
      const origHost = new URL(url).hostname;
      const finalHost = new URL(finalUrl).hostname;
      if (origHost !== finalHost) {
        throw new Error(`Rejected cross-origin redirect from HTTPS to HTTP: ${url} → ${finalUrl}`);
      }
    }

    // Allowlist: only text, JSON, XML
    const isAllowed = contentType.startsWith("text/") ||
                      contentType.includes("application/json") ||
                      contentType.includes("application/xml") ||
                      contentType.includes("application/xhtml");

    // Empty Content-Type is allowed to fall through — some servers omit it; assume HTML
    if (!isAllowed && contentType !== "") {
      throw new Error(`Cannot extract text from ${contentType}. Use a specialized tool.`);
    }

    if (contentType.includes("application/json")) {
      const json = await response.json();
      return { text: JSON.stringify(json, null, 2), finalUrl, contentType };
    }

    const html = await response.text();
    const text = stripHtml(html);

    return { text, finalUrl, contentType };
  };

  return withRetry(doFetch, { maxRetries: 2, baseDelayMs: 1500 });
}

// ─── Extension Factory ───────────────────────────────────────────────────────

export default function (pi: ExtensionAPI) {
  const cache = new WebFetchCache();

  pi.on("session_start", (_event, _ctx) => {
    pi.registerTool({
      name: "web_search",
      label: "Web Search",
      description: "Search the web and return ranked results. Use when the user asks about current practices, recent tools, libraries, or anything where an up-to-date answer matters. Prefer over answering from memory for anything time-sensitive (versions, pricing, news, API changes).",
      promptSnippet: "Search the web for current information, docs, or best practices",
      promptGuidelines: [
        "Use web_search when the user asks about current practices, recent tools, libraries, or anything where an up-to-date answer matters.",
        "Use proactively on broad 'how should I...' or 'what's the best way to...' questions.",
        "Prefer web_search over answering from memory for anything time-sensitive (versions, pricing, news, API changes).",
        "Do NOT use for basic syntax, standard library features, or well-established patterns.",
        "Run 2 parallel broad searches with different phrasings first. Only use domain scoping (e.g., 'topic site:developer.apple.com') as a fallback when those results are thin, outdated, or low-quality.",
      ],
      parameters: WebSearchParams,

      async execute(_toolCallId, params, signal, onUpdate) {
        const apiKey = process.env.BRAVE_API_KEY;
        if (!apiKey) {
          throw new Error("BRAVE_API_KEY not set. Sign up for a key at https://brave.com/search/api and export BRAVE_API_KEY=<your-key> in your shell.");
        }

        onUpdate?.({ content: [{ type: "text", text: `Searching: ${params.query}...` }], details: {} });

        const provider = new BraveSearchProvider(apiKey);
        const searchSignal = signal ? AbortSignal.any([signal, AbortSignal.timeout(30000)]) : AbortSignal.timeout(30000);
        const results = await provider.search(params.query, params.count ?? 5, searchSignal);

        if (results.length === 0) {
          return {
            content: [{ type: "text", text: `No results found for "${params.query}". Try rephrasing.` }],
            details: { query: params.query, resultCount: 0 },
          };
        }

        let text = results.map((r, i) =>
          `${i + 1}. ${r.title}\n   ${r.url}\n   ${r.description}`
        ).join("\n\n");

        const truncated = truncateHead(text, { maxLines: DEFAULT_MAX_LINES, maxBytes: DEFAULT_MAX_BYTES });

        return {
          content: [{ type: "text", text: truncated.content }],
          details: { query: params.query, resultCount: results.length, truncated: truncated.truncated },
        };
      },
    });

    pi.registerTool({
      name: "web_fetch",
      label: "Web Fetch",
      description: "Fetch a URL and extract content relevant to a specific question or topic. Always provide a focused prompt to keep context usage low and quality high.",
      promptSnippet: "Fetch and extract focused content from a web page",
      promptGuidelines: [
        "Use when the user provides a URL explicitly.",
        "Use after web_search when search snippets are insufficient.",
        "Always provide a focused prompt — this keeps context usage low and quality high.",
        "Do NOT use for URLs that require authentication or are known to block bots.",
      ],
      parameters: WebFetchParams,

      async execute(_toolCallId, params, signal, onUpdate) {
        onUpdate?.({ content: [{ type: "text", text: `Fetching ${params.url}...` }], details: {} });

        // Check extracted cache first
        const cached = cache.getExtracted(params.url, params.prompt);
        if (cached) {
          return {
            content: [{ type: "text", text: cached.text }],
            details: { url: params.url, prompt: params.prompt, cached: true, cacheAgeMs: cached.ageMs },
          };
        }

        // Check raw cache
        let rawText = cache.getRaw(params.url);
        let finalUrl = params.url;

        if (!rawText) {
          const result = await fetchAndExtract(params.url, signal);
          rawText = result.text;
          finalUrl = result.finalUrl;
          cache.setRaw(finalUrl, rawText);
        }

        // Short-circuit: if text is short, return raw
        if (rawText.length < 4096) {
          const output = `Source: ${finalUrl}\n\n${rawText}`;
          cache.setExtracted(finalUrl, params.prompt, output);
          return {
            content: [{ type: "text", text: output }],
            details: { url: finalUrl, prompt: params.prompt, rawLength: rawText.length },
          };
        }

        // v1: raw truncation. Future: call a cheap external API (e.g. Anthropic Haiku
        // via direct fetch) to extract only paragraphs relevant to params.prompt.
        const truncated = truncateHead(rawText, { maxLines: DEFAULT_MAX_LINES, maxBytes: DEFAULT_MAX_BYTES });

        let text = `Source: ${finalUrl}\n\n${truncated.content}`;
        if (params.url !== finalUrl) {
          text = `Redirected from: ${params.url}\n${text}`;
        }

        cache.setExtracted(finalUrl, params.prompt, text);

        return {
          content: [{ type: "text", text }],
          details: { url: finalUrl, prompt: params.prompt, truncated: truncated.truncated },
        };
      },
    });
  });
}