import type { ProviderModelConfig } from "@earendil-works/pi-coding-agent";
import { canonicalModelId } from "./models.js";
import { getCatalogText } from "./catalog-http.js";

export const PARLEY_DOCS_ORIGIN = "https://parley-docs.mit.edu";
export type Rates = Pick<ProviderModelConfig["cost"], "input" | "output" | "cacheRead" | "cacheWrite">;
export interface DocMetadata {
  prices: Record<string, Rates>; // normalized display name
  limits: Record<string, { contextWindow: number; maxTokens: number }>; // canonical model id
}

export function nameKey(name: string): string {
  return name.toLowerCase().replace(/[^a-z0-9]/g, "");
}

function htmlText(text: string): string {
  return text.replace(/<[^>]*>/g, "").replace(/&amp;/g, "&").replace(/&nbsp;/g, " ")
    .replace(/&#36;/g, "$").trim();
}

/** Read literal cells only. Never import, evaluate, or execute remote JavaScript. */
export function documentationTables(source: string): string[][][] {
  if (/<table[\s>]/i.test(source)) {
    return [...source.matchAll(/<table\b[^>]*>([\s\S]*?)<\/table>/gi)].map(table =>
      [...table[1].matchAll(/<tr\b[^>]*>([\s\S]*?)<\/tr>/gi)].map(row =>
        [...row[1].matchAll(/<t[hd]\b[^>]*>([\s\S]*?)<\/t[hd]>/gi)].map(cell => htmlText(cell[1]))));
  }
  // Docusaurus compiles Markdown tables to jsx(n.table, {children: ...}).
  // Only the known literal shape is accepted; a different build format falls back to cache.
  return source.split(/\b\w+\.table,\{children:/).slice(1).map(table =>
    table.split(/\b\w+\.tr,\{children:/).slice(1).map(row =>
      row.split(/\b\w+\.t[hd],\{children:/).slice(1).map(cell => {
        const literal = /^(?:\(0,\w+\.jsx\)\(\w+\.(?:strong|code|em),\{children:)?("(?:[^"\\]|\\.)*")/.exec(cell);
        if (!literal) return "";
        try { return JSON.parse(literal[1]) as string; } catch { return ""; }
      })));
}

function rate(text: string, cache = false): number | undefined {
  if (cache && text === "N/A") return 0;
  if (!/^\$\d+(?:\.\d+)?$/.test(text)) return undefined;
  const n = Number(text.slice(1));
  return Number.isFinite(n) && n <= 10000 ? n : undefined;
}

function tokens(text: string): number | undefined {
  const match = /^(\d+(?:\.\d+)?)([KM])?$/.exec(text);
  if (!match) return undefined;
  const n = Number(match[1]) * (match[2] === "M" ? 1000000 : match[2] === "K" ? 1000 : 1);
  return Number.isSafeInteger(n) && n >= 1024 && n <= 100000000 ? n : undefined;
}

export function parseDocMetadata(pricingSource: string, modelsSource: string): DocMetadata {
  const prices: DocMetadata["prices"] = {};
  const limits: DocMetadata["limits"] = {};
  for (const [header, ...rows] of documentationTables(pricingSource)) {
    if (header?.join("|") !== "Model|Input Rate|Output Rate|Cache Read|Cache Write") continue;
    for (const row of rows) {
      if (row.length !== 5) continue;
      const [name, inputText, outputText, readText, writeText] = row;
      const input = rate(inputText), output = rate(outputText), cacheRead = rate(readText, true), cacheWrite = rate(writeText, true);
      if (!name || input === undefined || output === undefined || cacheRead === undefined || cacheWrite === undefined) continue;
      // The docs combine Gemini Flash versions in one row; expand only this explicit form.
      const group = /^Gemini (\d+\.\d+(?: \/ \d+\.\d+)+) Flash$/.exec(name);
      const names = group ? group[1].split(" / ").map(version => `Gemini ${version} Flash`) : [name];
      for (const label of names) prices[nameKey(label)] = { input, output, cacheRead, cacheWrite };
    }
  }
  for (const [header, ...rows] of documentationTables(modelsSource)) {
    if (header?.join("|") !== "Model ID|Name|Context Window|Max Output") continue;
    for (const row of rows) {
      if (row.length !== 4) continue;
      const id = canonicalModelId(row[0]);
      const contextWindow = tokens(row[2]), maxTokens = tokens(row[3]);
      if (id && contextWindow && maxTokens && maxTokens <= contextWindow) limits[id] = { contextWindow, maxTokens };
    }
  }
  if (!Object.keys(prices).length || !Object.keys(limits).length) throw new Error("Unrecognized Parley documentation tables");
  return { prices, limits };
}

/** Cache is data too: whitelist fields and numbers rather than trusting JSON objects. */
export function validateDocMetadata(value: unknown): DocMetadata | undefined {
  if (!value || typeof value !== "object") return undefined;
  const doc = value as DocMetadata;
  if (!doc.prices || !doc.limits || typeof doc.prices !== "object" || typeof doc.limits !== "object") return undefined;
  const prices: DocMetadata["prices"] = {}, limits: DocMetadata["limits"] = {};
  for (const [name, c] of Object.entries(doc.prices)) {
    if (!/^[a-z0-9]{1,160}$/.test(name) || !c || typeof c !== "object") continue;
    if (![c.input, c.output, c.cacheRead, c.cacheWrite].every(n => typeof n === "number" && Number.isFinite(n) && n >= 0 && n <= 10000)) continue;
    prices[name] = { input: c.input, output: c.output, cacheRead: c.cacheRead, cacheWrite: c.cacheWrite };
  }
  for (const [id, l] of Object.entries(doc.limits)) {
    if (canonicalModelId(id) !== id || !l || typeof l !== "object") continue;
    if (![l.contextWindow, l.maxTokens].every(n => Number.isSafeInteger(n) && n >= 1024 && n <= 100000000) || l.maxTokens > l.contextWindow) continue;
    limits[id] = { contextWindow: l.contextWindow, maxTokens: l.maxTokens };
  }
  return Object.keys(prices).length && Object.keys(limits).length ? { prices, limits } : undefined;
}

function assetUrl(path: string, origin: string): URL {
  const url = new URL(path, origin);
  if (url.origin !== new URL(origin).origin || !/^\/assets\/js\/[\w.~\-]+\.js$/.test(url.pathname) || url.search || url.hash) {
    throw new Error("Unrecognized documentation asset URL");
  }
  return url;
}

export function docChunkPath(main: string, runtime: string, page: "cost-guidance" | "models"): string {
  const marker = `"@site/docs/${page}.md"`;
  const pos = main.indexOf(marker);
  if (pos < 0) throw new Error("Missing documentation module");
  const registration = main.slice(Math.max(0, pos - 200), pos);
  const chunkId = [...registration.matchAll(/\.e\((\d+)\)/g)].at(-1)?.[1];
  if (!chunkId) throw new Error("Missing documentation chunk");
  const parts = [...runtime.matchAll(new RegExp(`\\b${chunkId}:"([a-f0-9]+)"`, "g"))].map(m => m[1]);
  if (parts.length !== 2) throw new Error("Unrecognized documentation manifest");
  return `/assets/js/${parts[0]}.${parts[1]}.js`;
}

export async function fetchDocMetadata(signal: AbortSignal, origin = PARLEY_DOCS_ORIGIN): Promise<DocMetadata> {
  // Credentials are NEVER sent to the docs host. Public static assets contain the
  // tables even when the HTML shell waits for a browser authentication check.
  const html = await getCatalogText(new URL("/cost-guidance", origin), signal);
  if (/<table[\s>]/i.test(html)) {
    return parseDocMetadata(html, await getCatalogText(new URL("/models", origin), signal));
  }
  const scripts = [...html.matchAll(/<script\b[^>]*\bsrc="([^"]+)"/g)].map(m => m[1]);
  const mainPath = scripts.find(s => /\/main\.[\w]+\.js$/.test(s));
  const runtimePath = scripts.find(s => /\/runtime~main\.[\w]+\.js$/.test(s));
  if (!mainPath || !runtimePath) throw new Error("Unrecognized documentation page");
  const [main, runtime] = await Promise.all([
    getCatalogText(assetUrl(mainPath, origin), signal), getCatalogText(assetUrl(runtimePath, origin), signal),
  ]);
  const [pricing, models] = await Promise.all(["cost-guidance", "models"].map(page =>
    getCatalogText(assetUrl(docChunkPath(main, runtime, page as "cost-guidance" | "models"), origin), signal)));
  return parseDocMetadata(pricing, models);
}