import type { AskAnswer, PortfolioDocument } from "./types";

export interface LlmConfig {
  /** Env-like bag; defaults to process.env when available. */
  env?: Record<string, string | undefined>;
  timeoutMs?: number;
  /** Override model id (else a small fast default per provider). */
  model?: string;
  subject?: string;
}

type Provider = { name: string; model: string; call: (system: string, user: string, signal: AbortSignal) => Promise<string> };

function readEnv(cfg: LlmConfig): Record<string, string | undefined> {
  if (cfg.env) return cfg.env;
  return typeof process !== "undefined" && process.env ? (process.env as Record<string, string | undefined>) : {};
}

function openAICompatible(url: string, key: string, model: string): Provider["call"] {
  return async (system, user, signal) => {
    const res = await fetch(url, {
      method: "POST",
      signal,
      headers: { "content-type": "application/json", authorization: `Bearer ${key}` },
      body: JSON.stringify({
        model,
        temperature: 0.1,
        max_tokens: 380,
        messages: [
          { role: "system", content: system },
          { role: "user", content: user },
        ],
      }),
    });
    if (!res.ok) throw new Error(`llm_${res.status}`);
    const json = (await res.json()) as { choices?: { message?: { content?: string } }[] };
    return json.choices?.[0]?.message?.content?.trim() || "";
  };
}

/** Pick the first configured provider. Returns null when no key is present. */
export function resolveProvider(cfg: LlmConfig = {}): Provider | null {
  const env = readEnv(cfg);
  if (env.AI_GATEWAY_API_KEY) {
    const model = cfg.model || env.APS_MODEL || "openai/gpt-4o-mini";
    return { name: "ai-gateway", model, call: openAICompatible("https://ai-gateway.vercel.sh/v1/chat/completions", env.AI_GATEWAY_API_KEY, model) };
  }
  if (env.OPENAI_API_KEY) {
    const model = cfg.model || env.APS_MODEL || "gpt-4o-mini";
    return { name: "openai", model, call: openAICompatible("https://api.openai.com/v1/chat/completions", env.OPENAI_API_KEY, model) };
  }
  if (env.XAI_API_KEY) {
    const model = cfg.model || env.APS_MODEL || "grok-3-mini";
    return { name: "xai", model, call: openAICompatible("https://api.x.ai/v1/chat/completions", env.XAI_API_KEY, model) };
  }
  if (env.ANTHROPIC_API_KEY) {
    const key = env.ANTHROPIC_API_KEY;
    const model = cfg.model || env.APS_MODEL || "claude-3-5-haiku-latest";
    return {
      name: "anthropic",
      model,
      call: async (system, user, signal) => {
        const res = await fetch("https://api.anthropic.com/v1/messages", {
          method: "POST",
          signal,
          headers: { "content-type": "application/json", "x-api-key": key, "anthropic-version": "2023-06-01" },
          body: JSON.stringify({ model, max_tokens: 380, temperature: 0.1, system, messages: [{ role: "user", content: user }] }),
        });
        if (!res.ok) throw new Error(`llm_${res.status}`);
        const json = (await res.json()) as { content?: { text?: string }[] };
        return json.content?.map((c) => c.text || "").join("").trim() || "";
      },
    };
  }
  return null;
}

function contextBlock(docs: PortfolioDocument[]): string {
  return docs
    .map((d, i) => `[${i + 1}] ${d.title} (${d.category}${d.url ? ` · ${d.url}` : ""})\n${[d.summary, ...(d.facts || [])].join("\n")}`)
    .join("\n\n");
}

/**
 * Rewrite a deterministic answer with an LLM, grounded strictly in the same
 * cited documents. Returns null on any failure, timeout, or ungrounded output
 * so callers can keep the deterministic answer.
 */
export async function refineWithLlm(base: AskAnswer, docs: PortfolioDocument[], cfg: LlmConfig = {}): Promise<AskAnswer | null> {
  const provider = resolveProvider(cfg);
  if (!provider || !docs.length) return null;
  const subject = cfg.subject || "the portfolio owner";
  const system = [
    `You answer questions about ${subject} using ONLY the numbered sources provided.`,
    "Rules: never add facts, numbers, employers, dates or links that are not in the sources.",
    "Cite every sentence with [n] markers that refer to the sources. 2–4 short sentences, plain text, no markdown headings.",
    "If the sources do not answer the question, say so briefly.",
  ].join(" ");
  const user = `Question: ${base.question}\n\nSources:\n${contextBlock(docs)}`;
  const ctrl = new AbortController();
  const timer = setTimeout(() => ctrl.abort(), cfg.timeoutMs ?? 6000);
  const t0 = Date.now();
  try {
    const text = await provider.call(system, user, ctrl.signal);
    const refs = [...text.matchAll(/\[(\d+)\]/g)].map((m) => Number(m[1]));
    if (!text || !refs.length || refs.some((n) => n < 1 || n > docs.length)) return null;
    return { ...base, answer: text, mode: "llm", model: `${provider.name}:${provider.model}`, tookMs: base.tookMs + (Date.now() - t0) };
  } catch {
    return null;
  } finally {
    clearTimeout(timer);
  }
}