import type {
  AgentToolUpdateCallback,
  AgentToolResult,
  ExtensionAPI,
  ExtensionContext,
} from "@earendil-works/pi-coding-agent";
import { Type, type Static } from "typebox";
import { StringEnum } from "@earendil-works/pi-ai";
import { callTool } from "./mcp.ts";
import { getExaProvider, getProvider } from "./providers.ts";
import { formatResult } from "./utils.ts";

const SEARCH_TIMEOUT_MS = 25_000;
const ADVANCED_SEARCH_TIMEOUT_MS = 45_000;
const FETCH_TIMEOUT_MS = 30_000;
const MAX_FETCH_URLS = 20;
const DEFAULT_NUM_RESULTS = 8;
const DEFAULT_FETCH_MAX_CHARACTERS = 3_000;
const DEFAULT_ADVANCED_CONTEXT_MAX_CHARACTERS = 10_000;
const DEFAULT_PAPER_CONTEXT_MAX_CHARACTERS = 4_000;

const WebSearchSchema = Type.Object({
  query: Type.String({
    description:
      "The information to look up. For best results, describe the ideal page semantically ('blog post comparing React and Vue performance') rather than raw keywords.",
  }),
  queries: Type.Optional(
    Type.Array(Type.String(), {
      description:
        "Optional concise keyword queries (3-6 words each, 2-3 recommended). Used to run a multi-angle search with the Parallel backend; ignored by the Exa backend, which searches `query` alone.",
      minItems: 1,
      maxItems: 5,
    })
  ),
  numResults: Type.Optional(
    Type.Number({
      description: "Number of search results to return (default: 8)",
      minimum: 1,
      maximum: 25,
    })
  ),
});
type WebSearchInput = Static<typeof WebSearchSchema>;

const AdvancedSearchSchema = Type.Object({
  query: Type.String({
    description: "Semantic description of the ideal page to find.",
  }),
  category: Type.Optional(
    StringEnum(
      ["news", "company", "people", "publication", "personal site", "financial report"] as const,
      {
        description:
          "Restrict to a curated index: news (current events), company, people, publication (research papers/preprints), personal site (blogs), financial report (SEC filings, earnings).",
      }
    )
  ),
  includeDomains: Type.Optional(
    Type.Array(Type.String(), {
      description:
        "Only return results from these domains, e.g. ['reuters.com']. Do not combine with the publication category.",
      maxItems: 20,
    })
  ),
  excludeDomains: Type.Optional(
    Type.Array(Type.String(), {
      description: "Exclude results from these domains.",
      maxItems: 20,
    })
  ),
  startPublishedDate: Type.Optional(
    Type.String({
      description: "Only return pages published on or after this date (YYYY-MM-DD).",
    })
  ),
  endPublishedDate: Type.Optional(
    Type.String({
      description: "Only return pages published on or before this date (YYYY-MM-DD).",
    })
  ),
  maxAgeHours: Type.Optional(
    Type.Number({
      description:
        "Only return content crawled within this many hours. Use 0 to force live crawling, or 24-168 for recent news.",
      minimum: 0,
    })
  ),
  type: Type.Optional(
    StringEnum(["auto", "fast", "deep"] as const, {
      description:
        "Search quality/latency tradeoff. 'auto' is the default; 'deep' runs a slower, multi-step search better suited to complex research questions.",
    })
  ),
  numResults: Type.Optional(
    Type.Number({
      description: "Number of search results to return (default: 8)",
      minimum: 1,
      maximum: 25,
    })
  ),
  contextMaxCharacters: Type.Optional(
    Type.Number({
      description:
        "Cap on the LLM-optimized context returned per result (default: 10000). Lower it to save tokens when only key facts are needed.",
      minimum: 100,
    })
  ),
  enableHighlights: Type.Optional(
    Type.Boolean({
      description:
        "Return the most relevant excerpts per page instead of full text. Much cheaper in tokens; on by default for agent use.",
    })
  ),
});
type AdvancedSearchInput = Static<typeof AdvancedSearchSchema>;

