#!/usr/bin/env node

import { createHash } from "node:crypto";
import { readFile } from "node:fs/promises";
import process from "node:process";

const PROTOCOL_ID = "llm-atlas-deepseek-chat-task-bootstrap-crn-v1";
const CONDITIONS = ["s0_eos", "s1_eos", "s0_period", "s1_period"];
const EXPECTED_HASHES = {
  sampling: "ea0607f2b197fac3f794655c1538ce1f9a1cb072d637eb31d35573682e311809",
  evaluation: "82b2fc5d1f854a7e0cd7aba7c70ff74d733d24221522a4f92b962478c76177ab",
  replay: "6519947e2fa4327c1ba2cc506b0861f172a6bdbd4f2d6787447edaf8fbcac508",
  reproduction: "63ed39e5dcdbc2a30e516172f3657dfa453ff3b23d5bd5c49c734939242a70f5",
  analysis: "9ab17561ced930a668c082141f7c6e013cbda70e42b09de63d41f1b82c01a6ae",
  manifest: "6313e70536c464fe598a93035576752418f08016dfd60ac246437c3b43bf2ae1",
};
const DEFAULT_PATHS = {
  sampling: "src/data/deepseek-v2-lite-chat-task-bootstrap-crn.json",
  evaluation: "src/data/deepseek-v2-lite-chat-task-bootstrap-crn-eval.json",
  replay: "src/data/deepseek-v2-lite-chat-task-bootstrap-crn-replay.json",
  reproduction: "src/data/deepseek-v2-lite-chat-task-bootstrap-crn-reproduction.json",
  analysis: "src/data/deepseek-v2-lite-chat-task-bootstrap-crn-analysis.json",
  manifest: "research/DEEPSEEK_V2_LITE_CHAT_TASK_BOOTSTRAP_MANIFEST.json",
};

function sha256(value) {
  return createHash("sha256").update(value).digest("hex");
}

function canonicalIntegerArray(values) {
  return `[${values.map((value) => value.toString()).join(",")}]`;
}

function tapeUint64(tape, sourceId, step) {
  const payload = `${PROTOCOL_ID}\0uniform\0${tape}\0${sourceId}\0${step}`;
  return createHash("sha256").update(payload).digest().readBigUInt64BE(0);
}

function assert(value, message) {
  if (!value) {
    throw new Error(message);
  }
}

const loaded = {};
for (const [name, path] of Object.entries(DEFAULT_PATHS)) {
  const bytes = await readFile(path);
  assert(
    sha256(bytes) === EXPECTED_HASHES[name],
    `${name}: file SHA-256 differs`,
  );
  loaded[name] = JSON.parse(bytes);
}
const { sampling, evaluation, replay, reproduction, analysis, manifest } = loaded;
for (const [name, payload] of Object.entries(loaded)) {
  assert(payload.protocol_id === PROTOCOL_ID, `${name}: protocol ID differs`);
}

assert(sampling.execution_mode === "formal", "sampling: not formal mode");
assert(sampling.sources.length === 64, "sampling: source count differs");
assert(sampling.summary.runs === 88, "sampling: run count differs");
assert(sampling.summary.outputs === 352, "sampling: output count differs");
assert(sampling.summary.natural_eos === 343, "sampling: EOS count differs");
assert(sampling.summary.budget_truncated === 9, "sampling: truncation count differs");
assert(
  sampling.summary.torch_rng_unchanged_runs === 88,
  "sampling: RNG nonconsumption count differs",
);
assert(
  sampling.source_contract.prompt_hash_audit.exact === 256,
  "sampling: prompt audit differs",
);

