import { createHash } from "node:crypto";
import { readFileSync, statSync, writeFileSync } from "node:fs";
import { resolve } from "node:path";

const root = resolve(import.meta.dirname, "..");
const formalPath = resolve(
  root,
  "src/data/deepseek-v2-lite-chat-behavior.json",
);
const reproPath = resolve(
  root,
  "src/data/deepseek-v2-lite-chat-behavior-repro-1pd.json",
);
const outputPath = resolve(
  root,
  "src/data/deepseek-v2-lite-chat-behavior-compact.json",
);

const sha256 = (path) => createHash("sha256")
  .update(readFileSync(path))
  .digest("hex");

const formal = JSON.parse(readFileSync(formalPath, "utf8"));
const repro = JSON.parse(readFileSync(reproPath, "utf8"));
const formalSha256 = sha256(formalPath);
const reproSha256 = sha256(reproPath);

const formalOutputByKey = new Map(
  formal.sources.flatMap((source) => source.outputs.map((output) => [
    `${source.id}\0${output.condition}`,
    { source, output },
  ])),
);
const reproduction = {
  sources: repro.sources.length,
  cells: 0,
  promptHashExact: 0,
  generatedTokenIdsExact: 0,
  generatedTextExact: 0,
  eosStateExact: 0,
};
for (const source of repro.sources) {
  for (const output of source.outputs) {
    const key = `${source.id}\0${output.condition}`;
    const reference = formalOutputByKey.get(key);
    if (!reference) throw new Error(`formal output missing: ${key}`);
    reproduction.cells += 1;
    reproduction.promptHashExact += (
      reference.output.prompt_token_ids_sha256
      === output.prompt_token_ids_sha256
    );
    reproduction.generatedTokenIdsExact += (
      JSON.stringify(reference.output.generated_token_ids)
      === JSON.stringify(output.generated_token_ids)
    );
    reproduction.generatedTextExact += (
      reference.output.text === output.text
    );
    reproduction.eosStateExact += (
      reference.output.hit_eos === output.hit_eos
    );
  }
}
if (
  reproduction.cells !== 32
  || reproduction.promptHashExact !== reproduction.cells
  || reproduction.generatedTokenIdsExact !== reproduction.cells
  || reproduction.generatedTextExact !== reproduction.cells
  || reproduction.eosStateExact !== reproduction.cells
) {
  throw new Error(
    `Chat behavior reproduction mismatch: ${JSON.stringify(reproduction)}`,
  );
}

const formalModelHashes = Object.fromEntries(
  Object.entries(formal.model.files).map(([name, value]) => [
    name,
    value.sha256,
  ]),
);
const reproModelHashes = Object.fromEntries(
  Object.entries(repro.model.files).map(([name, value]) => [
    name,
    value.sha256,
  ]),
);
if (JSON.stringify(formalModelHashes) !== JSON.stringify(reproModelHashes)) {
  throw new Error("formal and reproduction model-file hashes differ");
}

const conditions = formal.generation_contract.conditions;
const edgeOrder = [
  "system_eos",
  "system_bos",
  "system_x",
  "system_period",
  "bos_at_s0",
  "bos_at_s1",
  "x_at_s0",
  "x_at_s1",
  "period_at_s0",
  "period_at_s1",
];
const edgeLabels = {
  system_eos: "System on − off · EOS",
  system_bos: "System on − off · BOS",
  system_x: "System on − off · x",
  system_period: "System on − off · 句点",
  bos_at_s0: "BOS − EOS · system off",
  bos_at_s1: "BOS − EOS · system on",
  x_at_s0: "x − EOS · system off",
  x_at_s1: "x − EOS · system on",
  period_at_s0: "句点 − EOS · system off",
  period_at_s1: "句点 − EOS · system on",
};

const compactOutput = (output) => ({
  condition: output.condition,
  factors: output.factors,
  promptTokens: output.prompt_tokens,
  leftPaddingTokens: output.left_padding_tokens,
  promptTokenIdsSha256: output.prompt_token_ids_sha256,
  generatedTokens: output.generated_tokens,
  generatedTokenIdsSha256: output.generated_token_ids_sha256,
  hitEos: output.hit_eos,
  stoppedAtMaxNewTokens: output.stopped_at_max_new_tokens,
  text: output.text,
  textSha256: output.text_sha256,
  taskScore: output.task_score,
});

