import { readFile, readdir } from "node:fs/promises";
import { join } from "node:path";
import { loadConfig, loadConfigVersion, writeJsonAtomic, type UltraPaths } from "../config/loader.js";
import { loadOrCreateKey } from "../security/privacy.js";
import { EventStore } from "../telemetry/event-store.js";
import { TelemetryCollector } from "../telemetry/collector.js";
import type { BaseEvent, UltraConfig } from "../types.js";
import { experimentAssignment, experimentExclusionReason, type ExperimentExclusionReason, type ExperimentSafetyInput } from "./assignment.js";
import { experimentStopReason, type ExperimentStopReason } from "./stop-loss.js";
import { effectiveInitial } from "../controller/progressive-widening.js";

export interface ExperimentEnvironment { provider: string; model: string }
export interface ExperimentBaseline { champion: ExperimentEnvironment; challenger: ExperimentEnvironment }

export interface ExperimentRecord {
  schemaVersion: 1;
  experimentId: string;
  name: string;
  status: "running" | "stopped";
  championVersion: string;
  challengerVersion: string;
  allocation: number;
  factorPath: string;
  thresholds: { successDrop: number; costIncrease: number };
  baseline: ExperimentBaseline;
  startedAt: string;
  stoppedAt?: string;
  stopReason?: ExperimentStopReason;
}

export interface ExperimentEvent {
  eventType: "experiment.started" | "experiment.assigned" | "experiment.evaluated" | "experiment.stopped";
  fields: Record<string, string | number | boolean>;
}

export interface ExperimentLifecycleResult { record: ExperimentRecord; event: ExperimentEvent }

export interface ExperimentSampleCounts { championTasks: number; challengerTasks: number; verifiedComparablePairs: number }

export interface ExperimentSample extends ExperimentSampleCounts {
  successDrop: number;
  reworkIncrease: number;
  costIncrease: number;
  qualityBenefit: boolean;
  safetyIncident: boolean;
  sampleCorruption: boolean;
  observed: ExperimentBaseline;
}

export interface MinimumSampleSummary {
  minCohortTasks: number;
  verifiedComparablePairs: number;
  sampleSufficient: boolean;
  confidenceCeiling: "medium" | "high";
  reasons: string[];
}

const CONFIG_METADATA = new Set(["configVersion", "policyVersion", "changeReason", "createdAt"]);
const ID = /^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$/;
const CONFIG_VERSION = /^\d+\.\d+\.\d+(?:[-+][A-Za-z0-9.-]+)?$/;
const STOP_REASONS = new Set<ExperimentStopReason>(["verified-success-drop", "rework-increase", "cost-increase-without-quality-benefit", "safety-incident", "sample-corruption", "provider-drift", "model-drift", "manual", "completed"]);
const MODEL_SECTIONS = new Set(["root", "scout", "boundedWriter", "repair", "deep", "arbitration"]);
type ModelSection = "root" | "scout" | "boundedWriter" | "repair" | "deep" | "arbitration";

function assertId(experimentId: string): void {
  if (!ID.test(experimentId)) throw new Error("Experiment id must be a safe local identifier");
}

function isRecord(value: unknown): value is Record<string, unknown> {
  return Boolean(value) && typeof value === "object" && !Array.isArray(value);
}

function changedPaths(left: unknown, right: unknown, prefix = "", output: string[] = []): string[] {
  if (Object.is(left, right)) return output;
  if (isRecord(left) && isRecord(right)) {
    for (const key of [...new Set([...Object.keys(left), ...Object.keys(right)])].sort()) {
      if (!prefix && CONFIG_METADATA.has(key)) continue;
      changedPaths(left[key], right[key], prefix ? `${prefix}.${key}` : key, output);
    }
    return output;
  }
  if (JSON.stringify(left) !== JSON.stringify(right)) output.push(prefix || "<root>");
  return output;
}

async function readJson(file: string): Promise<unknown | undefined> {
  try { return JSON.parse(await readFile(file, "utf8")) as unknown; }
  catch (error) { if ((error as NodeJS.ErrnoException).code === "ENOENT") return undefined; throw error; }
}

function assertEnvironment(environment: ExperimentEnvironment): void {
  if (!environment.provider.trim() || !environment.model.trim()) throw new Error("Experiment baseline provider and model are required");
}