const manifestBySource = new Map(
  manifest.sources.map((source) => [source.id, source]),
);
const samplingKeys = new Set();
let uniformOutputHashesExact = 0;
let diagnosticSources = 0;
for (const source of sampling.sources) {
  const frozen = manifestBySource.get(source.id);
  assert(frozen, `${source.id}: absent from manifest`);
  assert(source.domain === frozen.domain, `${source.id}: domain differs`);
  assert(
    source.within_domain_index === frozen.domain_index,
    `${source.id}: domain index differs`,
  );
  const expectedTapes = [...frozen.main_tapes, ...frozen.diagnostic_tapes];
  const observedTapes = source.runs.map((run) => run.tape_label);
  assert(
    JSON.stringify(observedTapes) === JSON.stringify(expectedTapes),
    `${source.id}: tape assignment differs`,
  );
  diagnosticSources += source.runs.length === 4;
  for (const run of source.runs) {
    assert(run.torch_rng_unchanged, `${source.id}/${run.tape_label}: RNG changed`);
    assert(
      JSON.stringify(run.outputs.map((output) => output.condition))
        === JSON.stringify(CONDITIONS),
      `${source.id}/${run.tape_label}: condition order differs`,
    );
    for (const output of run.outputs) {
      const key = `${source.id}\0${run.tape_label}\0${output.condition}`;
      assert(!samplingKeys.has(key), `${key}: duplicate sampling key`);
      samplingKeys.add(key);
      assert(
        output.prompt_token_ids_sha256
          === frozen.chat_generation_prompt_token_ids_sha256[output.condition],
        `${key}: prompt hash differs`,
      );
      const uniforms = Array.from(
        { length: output.uniform_steps_consumed },
        (_, step) => tapeUint64(run.tape_label, source.id, step),
      );
      assert(
        sha256(canonicalIntegerArray(uniforms))
          === output.uniform_uint64_prefix_sha256,
        `${key}: uniform prefix hash differs`,
      );
      uniformOutputHashesExact += 1;
    }
  }
}
assert(diagnosticSources === 8, "sampling: diagnostic source count differs");
assert(samplingKeys.size === 352, "sampling: unique key count differs");
assert(uniformOutputHashesExact === 352, "sampling: uniform audit differs");

assert(evaluation.rows.length === 352, "evaluation: row count differs");
const evaluationKeys = new Set(
  evaluation.rows.map(
    (row) => `${row.source_id}\0${row.tape_label}\0${row.condition}`,
  ),
);
assert(evaluationKeys.size === 352, "evaluation: duplicate keys");
assert(
  [...evaluationKeys].every((key) => samplingKeys.has(key)),
  "evaluation: cell key absent from sampling",
);
assert(
  evaluation.summary.main_t0.outputs === 256,
  "evaluation: T0 output count differs",
);
assert(
  evaluation.summary.main_t0.by_domain.code.fixed_budget_success === 59,
  "evaluation: code success count differs",
);
assert(
  evaluation.summary.main_t0.by_domain.math.fixed_budget_success === 71,
  "evaluation: math success count differs",
);

assert(replay.execution_mode === "replay", "replay: mode differs");
assert(replay.sources.length === 16, "replay: source count differs");
assert(replay.summary.outputs === 64, "replay: output count differs");
assert(
  replay.summary.torch_rng_unchanged_runs === 16,
  "replay: RNG nonconsumption count differs",
);
assert(
  reproduction.summary.cells === 64
    && reproduction.summary.all_preregistered_fields_exact === 64,
  "reproduction: exact cell count differs",
);
assert(
  Object.values(reproduction.summary.by_field).every((count) => count === 64),
  "reproduction: a field is not 64/64 exact",
);

assert(
  analysis.bootstrap_contract.resamples === 10000
    && analysis.bootstrap_contract.seed === 1364512825,
  "analysis: bootstrap contract differs",
);
let crnContrastChecks = 0;
for (const domain of ["code", "math"]) {
  const domainAnalysis = analysis.main_t0_selected_task_analysis[domain];
  assert(domainAnalysis.tasks === 32, `${domain}: task count differs`);
  for (const contrast of Object.values(domainAnalysis.contrasts)) {
    assert(
      contrast.trajectory.shared_uniform_prefix_exact === 32,
      `${domain}: contrast CRN audit differs`,
    );
    crnContrastChecks += contrast.trajectory.shared_uniform_prefix_exact;
  }
}
assert(crnContrastChecks === 256, "analysis: CRN contrast total differs");

process.stdout.write(
  `${JSON.stringify(
    {
      passed: true,
      protocol_id: PROTOCOL_ID,
      files: EXPECTED_HASHES,
      formal: {
        sources: sampling.sources.length,
        runs: sampling.summary.runs,
        outputs: sampling.summary.outputs,
        prompt_hashes_exact: sampling.source_contract.prompt_hash_audit.exact,
        uniform_output_hashes_exact: uniformOutputHashesExact,
        torch_rng_unchanged_runs: sampling.summary.torch_rng_unchanged_runs,
      },
      evaluation: {
        rows: evaluation.rows.length,
        main_t0_outputs: evaluation.summary.main_t0.outputs,
        code_success: evaluation.summary.main_t0.by_domain.code.fixed_budget_success,
        math_success: evaluation.summary.main_t0.by_domain.math.fixed_budget_success,
      },
      reproduction: reproduction.summary,
      analysis: {
        bootstrap_resamples: analysis.bootstrap_contract.resamples,
        crn_contrast_checks: crnContrastChecks,
      },
    },
    null,
    2,
  )}\n`,
);