const SearchPapersSchema = Type.Object({
  query: Type.String({
    description:
      "Semantic description of the paper or research question, e.g. 'retrieval-augmented generation survey 2025' or 'diffusion models for time series forecasting'.",
  }),
  numResults: Type.Optional(
    Type.Number({
      description: "Number of papers to return (default: 8)",
      minimum: 1,
      maximum: 25,
    })
  ),
  startPublishedDate: Type.Optional(
    Type.String({
      description: "Only return papers published on or after this date (YYYY-MM-DD).",
    })
  ),
  endPublishedDate: Type.Optional(
    Type.String({
      description: "Only return papers published on or before this date (YYYY-MM-DD).",
    })
  ),
  type: Type.Optional(
    StringEnum(["auto", "fast", "deep"] as const, {
      description:
        "Search quality/latency tradeoff. 'deep' runs a slower, multi-step search suited to complex literature reviews.",
    })
  ),
  contextMaxCharacters: Type.Optional(
    Type.Number({
      description:
        "Cap on the returned content per paper (default: 4000). Lower it to save tokens when only titles and abstracts are needed.",
      minimum: 100,
    })
  ),
  enableHighlights: Type.Optional(
    Type.Boolean({
      description: "Return the most relevant excerpts per paper instead of full text (on by default).",
    })
  ),
  enableSummary: Type.Optional(
    Type.Boolean({
      description:
        "Return an LLM-generated summary per paper. Useful for rapid triage of many results; adds latency and tokens.",
    })
  ),
});
type SearchPapersInput = Static<typeof SearchPapersSchema>;

const WebFetchSchema = Type.Object({
  urls: Type.Array(Type.String(), {
    description: "URLs to fetch content from. Batch multiple related URLs in one call.",
    minItems: 1,
    maxItems: MAX_FETCH_URLS,
  }),
  objective: Type.Optional(
    Type.String({
      description:
        "What to extract from the pages, to focus the returned excerpts on the relevant parts (Parallel backend; ignored by Exa).",
    })
  ),
  maxCharacters: Type.Optional(
    Type.Number({
      description: "Maximum characters to extract per page (default: 3000)",
      minimum: 1,
    })
  ),
  fullContent: Type.Optional(
    Type.Boolean({
      description:
        "Return the full page content instead of focused excerpts. Only set for long articles or when exhaustive reading is required (Parallel backend; ignored by Exa).",
    })
  ),
});
type WebFetchInput = Static<typeof WebFetchSchema>;

function sessionId(ctx: ExtensionContext): string {
  try {
    const id = ctx.sessionManager?.getSessionId();
    if (id) return id;
  } catch {
    // Fall back to the working directory when no session is active.
  }
  return ctx.cwd;
}

function modelName(ctx: ExtensionContext): string | undefined {
  return ctx.model?.id ? ctx.model.id.slice(0, 100) : undefined;
}

function validateUrls(urls: string[]): void {
  for (const url of urls) {
    let parsed: URL;
    try {
      parsed = new URL(url);
    } catch {
      throw new Error(`Invalid URL (must start with http:// or https://): ${url}`);
    }
    if (parsed.protocol !== "http:" && parsed.protocol !== "https:") {
      throw new Error(`Invalid URL (must start with http:// or https://): ${url}`);
    }
  }
}

function validateAdvancedSearch(params: AdvancedSearchInput): void {
  const hasDomainFilters = (params.includeDomains?.length || 0) > 0 || (params.excludeDomains?.length || 0) > 0;
  if (params.category === "publication" && hasDomainFilters) {
    throw new Error("The Exa publication index does not support domain filters. Use search_papers with date filters instead.");
  }

  const isEntityCategory = params.category === "company" || params.category === "people";
  if (isEntityCategory && (params.startPublishedDate || params.endPublishedDate || params.excludeDomains?.length)) {
    throw new Error(`The Exa ${params.category} index does not support publication dates or excludeDomains.`);
  }
}

