/**
 * Build a decision certificate from a seed JSONL file.
 * `{ id?, input, expect: { <question>: value }, locale? }`.
 */

import type { AiDecisionQuestion } from "../declare.ts";
import {
  applyTemperature,
  learnThenTest,
  questionHash,
  type DecisionCalibrator,
  type DecisionCertSlice,
  type DecisionDeciderCert,
  type DecisionLabel,
} from "./certificate.ts";
import type { DecisionResponse } from "./provider.ts";

/** One seed row. */
export interface DecisionSeedCase {
  readonly id?: string;
  readonly input: unknown;
  readonly expect: Readonly<Record<string, unknown>>;
  readonly locale?: string;
}

/** One provider call used while certifying. */
export type DecisionEvaluate = (input: unknown) => Promise<DecisionResponse>;

/** Printed for one question after certify. */
export interface CertifyQuestionReport {
  readonly accuracy: number;
  readonly ece: number;
  readonly coverage: number;
  readonly p50Ms: number;
  readonly p95Ms: number;
  readonly cost: number;
}

/** Certificate plus the comparison row. */
export interface CertifySeedResult {
  readonly cert: DecisionDeciderCert;
  readonly report: Readonly<Record<string, CertifyQuestionReport>>;
  /** True when the certificate is an expiring alias. */
  readonly unpinned: boolean;
}

/**
 * Parse decision seed JSONL. Blank lines are skipped.
 *
 * @param text - File contents
 */
export function parseDecisionSeed(text: string): DecisionSeedCase[] {
  const cases: DecisionSeedCase[] = [];
  for (const line of text.split("\n")) {
    const trimmed = line.trim();
    if (!trimmed) continue;
    const row = JSON.parse(trimmed) as {
      id?: unknown;
      input?: unknown;
      expect?: unknown;
      locale?: unknown;
    };
    if (!row.expect || typeof row.expect !== "object" || Array.isArray(row.expect)) {
      throw new TypeError("decision seed row needs expect");
    }
    cases.push({
      ...(typeof row.id === "string" ? { id: row.id } : {}),
      input: row.input,
      expect: row.expect as Readonly<Record<string, unknown>>,
      ...(typeof row.locale === "string" ? { locale: row.locale } : {}),
    });
  }
  return cases;
}

/** One labeled probability used to fit a slice. */
interface FitRow {
  readonly probs: readonly number[];
  readonly labelIndex: number;
  readonly noul?: number;
  readonly booleanLabel?: boolean;
  /** Inverse propensity. Audit rows count more than a full review. */
  readonly weight: number;
}

/**
 * Certificate from a seed file. Calls the provider, fits a calibrator per
 * question, and keeps a threshold only when Learn-then-Test passes.
 *
 * @param options - Decision, risk, seed, and provider
 */
export async function certifySeed(options: {
  readonly model: string;
  readonly pinned?: boolean;
  readonly expiresAt?: number;
  readonly maxError: number;
  readonly delta?: number;
  readonly jsonl: string;
  readonly ask: Readonly<Record<string, AiDecisionQuestion>>;
  readonly evaluate: DecisionEvaluate;
  readonly now?: () => number;
}): Promise<CertifySeedResult> {
  const cases = parseDecisionSeed(options.jsonl);
  const buckets = new Map<string, FitRow[]>();
  const samples = new Map<string, { correct: boolean; confidence: number }[]>();
  const latencies: number[] = [];
  let cost = 0;
  let resolved = options.model;
  const now = options.now ?? Date.now;
  for (const row of cases) {
    const started = now();
    const response = await options.evaluate(row.input);
    latencies.push(Math.max(0, now() - started));
    cost += response.usage.cost ?? 0;
    if (response.model) resolved = response.model;
    const locale = row.locale ?? "";
    for (const [id, expected] of Object.entries(row.expect)) {
      const question = options.ask[id];
      if (!question) continue;
      const answers = response.answers;
      const answer = Array.isArray(answers)
        ? undefined
        : (answers as Readonly<Record<string, unknown>>)[id];
      const fit = fitRow(question, answer, expected);
      if (!fit) continue;
      const key = `${id}\0${locale}`;
      const list = buckets.get(key) ?? [];
      list.push(fit);
      buckets.set(key, list);
      const confidence =
        fit.noul !== undefined ? Math.max(fit.noul, 1 - fit.noul) : Math.max(...fit.probs, 0);
      const correct =
        fit.booleanLabel !== undefined
          ? fit.noul !== undefined && fit.noul >= 0.5 === fit.booleanLabel
          : fit.probs[fit.labelIndex] === Math.max(...fit.probs);
      const bag = samples.get(id) ?? [];
      bag.push({ correct, confidence });
      samples.set(id, bag);
    }
  }
  const entry = entryFromBuckets(
    resolved,
    options.ask,
    buckets,
    options.maxError,
    options.delta ?? 0.1,
  );
  const pinned = options.pinned !== false;
  const cert: DecisionDeciderCert = {
    ...entry,
    pinned,
    ...(options.expiresAt !== undefined ? { expiresAt: options.expiresAt } : {}),
    certifiedAt: Date.now(),
  };
  const report: Record<string, CertifyQuestionReport> = {};
  for (const [id, rows] of samples) {
    const threshold = cert.questions[id]?.[""]?.threshold;
    const covered =
      threshold === undefined
        ? 0
        : rows.filter((row) => row.confidence >= threshold).length / rows.length;
    report[id] = {
      accuracy: rows.filter((row) => row.correct).length / rows.length,
      ece: expectedCalibrationError(rows),
      coverage: covered,
      p50Ms: percentile(latencies, 0.5),
      p95Ms: percentile(latencies, 0.95),
      cost,
    };
  }
  return { cert, report, unpinned: !pinned };
}