function environment(model: string, section: string): ExperimentEnvironment {
  const separator = model.indexOf("/");
  if (separator <= 0 || separator === model.length - 1) throw new Error(`Invalid experiment ${section} model ${model}`);
  return { provider: model.slice(0, separator), model: model.slice(separator + 1) };
}

function modelSection(factorPath: string): ModelSection {
  const candidate = factorPath.split(".", 1)[0] ?? "root";
  return MODEL_SECTIONS.has(candidate) ? candidate as ModelSection : "root";
}

function baselineFor(champion: UltraConfig, challenger: UltraConfig, factorPath: string): ExperimentBaseline {
  const section = modelSection(factorPath);
  return { champion: environment(champion[section].model, section), challenger: environment(challenger[section].model, section) };
}

function assertRecord(value: unknown): asserts value is ExperimentRecord {
  if (!isRecord(value) || value.schemaVersion !== 1 || typeof value.experimentId !== "string" || typeof value.name !== "string") throw new Error("Invalid experiment record");
  if (value.status !== "running" && value.status !== "stopped") throw new Error("Invalid experiment status");
  if (typeof value.championVersion !== "string" || typeof value.challengerVersion !== "string" || typeof value.factorPath !== "string" || typeof value.startedAt !== "string") throw new Error("Invalid experiment record");
  if (typeof value.allocation !== "number" || !isRecord(value.thresholds) || !isRecord(value.baseline)) throw new Error("Invalid experiment record");
  if (!ID.test(value.experimentId) || !value.name.trim() || !value.factorPath || !CONFIG_VERSION.test(value.championVersion) || !CONFIG_VERSION.test(value.challengerVersion)) throw new Error("Invalid experiment record");
  if (!Number.isFinite(value.allocation) || value.allocation < 0 || value.allocation > 1 || typeof value.thresholds.successDrop !== "number" || typeof value.thresholds.costIncrease !== "number" || !Number.isFinite(value.thresholds.successDrop) || !Number.isFinite(value.thresholds.costIncrease) || value.thresholds.successDrop < 0 || value.thresholds.costIncrease < 0) throw new Error("Invalid experiment record");
  if (!isRecord(value.baseline.champion) || !isRecord(value.baseline.challenger) || typeof value.baseline.champion.provider !== "string" || typeof value.baseline.champion.model !== "string" || typeof value.baseline.challenger.provider !== "string" || typeof value.baseline.challenger.model !== "string") throw new Error("Invalid experiment record");
  assertEnvironment({ provider: value.baseline.champion.provider, model: value.baseline.champion.model });
  assertEnvironment({ provider: value.baseline.challenger.provider, model: value.baseline.challenger.model });
  if (Number.isNaN(Date.parse(value.startedAt))) throw new Error("Invalid experiment record");
  if (value.status === "stopped" && (typeof value.stoppedAt !== "string" || !STOP_REASONS.has(value.stopReason as ExperimentStopReason))) throw new Error("Invalid experiment record");
  if (value.status === "running" && (value.stoppedAt !== undefined || value.stopReason !== undefined)) throw new Error("Invalid experiment record");
}

function recordFields(record: ExperimentRecord): Record<string, string | number | boolean> {
  return {
    experimentId: record.experimentId,
    name: record.name,
    status: record.status,
    championVersion: record.championVersion,
    challengerVersion: record.challengerVersion,
    allocation: record.allocation,
    factorPath: record.factorPath,
    startedAt: record.startedAt,
    ...(record.stoppedAt ? { stoppedAt: record.stoppedAt } : {}),
    ...(record.stopReason ? { stopReason: record.stopReason } : {}),
  };
}

export async function readExperiment(paths: UltraPaths, experimentId: string): Promise<ExperimentRecord> {
  assertId(experimentId);
  const value = await readJson(join(paths.experiments, `${experimentId}.json`));
  if (!value) throw new Error(`Unknown experiment ${experimentId}`);
  assertRecord(value);
  if (value.experimentId !== experimentId) throw new Error("Experiment record id mismatch");
  return value;
}