async function webSearch(
  _id: string,
  params: WebSearchInput,
  signal: AbortSignal | undefined,
  onUpdate: AgentToolUpdateCallback | undefined,
  ctx: ExtensionContext
): Promise<AgentToolResult<unknown>> {
  const currentSessionId = sessionId(ctx);
  const currentModel = modelName(ctx);
  const provider = getProvider(currentSessionId);
  const numResults = params.numResults ?? DEFAULT_NUM_RESULTS;
  onUpdate?.({
    content: [{ type: "text", text: `Searching the web with ${provider.label}...` }],
    details: {},
  });

  const raw = await callTool({
    url: provider.url,
    tool: provider.searchTool,
    args: provider.buildSearchArgs({
      query: params.query,
      queries: params.queries,
      numResults,
      sessionId: currentSessionId,
      modelName: currentModel,
    }),
    headers: provider.headers,
    timeoutMs: SEARCH_TIMEOUT_MS,
    signal,
  });
  return formatResult(provider.normalizeOutput(raw, numResults), {
    provider: provider.name,
    query: params.query,
    queries: params.queries,
    numResults,
    model: currentModel,
  });
}

async function webSearchAdvanced(
  _id: string,
  params: AdvancedSearchInput,
  signal: AbortSignal | undefined,
  onUpdate: AgentToolUpdateCallback | undefined,
  ctx: ExtensionContext
): Promise<AgentToolResult<unknown>> {
  const provider = getExaProvider();
  const currentModel = modelName(ctx);
  const numResults = params.numResults ?? DEFAULT_NUM_RESULTS;
  const contextMaxCharacters = params.contextMaxCharacters ?? DEFAULT_ADVANCED_CONTEXT_MAX_CHARACTERS;
  const enableHighlights = params.enableHighlights ?? true;
  onUpdate?.({
    content: [{ type: "text", text: `Running an advanced web search with ${provider.label}...` }],
    details: {},
  });

  validateAdvancedSearch(params);
  const raw = await callTool({
    url: provider.url,
    tool: provider.advancedSearchTool,
    args: provider.buildAdvancedSearchArgs({
      query: params.query,
      category: params.category,
      includeDomains: params.includeDomains,
      excludeDomains: params.excludeDomains,
      startPublishedDate: params.startPublishedDate,
      endPublishedDate: params.endPublishedDate,
      maxAgeHours: params.maxAgeHours,
      type: params.type,
      numResults,
      contextMaxCharacters,
      enableHighlights,
    }),
    headers: provider.headers,
    timeoutMs: ADVANCED_SEARCH_TIMEOUT_MS,
    signal,
  });
  return formatResult(provider.normalizeOutput(raw), {
    provider: provider.name,
    query: params.query,
    category: params.category,
    maxAgeHours: params.maxAgeHours,
    numResults,
    contextMaxCharacters,
    model: currentModel,
  });
}

async function searchPapers(
  _id: string,
  params: SearchPapersInput,
  signal: AbortSignal | undefined,
  onUpdate: AgentToolUpdateCallback | undefined,
  ctx: ExtensionContext
): Promise<AgentToolResult<unknown>> {
  const provider = getExaProvider();
  const currentModel = modelName(ctx);
  const numResults = params.numResults ?? DEFAULT_NUM_RESULTS;
  const contextMaxCharacters = params.contextMaxCharacters ?? DEFAULT_PAPER_CONTEXT_MAX_CHARACTERS;
  const enableHighlights = params.enableHighlights ?? true;
  onUpdate?.({
    content: [{ type: "text", text: `Searching the scientific literature with ${provider.label}...` }],
    details: {},
  });

  const raw = await callTool({
    url: provider.url,
    tool: provider.advancedSearchTool,
    args: provider.buildAdvancedSearchArgs({
      query: params.query,
      category: "publication",
      numResults,
      startPublishedDate: params.startPublishedDate,
      endPublishedDate: params.endPublishedDate,
      type: params.type,
      contextMaxCharacters,
      enableHighlights,
      enableSummary: params.enableSummary,
    }),
    headers: provider.headers,
    timeoutMs: ADVANCED_SEARCH_TIMEOUT_MS,
    signal,
  });
  return formatResult(provider.normalizeOutput(raw), {
    provider: provider.name,
    query: params.query,
    category: "publication",
    numResults,
    contextMaxCharacters,
    model: currentModel,
  });
}