/**
 * The same certificate, built from persisted labels that already carry a score.
 *
 * @param options - Model, questions, labels, and risk
 */
export function certifyLabels(options: {
  readonly model: string;
  readonly maxError: number;
  readonly delta?: number;
  readonly ask: Readonly<Record<string, AiDecisionQuestion>>;
  readonly labels: readonly DecisionLabel[];
}): DecisionDeciderCert {
  const groups = new Map<string, DecisionLabel[]>();
  for (const label of options.labels) {
    if (label.raw === undefined) continue;
    if (options.model && label.model !== options.model) continue;
    const key = `${label.question}\0${label.locale ?? ""}`;
    const list = groups.get(key) ?? [];
    list.push(label);
    groups.set(key, list);
  }
  const questions: Record<string, Record<string, DecisionCertSlice>> = {};
  const model = options.labels.find((label) => label.model)?.model ?? options.model;
  const delta = options.delta ?? 0.1;
  for (const [key, rows] of groups) {
    const split = key.indexOf("\0");
    const id = key.slice(0, split);
    const locale = key.slice(split + 1);
    const question = options.ask[id];
    if (!question) continue;
    const fitted = rows.flatMap((row) => {
      const fit = fitRow(question, row.raw, row.value);
      return fit ? [{ ...fit, weight: propensityWeight(row.propensity) }] : [];
    });
    const slice = sliceFromRows(question, fitted, options.maxError, delta);
    if (!slice) continue;
    const bucket = questions[id] ?? {};
    bucket[locale] = slice;
    questions[id] = bucket;
  }
  return { model, pinned: true, certifiedAt: Date.now(), questions };
}

function entryFromBuckets(
  model: string,
  ask: Readonly<Record<string, AiDecisionQuestion>>,
  buckets: ReadonlyMap<string, readonly FitRow[]>,
  maxError: number,
  delta: number,
): DecisionDeciderCert {
  const questions: Record<string, Record<string, DecisionCertSlice>> = {};
  for (const [key, rows] of buckets) {
    const split = key.indexOf("\0");
    const id = key.slice(0, split);
    const locale = key.slice(split + 1);
    const question = ask[id];
    if (!question) continue;
    const slice = sliceFromRows(question, rows, maxError, delta);
    if (!slice) continue;
    const bucket = questions[id] ?? {};
    bucket[locale] = slice;
    questions[id] = bucket;
  }
  return { model, pinned: true, questions };
}

function sliceFromRows(
  question: AiDecisionQuestion,
  rows: readonly FitRow[],
  maxError: number,
  delta: number,
): DecisionCertSlice | undefined {
  const calibrator = fitCalibrator(question, rows);
  const scored = rows.map((row) => scoreFit(question, row, calibrator));
  const threshold = learnThenTest(scored, maxError, delta);
  if (threshold === null) return undefined;
  const errors = scored.filter((row) => row.score >= threshold && row.loss > 0).length;
  const accepted = scored.filter((row) => row.score >= threshold).length;
  return {
    hash: questionHash(question),
    calibrator,
    threshold,
    metrics: { labels: rows.length, accepted, errors, maxError, delta },
  };
}

function fitCalibrator(question: AiDecisionQuestion, rows: readonly FitRow[]): DecisionCalibrator {
  if (question.kind === "boolean") return fitPlatt(rows);
  return { kind: "temperature", t: fitTemperature(rows) };
}

function fitTemperature(rows: readonly FitRow[]): number {
  let bestT = 1;
  let best = Number.POSITIVE_INFINITY;
  for (let step = 1; step <= 20; step++) {
    const t = step / 10;
    let nll = 0;
    for (const row of rows) {
      const scaled = applyTemperature(row.probs, t);
      nll -= row.weight * Math.log(Math.max(scaled[row.labelIndex] ?? 1e-12, 1e-12));
    }
    if (nll < best) {
      best = nll;
      bestT = t;
    }
  }
  return bestT;
}