export async function listExperiments(paths: UltraPaths): Promise<ExperimentRecord[]> {
  return (await Promise.all((await readdir(paths.experiments, { withFileTypes: true }))
    .filter((entry) => entry.isFile() && entry.name.endsWith(".json"))
    .map((entry) => readExperiment(paths, entry.name.slice(0, -5)))))
    .sort((left, right) => left.experimentId.localeCompare(right.experimentId));
}

export async function recordExperimentLifecycle(paths: UltraPaths, result: ExperimentLifecycleResult, now = new Date(result.record.stoppedAt ?? result.record.startedAt)): Promise<void> {
  if (result.event.eventType !== "experiment.evaluated" && (await new EventStore(paths.events).all()).some((event) => event.experimentId === result.record.experimentId && event.eventType === result.event.eventType)) return;
  const config = await loadConfigVersion(paths, result.record.championVersion);
  const runId = `experiment:${result.record.experimentId}`;
  const collector = new TelemetryCollector(paths.events, paths.database, paths.hmacKey, config, config.profile);
  try { await collector.record({ sessionId: "system", taskId: runId, runId, experimentId: result.record.experimentId }, result.event.eventType, result.event.fields, now); }
  finally { collector.close(); }
}

export async function startExperiment(paths: UltraPaths, input: { experimentId: string; name: string; championVersion: string; challengerVersion: string }, now = new Date()): Promise<ExperimentLifecycleResult> {
  assertId(input.experimentId);
  if (!input.name.trim() || input.name.length > 200) throw new Error("Experiment name must be 1-200 characters");
  if (await readJson(join(paths.experiments, `${input.experimentId}.json`))) throw new Error(`Experiment ${input.experimentId} already exists`);
  const activeChampion = await loadConfig(paths);
  const running = (await listExperiments(paths)).find((record) => record.status === "running");
  if (running) throw new Error(`Experiment already running: ${running.experimentId}`);
  const [champion, challenger] = await Promise.all([loadConfigVersion(paths, input.championVersion), loadConfigVersion(paths, input.challengerVersion)]);
  if (activeChampion.configVersion !== input.championVersion) throw new Error(`Experiment champion ${input.championVersion} does not match active champion ${activeChampion.configVersion}`);
  if (champion.profile !== challenger.profile) throw new Error("Experiment cannot change privacy profile");
  const factors = changedPaths(champion, challenger);
  if (factors.length !== 1) throw new Error(`Experiment must change exactly one primary factor; found ${factors.length}`);
  if (factors[0] === "scout.initial" && effectiveInitial(champion, champion.policy) === effectiveInitial(challenger, challenger.policy)) throw new Error("Experiment has no effective scout fan-out difference under the configured policy");
  const baseline = baselineFor(champion, challenger, factors[0]!);
  assertEnvironment(baseline.champion);
  assertEnvironment(baseline.challenger);
  const allocation = champion.experiment.challengerAllocation;
  experimentAssignment("validation", "validation", input.experimentId, allocation);
  const record: ExperimentRecord = {
    schemaVersion: 1,
    experimentId: input.experimentId,
    name: input.name.trim(),
    status: "running",
    championVersion: input.championVersion,
    challengerVersion: input.challengerVersion,
    allocation,
    factorPath: factors[0]!,
    thresholds: { successDrop: champion.experiment.stopLossSuccessDrop, costIncrease: champion.experiment.stopLossCostIncrease },
    baseline,
    startedAt: now.toISOString(),
  };
  await writeJsonAtomic(join(paths.experiments, `${record.experimentId}.json`), record);
  return { record, event: { eventType: "experiment.started", fields: recordFields(record) } };
}

export interface ExperimentAssignmentResult {
  cohort: "champion" | "challenger";
  configVersion: string;
  excluded: boolean;
  exclusionReason?: ExperimentExclusionReason | "inactive-experiment";
  event: ExperimentEvent;
}

export async function assignExperiment(paths: UltraPaths, experimentId: string, taskFingerprint: string, safety: ExperimentSafetyInput): Promise<ExperimentAssignmentResult> {
  const record = await readExperiment(paths, experimentId);
  const exclusionReason = record.status === "running" ? experimentExclusionReason(safety) : "inactive-experiment";
  const cohort = exclusionReason ? "champion" : experimentAssignment(await loadOrCreateKey(paths.hmacKey), taskFingerprint, experimentId, record.allocation);
  const configVersion = cohort === "challenger" ? record.challengerVersion : record.championVersion;
  const fields = { experimentId, cohort, configVersion, excluded: Boolean(exclusionReason), ...(exclusionReason ? { exclusionReason } : {}) };
  return { cohort, configVersion, excluded: Boolean(exclusionReason), ...(exclusionReason ? { exclusionReason } : {}), event: { eventType: "experiment.assigned", fields } };
}