async function webFetch(
  _id: string,
  params: WebFetchInput,
  signal: AbortSignal | undefined,
  onUpdate: AgentToolUpdateCallback | undefined,
  ctx: ExtensionContext
): Promise<AgentToolResult<unknown>> {
  const currentSessionId = sessionId(ctx);
  const currentModel = modelName(ctx);
  const provider = getProvider(currentSessionId);
  const maxCharacters = params.maxCharacters ?? DEFAULT_FETCH_MAX_CHARACTERS;
  onUpdate?.({
    content: [{ type: "text", text: `Fetching ${params.urls.length} URL(s) with ${provider.label}...` }],
    details: {},
  });

  validateUrls(params.urls);
  const raw = await callTool({
    url: provider.url,
    tool: provider.fetchTool,
    args: provider.buildFetchArgs({
      urls: params.urls,
      objective: params.objective,
      maxCharacters,
      fullContent: params.fullContent,
      sessionId: currentSessionId,
      modelName: currentModel,
    }),
    headers: provider.headers,
    timeoutMs: FETCH_TIMEOUT_MS,
    signal,
  });
  return formatResult(provider.normalizeOutput(raw), {
    provider: provider.name,
    urls: params.urls,
    objective: params.objective,
    maxCharacters,
    model: currentModel,
  });
}

export default function (pi: ExtensionAPI): void {
  pi.registerTool({
    name: "web_search",
    label: "Web Search",
    description:
      "Search the web for current information and return clean, ready-to-use content from the top results. Use for facts, news, documentation, people, and companies. The returned excerpts are usually sufficient to answer directly — avoid fetching every result; only read a page in full with web_fetch when excerpts conflict or are insufficient. Describe the ideal page semantically rather than using bare keywords, and pass 2-3 concise `queries` for broad or multi-angle topics (used by the Parallel backend; the Exa backend searches `query` alone). For time-sensitive facts, say so in the query or use web_search_advanced with maxAgeHours.",
    promptSnippet: "Search the web for current information (web_search) and read pages (web_fetch)",
    parameters: WebSearchSchema,
    execute: webSearch,
  });

  pi.registerTool({
    name: "web_search_advanced",
    label: "Advanced Web Search",
    description:
      "Search the web with precise controls: category-restricted indexes (news, company, people, publication, financial report), domain allow/deny lists, publication date ranges, freshness windows (maxAgeHours for breaking news), result count, and token-efficient highlights. Use when you need current, filtered, or academically-grounded results that web_search cannot express. Uses Exa's advanced search index.",
    promptSnippet: "Search the web with filters like categories, domains, dates, and freshness (web_search_advanced)",
    parameters: AdvancedSearchSchema,
    execute: webSearchAdvanced,
  });

  pi.registerTool({
    name: "search_papers",
    label: "Search Scientific Literature",
    description:
      "Search 350M+ scholarly publications: research papers, preprints, and journal articles (arXiv, PubMed, and other repositories). Returns structured metadata per paper — title, authors, publication date, venue — with excerpts or summaries. Use for literature reviews, finding papers on a topic, or locating a specific paper. Narrow by publication year with startPublishedDate/endPublishedDate. Uses Exa's publication index.",
    promptSnippet: "Search scientific papers, preprints, and journal articles (search_papers)",
    parameters: SearchPapersSchema,
    execute: searchPapers,
  });

  pi.registerTool({
    name: "web_fetch",
    label: "Web Fetch",
    description:
      "Fetch the content of one or more public URLs as clean markdown or focused excerpts (up to 20 URLs per call). Use only when search excerpts are insufficient: exact wording, full-page analysis, or reading documentation. Pass an `objective` to focus excerpts on the relevant parts, and set `fullContent` only when the entire page is needed (both are Parallel backend options; the Exa backend caps size with `maxCharacters` only).",
    promptSnippet: "Fetch web pages as clean markdown (web_fetch)",
    parameters: WebFetchSchema,
    execute: webFetch,
  });
}