function fitPlatt(rows: readonly FitRow[]): DecisionCalibrator {
  let best = { a: 1, b: 0, nll: Number.POSITIVE_INFINITY };
  for (let a = -2; a <= 2.01; a += 0.5) {
    for (let b = -2; b <= 2.01; b += 0.5) {
      let nll = 0;
      for (const row of rows) {
        const p = platt(row.noul ?? 0.5, a, b);
        const y = row.booleanLabel ? 1 : 0;
        nll -=
          row.weight *
          (y * Math.log(Math.max(p, 1e-12)) + (1 - y) * Math.log(Math.max(1 - p, 1e-12)));
      }
      if (nll < best.nll) best = { a, b, nll };
    }
  }
  return { kind: "platt", a: best.a, b: best.b };
}

function platt(noul: number, a: number, b: number): number {
  const p = Math.min(1, Math.max(0, noul));
  const logit = Math.log(Math.max(p, 1e-12) / Math.max(1 - p, 1e-12));
  return 1 / (1 + Math.exp(-(a * logit + b)));
}

function scoreFit(
  question: AiDecisionQuestion,
  row: FitRow,
  calibrator: DecisionCalibrator,
): { readonly score: number; readonly loss: number } {
  if (question.kind === "boolean" && calibrator.kind !== "temperature") {
    const p = platt(
      row.noul ?? 0.5,
      calibrator.kind === "platt" ? calibrator.a : 1,
      calibrator.kind === "platt" ? calibrator.b : 0,
    );
    const predicted = p >= 0.5;
    return { score: Math.max(p, 1 - p), loss: predicted === row.booleanLabel ? 0 : 1 };
  }
  const t = calibrator.kind === "temperature" ? calibrator.t : 1;
  const scaled = applyTemperature(row.probs, t);
  let best = 0;
  for (let i = 1; i < scaled.length; i++) {
    if ((scaled[i] ?? 0) > (scaled[best] ?? 0)) best = i;
  }
  return { score: scaled[best] ?? 0, loss: best === row.labelIndex ? 0 : 1 };
}

/**
 * Provider answers wrap the distribution in `probabilities`.
 * A stored label writes that same distribution as `raw`.
 *
 * @param record - Answer or stored raw
 */
function percentile(values: readonly number[], p: number): number {
  if (values.length === 0) return 0;
  const sorted = [...values].sort((a, b) => a - b);
  const index = Math.min(sorted.length - 1, Math.max(0, Math.ceil(p * sorted.length) - 1));
  return sorted[index] ?? 0;
}

function expectedCalibrationError(
  rows: readonly { readonly confidence: number; readonly correct: boolean }[],
): number {
  if (rows.length === 0) return 0;
  const bins = Array.from({ length: 10 }, () => ({ n: 0, confidence: 0, correct: 0 }));
  for (const row of rows) {
    const index = Math.min(9, Math.floor(row.confidence * 10));
    const bin = bins[index];
    if (!bin) continue;
    bin.n += 1;
    bin.confidence += row.confidence;
    bin.correct += row.correct ? 1 : 0;
  }
  return bins.reduce((sum, bin) => {
    if (bin.n === 0) return sum;
    return sum + (bin.n / rows.length) * Math.abs(bin.correct / bin.n - bin.confidence / bin.n);
  }, 0);
}

function storedDistribution(record: Record<string, unknown>): unknown {
  if (record.probabilities !== undefined) return record.probabilities;
  return record;
}

function fitRow(
  question: AiDecisionQuestion,
  answer: unknown,
  expected: unknown,
): FitRow | undefined {
  const record = answer && typeof answer === "object" ? (answer as Record<string, unknown>) : {};
  if (record.type === "refusal" || record.type === "malformed") return undefined;
  if (question.kind === "boolean") {
    if (typeof expected !== "boolean") return undefined;
    return {
      probs: [],
      labelIndex: expected ? 1 : 0,
      noul:
        typeof record.probability === "number"
          ? record.probability
          : typeof record.noul === "number"
            ? record.noul
            : 0,
      booleanLabel: expected,
      weight: 1,
    };
  }
  const keys =
    question.kind === "choice"
      ? [...Object.keys(question.options), "none_of_these"]
      : [...question.levels];
  const labelIndex = keys.indexOf(String(expected));
  if (labelIndex < 0) return undefined;
  const raw = storedDistribution(record);
  const probs = keys.map((key, index) => {
    if (raw && typeof raw === "object" && !Array.isArray(raw)) {
      const value = (raw as Record<string, unknown>)[key];
      return typeof value === "number" ? value : 0;
    }
    return Array.isArray(raw) && typeof raw[index] === "number" ? raw[index] : 0;
  });
  return { probs, labelIndex, weight: 1 };
}

function propensityWeight(propensity: number): number {
  if (!(propensity > 0)) return 1;
  return 1 / propensity;
}