function assertCounts(counts: ExperimentSampleCounts): void {
  for (const value of [counts.championTasks, counts.challengerTasks, counts.verifiedComparablePairs]) if (!Number.isSafeInteger(value) || value < 0) throw new Error("Experiment sample counts must be non-negative integers");
}

export function minimumSampleSummary(counts: ExperimentSampleCounts): MinimumSampleSummary {
  assertCounts(counts);
  const minCohortTasks = Math.min(counts.championTasks, counts.challengerTasks);
  const reasons = [...(minCohortTasks < 8 ? ["cohort-tasks<8"] : []), ...(counts.verifiedComparablePairs < 5 ? ["comparable-pairs<5"] : [])];
  return { minCohortTasks, verifiedComparablePairs: counts.verifiedComparablePairs, sampleSufficient: reasons.length === 0, confidenceCeiling: reasons.length ? "medium" : "high", reasons };
}

export async function stopExperiment(paths: UltraPaths, experimentId: string, reason: ExperimentStopReason, now = new Date()): Promise<ExperimentLifecycleResult> {
  if (!STOP_REASONS.has(reason)) throw new Error("Invalid experiment stop reason");
  const current = await readExperiment(paths, experimentId);
  if (current.status === "stopped") {
    if (current.stopReason !== reason) throw new Error(`Experiment ${experimentId} already stopped: ${current.stopReason}`);
    return { record: current, event: { eventType: "experiment.stopped", fields: recordFields(current) } };
  }
  const record: ExperimentRecord = { ...current, status: "stopped", stoppedAt: now.toISOString(), stopReason: reason };
  await writeJsonAtomic(join(paths.experiments, `${experimentId}.json`), record);
  return { record, event: { eventType: "experiment.stopped", fields: recordFields(record) } };
}

function sampleFields(sample: ExperimentSample, summary: MinimumSampleSummary, providerDrift: boolean, modelDrift: boolean): Record<string, string | number | boolean> {
  return {
    championTasks: sample.championTasks,
    challengerTasks: sample.challengerTasks,
    verifiedComparablePairs: sample.verifiedComparablePairs,
    minCohortTasks: summary.minCohortTasks,
    sampleSufficient: summary.sampleSufficient,
    confidenceCeiling: summary.confidenceCeiling,
    successDrop: sample.successDrop,
    reworkIncrease: sample.reworkIncrease,
    costIncrease: sample.costIncrease,
    qualityBenefit: sample.qualityBenefit,
    safetyIncident: sample.safetyIncident,
    sampleCorruption: sample.sampleCorruption,
    providerDrift,
    modelDrift,
  };
}

export async function evaluateExperiment(paths: UltraPaths, experimentId: string, sample: ExperimentSample, now = new Date()): Promise<ExperimentLifecycleResult> {
  const record = await readExperiment(paths, experimentId);
  if (record.status !== "running") throw new Error(`Experiment ${experimentId} is not running`);
  const summary = minimumSampleSummary(sample);
  for (const value of [sample.successDrop, sample.reworkIncrease, sample.costIncrease]) if (!Number.isFinite(value)) throw new Error("Experiment deltas must be finite");
  const arms = ["champion", "challenger"] as const;
  const providerDrift = arms.some((arm) => sample.observed[arm].provider !== record.baseline[arm].provider);
  const modelDrift = arms.some((arm) => sample.observed[arm].model !== record.baseline[arm].model);
  const fields = { experimentId, ...sampleFields(sample, summary, providerDrift, modelDrift) };
  const stopReason = experimentStopReason({ ...sample, providerDrift, modelDrift }, record.thresholds);
  if (!stopReason) return { record, event: { eventType: "experiment.evaluated", fields } };
  const stopped = await stopExperiment(paths, experimentId, stopReason, now);
  return { ...stopped, event: { ...stopped.event, fields: { ...stopped.event.fields, ...fields } } };
}