const compact = {
  schemaVersion: 1,
  source: {
    formalSha256,
    reproSha256,
    formalBytes: statSync(formalPath).size,
    reproBytes: statSync(reproPath).size,
    reproduction,
  },
  model: {
    repo: formal.model.repo,
    revision: formal.model.revision,
    checkpointIdentity: formal.model.checkpoint_identity,
    architecture: formal.model.architecture,
    dtype: formal.model.dtype,
    checkpointTensorBytes: formal.model.checkpoint_tensor_bytes,
    shardFileBytesIncludingHeaders: (
      formal.model.shard_file_bytes_including_headers
    ),
    allFilesSameRevision: (
      formal.model.download_revision_contract.all_files_same_revision
    ),
    fileCount: Object.keys(formal.model.files).length,
    fileHashes: formalModelHashes,
  },
  tokenizer: formal.tokenizer_contract,
  sources: formal.sources.map((source) => ({
    id: source.id,
    domain: source.domain,
    label: source.label,
    withinDomainIndex: source.within_domain_index,
    selectionRank: source.selection_rank,
    sourceTextSha256: source.source_text_sha256,
    sourceCharacters: source.source_characters,
    sourceTokens: source.source_tokens,
    promptTokensByCondition: source.prompt_tokens_by_condition,
    batchPromptTokens: source.batch_prompt_tokens_after_left_padding,
    batchPaddingSide: source.batch_padding_side,
    outputs: source.outputs.map(compactOutput),
  })),
  contract: {
    source: formal.source_contract,
    conditions,
    edgeOrder,
    edgeLabels,
    officialSerialization: (
      formal.generation_contract.official_serialization
    ),
    decode: formal.generation_contract.decode,
    doSample: formal.generation_contract.do_sample,
    maxNewTokens: formal.generation_contract.max_new_tokens,
    useCache: formal.generation_contract.use_cache,
    batchPadding: formal.generation_contract.batch_padding,
    counterfactualBoundary: (
      formal.generation_contract.counterfactual_boundary
    ),
    officialGenerationConfigRecordedNotUsed: (
      formal.generation_contract
        .official_generation_config_recorded_not_used
    ),
  },
  execution: {
    python: formal.execution.python,
    torch: formal.execution.torch,
    transformers: formal.execution.transformers,
    accelerate: formal.execution.accelerate,
    safetensors: formal.execution.safetensors,
    gpu: formal.execution.gpu,
    gpuMemoryLimit: formal.execution.gpu_memory_limit,
    cpuMemoryLimit: formal.execution.cpu_memory_limit,
    inputDevice: formal.execution.input_device,
    deviceMap: formal.execution.device_map,
    parameterBytesByRuntimeParameterDevice: (
      formal.execution.parameter_bytes_by_runtime_parameter_device
    ),
    parameterBytesByDtype: formal.execution.parameter_bytes_by_dtype,
    offloadParameterDeviceNote: (
      formal.execution.offload_parameter_device_note
    ),
    loadSeconds: formal.execution.load_seconds,
    generationSeconds: formal.sources.reduce(
      (sum, source) => sum + source.generation_seconds,
      0,
    ),
    loadPeakCudaMemoryAllocatedBytes: (
      formal.execution.load_peak_cuda_memory_allocated_bytes
    ),
    peakCudaMemoryAllocatedBytes: (
      formal.execution.peak_cuda_memory_allocated_bytes
    ),
    processMaxRssKib: formal.execution.process_max_rss_kib,
    officialSingleGpuBf16Requirement: (
      formal.execution.official_single_gpu_bf16_requirement
    ),
    localSingleGpuCapacityMib: (
      formal.execution.local_single_gpu_capacity_mib
    ),
    offloadRequiredByLocalCapacity: (
      formal.execution.offload_required_by_local_capacity
    ),
  },
  summary: {
    aggregates: Object.fromEntries(
      edgeOrder.map((edge) => [
        edge,
        formal.summary.aggregates[edge],
      ]),
    ),
    aggregatesByDomain: formal.summary.aggregates_by_domain,
    outputByCondition: formal.summary.output_by_condition,
    taskByCondition: formal.summary.task_by_condition,
    pairwise: Object.fromEntries(
      edgeOrder.map((edge) => [
        edge,
        formal.summary.pairwise[edge],
      ]),
    ),
    completedOutputs: formal.sources.reduce(
      (sum, source) => sum + source.outputs.filter(
        (output) => output.hit_eos,
      ).length,
      0,
    ),
    truncatedOutputs: formal.sources.reduce(
      (sum, source) => sum + source.outputs.filter(
        (output) => output.stopped_at_max_new_tokens,
      ).length,
      0,
    ),
  },
  claimBoundary: formal.claim_boundary,
};

writeFileSync(
  outputPath,
  `${JSON.stringify(compact, null, 2)}\n`,
  "utf8",
);
process.stdout.write(
  `${outputPath}\n`
  + `${formalSha256}\n`
  + `${statSync(formalPath).size} bytes formal → `
  + `${statSync(outputPath).size} bytes compact\n`
  + `${reproduction.generatedTokenIdsExact}`
  + ` / ${reproduction.cells} long-sequence cells exact\n`,
);