#!/usr/bin/env node

import { createHash } from "node:crypto";
import { readFile, writeFile } from "node:fs/promises";
import { resolve } from "node:path";
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 DOMAINS = ["code", "math"];
const DIAGNOSTIC_INDICES = new Set([0, 8, 16, 24]);
const REPLAY_INDICES = new Set([0, 4, 8, 12, 16, 20, 24, 28]);

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

function uint32(label) {
  return createHash("sha256").update(label).digest().readUInt32BE(0);
}

function canonical(value) {
  if (Array.isArray(value)) {
    return `[${value.map(canonical).join(",")}]`;
  }
  if (value !== null && typeof value === "object") {
    return `{${Object.keys(value)
      .sort()
      .map((key) => `${JSON.stringify(key)}:${canonical(value[key])}`)
      .join(",")}}`;
  }
  return JSON.stringify(value);
}

function parseArgs(argv) {
  const args = {};
  for (let index = 2; index < argv.length; index += 2) {
    const key = argv[index];
    const value = argv[index + 1];
    if (!key?.startsWith("--") || value === undefined) {
      throw new Error(`Expected --name value arguments; got ${key ?? "<end>"}`);
    }
    args[key.slice(2)] = value;
  }
  if (!args.reference || !args.output) {
    throw new Error("Usage: build-deepseek-chat-task-bootstrap-manifest.mjs --reference FILE --output FILE");
  }
  return args;
}

const args = parseArgs(process.argv);
const referencePath = resolve(args.reference);
const outputPath = resolve(args.output);
const referenceBytes = await readFile(referencePath);
const reference = JSON.parse(referenceBytes);
const selected = reference.corpus_contract?.selected ?? [];

const sources = DOMAINS.flatMap((domain) => {
  const rows = selected
    .filter((row) => row.domain === domain)
    .sort((left, right) => left.within_domain_index - right.within_domain_index);
  if (rows.length !== 32) {
    throw new Error(`${domain}: expected 32 frozen sources, observed ${rows.length}`);
  }
  return rows.map((row, domainIndex) => {
    if (row.within_domain_index !== domainIndex) {
      throw new Error(`${domain}: non-contiguous within_domain_index at ${domainIndex}`);
    }
    const promptHashes = Object.fromEntries(
      CONDITIONS.map((condition) => {
        const hash = row.conditions?.[condition]?.token_ids_sha256;
        if (!/^[0-9a-f]{64}$/.test(hash ?? "")) {
          throw new Error(`${row.id}/${condition}: missing prompt hash`);
        }
        return [condition, hash];
      }),
    );
    return {
      id: row.id,
      domain,
      domain_index: domainIndex,
      selection_rank: row.selection_rank,
      source_text_sha256: row.text_sha256,
      routing_probe_prompt_token_ids_sha256: promptHashes,
      main_tapes: ["T0"],
      diagnostic_tapes: DIAGNOSTIC_INDICES.has(domainIndex)
        ? ["T1", "T2", "T3"]
        : [],
      independent_replay: REPLAY_INDICES.has(domainIndex),
    };
  });
});

const tapes = Array.from({ length: 4 }, (_, index) => {
  const label = `T${index}`;
  const derivation = `${PROTOCOL_ID}/tape/${index}`;
  return {
    label,
    index,
    display_seed: uint32(derivation),
    derivation,
    derivation_sha256: sha256(derivation),
  };
});

const manifest = {
  schema_version: 1,
  protocol_id: PROTOCOL_ID,
  status: "preregistered_before_any_protocol_output",
  source_reference: {
    path: args.reference,
    sha256: sha256(referenceBytes),
    sample_salt: reference.corpus_contract.sample_salt,
    selection_rule:
      "For each domain, all 32 rows ordered by frozen within_domain_index; no outcome-based selection.",
  },
  model: {
    repo: "deepseek-ai/DeepSeek-V2-Lite-Chat",
    revision: "85864749cd611b4353ce1decdb286193298f64c7",
  },
  conditions: CONDITIONS,
  domains: DOMAINS,
  sources,
  source_counts: Object.fromEntries(
    DOMAINS.map((domain) => [
      domain,
      sources.filter((source) => source.domain === domain).length,
    ]),
  ),
  tape_contract: {
    tapes,
    main_tape: "T0",
    diagnostic_source_indices_per_domain: [...DIAGNOSTIC_INDICES],
    diagnostic_tapes: ["T0", "T1", "T2", "T3"],
    uniform_uint64:
      "first 8 SHA-256 bytes, big-endian, of protocol_id + '\\0uniform\\0' + tape_label + '\\0' + source_id + '\\0' + decimal(step)",
    uniform_open_interval:
      "u_t = (uniform_uint64 + 0.5) / 2^64; cast to torch.float32 before searchsorted",
    common_random_number_scope:
      "For the same source and tape, all four conditions consume the identical u_t at generation step t while each trajectory remains active.",
  },
  execution_grid: {
    main: {
      sources: 64,
      tapes_per_source: 1,
      conditions: 4,
      outputs: 256,
    },
    additional_tape_diagnostic: {
      sources: 8,
      additional_tapes_per_source: 3,
      conditions: 4,
      outputs: 96,
    },
    formal_outputs: 352,
    independent_replay: {
      source_indices_per_domain: [...REPLAY_INDICES],
      tapes: ["T0"],
      conditions: 4,
      outputs: 64,
    },
  },
  bootstrap_contract: {
    resamples: 10000,
    domain_separated: true,
    main_task_bootstrap_seed: uint32(
      `${PROTOCOL_ID}/bootstrap/task-bootstrap`,
    ),
    diagnostic_task_bootstrap_seed: uint32(
      `${PROTOCOL_ID}/bootstrap/diagnostic-task-bootstrap`,
    ),
    tape_bootstrap_seed: uint32(
      `${PROTOCOL_ID}/bootstrap/tape-bootstrap`,
    ),
    estimand_boundary:
      "Selected-task resampling bands for this fixed 32-task frame and one frozen tape; not benchmark-population or generation-seed confidence intervals.",
  },
  prompt_contract: {
    routing_probe_hashes_are_not_chat_generation_hashes: true,
    chat_generation_hashes:
      "Added by freeze-deepseek-chat-task-bootstrap-prompts.py before any protocol output.",
  },
};

manifest.canonical_content_sha256 = sha256(canonical(manifest));
await writeFile(outputPath, `${JSON.stringify(manifest, null, 2)}\n`);

const outputBytes = await readFile(outputPath);
process.stdout.write(
  `${JSON.stringify(
    {
      output: outputPath,
      bytes: outputBytes.length,
      sha256: sha256(outputBytes),
      canonical_content_sha256: manifest.canonical_content_sha256,
      sources: sources.length,
      formal_outputs: manifest.execution_grid.formal_outputs,
      replay_outputs: manifest.execution_grid.independent_replay.outputs,
    },
    null,
    2,
  )}\n`,
);