function last(events: readonly BaseEvent[], type: string): BaseEvent | undefined {
  return [...events].reverse().find((event) => event.eventType === type);
}

function rate(tasks: readonly BaseEvent[][], predicate: (events: readonly BaseEvent[]) => boolean): number {
  return tasks.length ? tasks.filter(predicate).length / tasks.length : 0;
}

function isModelCall(event: BaseEvent): boolean {
  return (event.eventType === "model.call.completed" || event.eventType === "model.call.failed") && typeof event.provider === "string" && typeof event.model === "string";
}

function factorCalls(events: readonly BaseEvent[], section: ModelSection): BaseEvent[] {
  const selected: BaseEvent[] = [];
  const addNext = (index: number, roles: readonly string[], stop: (event: BaseEvent) => boolean = () => false) => {
    const call = events.slice(index + 1).find((event) => stop(event) || (isModelCall(event) && roles.includes(String(event.role))));
    if (call && !stop(call) && !selected.some((event) => event.eventId === call.eventId)) selected.push(call);
  };
  if (section === "arbitration") {
    events.forEach((event, index) => {
      if (event.eventType === "escalation.triggered" && event.reason === "contradictory-scout-evidence") addNext(index, ["root", "writer"], (candidate) => candidate.eventType === "repair.attempted" || candidate.eventType === "result.completed" || candidate.eventType === "escalation.triggered");
    });
    return selected;
  }
  if (section === "boundedWriter") {
    events.forEach((event, index) => {
      if (event.eventType === "route.validated" && event.finalTopology === "swarm") addNext(index, ["writer"], (candidate) => candidate.eventType === "repair.attempted" || candidate.eventType === "result.completed" || (candidate.eventType === "escalation.triggered" && candidate.reason === "contradictory-scout-evidence"));
    });
    return selected;
  }
  if (section === "repair") {
    events.forEach((event, index) => {
      if (event.eventType !== "repair.attempted") return;
      const priorMarker = events.slice(0, index).findLast((candidate) => candidate.eventType === "repair.attempted" || candidate.eventType === "escalation.triggered");
      if (priorMarker?.eventType !== "escalation.triggered" || priorMarker.reason !== "same-fingerprint-twice") addNext(index, ["writer"], (candidate) => candidate.eventType === "repair.attempted" || candidate.eventType === "result.completed" || candidate.eventType === "escalation.triggered");
    });
    return selected;
  }
  if (section === "deep") {
    events.forEach((event, index) => {
      if (event.eventType === "route.validated" && event.finalTopology === "deep") addNext(index, ["writer"], (candidate) => candidate.eventType === "repair.attempted" || candidate.eventType === "result.completed" || candidate.eventType === "escalation.triggered");
      if (event.eventType === "escalation.triggered" && (event.reason === "same-fingerprint-twice" || event.reason === "warroom-no-convergence")) addNext(index, ["writer"], (candidate) => candidate.eventType === "result.completed" || candidate.eventType === "escalation.triggered");
    });
    return selected;
  }
  const arbitration = new Set(factorCalls(events, "arbitration").map((event) => event.eventId));
  return events.filter((event) => isModelCall(event) && (section === "root"
    ? (event.role === "root" && !arbitration.has(event.eventId)) || String(event.nodeId).endsWith(".branches.lead")
    : event.role === "scout" && !String(event.nodeId).endsWith(".branches.lead")));
}

function observedEnvironments(cohorts: { champion: readonly BaseEvent[][]; challenger: readonly BaseEvent[][] }, baseline: ExperimentBaseline, factorPath: string): ExperimentBaseline {
  const observed = structuredClone(baseline);
  const section = modelSection(factorPath);
  for (const arm of ["champion", "challenger"] as const) {
    const environments: ExperimentEnvironment[] = [];
    for (const event of cohorts[arm].flatMap((events) => factorCalls(events, section))) {
      const provider = event.provider as string;
      const recordedModel = event.model as string;
      const model = recordedModel.startsWith(`${provider}/`) ? recordedModel.slice(provider.length + 1) : recordedModel;
      if (!environments.some((entry) => entry.provider === provider && entry.model === model)) environments.push({ provider, model });
    }
    observed[arm] = environments.find((entry) => entry.provider !== baseline[arm].provider || entry.model !== baseline[arm].model) ?? observed[arm];
  }
  return observed;
}

export async function evaluateRunningExperiments(paths: UltraPaths, now = new Date()): Promise<ExperimentLifecycleResult[]> {
  const records = await Promise.all((await readdir(paths.experiments, { withFileTypes: true }))
    .filter((entry) => entry.isFile() && entry.name.endsWith(".json"))
    .map((entry) => readExperiment(paths, entry.name.slice(0, -5))));
  const running = records.filter((record) => record.status === "running");
  if (!running.length) return [];
  const store = new EventStore(paths.events);
  const allEvents = await store.all();
  const output: ExperimentLifecycleResult[] = [];
  for (const record of running) {
    const events = allEvents.filter((event) => event.experimentId === record.experimentId);
    const lastEvaluation = [...events].reverse().find((event) => event.eventType === "experiment.evaluated" || event.eventType === "experiment.stopped");
    const grouped = new Map<string, BaseEvent[]>();
    for (const event of events) if (!event.taskId.startsWith("experiment:")) grouped.set(event.taskId, [...(grouped.get(event.taskId) ?? []), event]);
    const tasks = [...grouped.values()].filter((taskEvents) => last(taskEvents, "result.completed"));
    const newestEvidence = tasks.flat().sort((left, right) => right.timestamp.localeCompare(left.timestamp))[0];
    if (!newestEvidence || (lastEvaluation && newestEvidence.timestamp <= lastEvaluation.timestamp)) continue;

    let sampleCorruption = false;
    const cohorts = { champion: [] as BaseEvent[][], challenger: [] as BaseEvent[][] };
    for (const taskEvents of tasks) {
      const assignments = taskEvents.filter((event) => event.eventType === "experiment.assigned").map((event) => event.cohort);
      const unique = [...new Set(assignments)];
      if (unique.length !== 1 || (unique[0] !== "champion" && unique[0] !== "challenger")) { sampleCorruption = true; continue; }
      cohorts[unique[0]].push(taskEvents);
    }
    const verified = (taskEvents: readonly BaseEvent[]) => { const result = last(taskEvents, "result.completed"); return result?.success === true && result.verified === true; };
    const reworked = (taskEvents: readonly BaseEvent[]) => taskEvents.some((event) => (event.eventType === "user.feedback" || event.eventType === "feedback.explicit") && (event.kind === "fixed" || event.kind === "bad"));
    const credits = (taskEvents: readonly BaseEvent[]) => taskEvents.reduce((total, event) => total + ((event.eventType === "model.call.completed" || event.eventType === "model.call.failed") && typeof event.creditsEstimated === "number" ? event.creditsEstimated : 0), 0);
    const championSuccess = rate(cohorts.champion, verified);
    const challengerSuccess = rate(cohorts.challenger, verified);
    const championCredits = cohorts.champion.length ? cohorts.champion.reduce((total, task) => total + credits(task), 0) / cohorts.champion.length : 0;
    const challengerCredits = cohorts.challenger.length ? cohorts.challenger.reduce((total, task) => total + credits(task), 0) / cohorts.challenger.length : 0;
    const comparable = cohorts.champion.length > 0 && cohorts.challenger.length > 0;
    const sample: ExperimentSample = {
      championTasks: cohorts.champion.length,
      challengerTasks: cohorts.challenger.length,
      verifiedComparablePairs: Math.min(cohorts.champion.length, cohorts.challenger.length),
      successDrop: comparable ? championSuccess - challengerSuccess : 0,
      reworkIncrease: comparable ? rate(cohorts.challenger, reworked) - rate(cohorts.champion, reworked) : 0,
      costIncrease: comparable && championCredits > 0 ? (challengerCredits - championCredits) / championCredits : 0,
      qualityBenefit: comparable && challengerSuccess > championSuccess,
      safetyIncident: cohorts.challenger.flat().some((event) => event.eventType === "safety.incident" || event.safetyIncident === true),
      sampleCorruption,
      observed: observedEnvironments(cohorts, record.baseline, record.factorPath),
    };
    const result = await evaluateExperiment(paths, record.experimentId, sample, now);
    await recordExperimentLifecycle(paths, result, now);
    output.push(result);
  }
  return output;
}