#!/usr/bin/env python3
"""Create the immutable, blinded v2 challenge release artifacts.

This is a custodial operation only: it copies selected source formulas into an
anonymous evaluator payload and writes the labels to a separate restricted key.
It never loads model embeddings or computes evaluation outcomes.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import random
from pathlib import Path
from typing import Any


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def canonical_formula(row: dict[str, Any]) -> list[dict[str, Any]]:
    """Retain formula inputs while stripping labels and source metadata."""
    return [
        {
            key: component[key]
            for key in ("cas", "name", "smiles", "weight_fraction")
            if key in component
        }
        for component in row.get("formula", [])
    ]


def archive(challenge_path: Path, dataset_path: Path, output_dir: Path, seed: int) -> dict[str, Any]:
    challenge = json.loads(challenge_path.read_text(encoding="utf-8"))
    records = challenge["challenge_records"]
    if len(records) % 4:
        raise ValueError("challenge records must contain complete four-record blocks")

    wanted = {int(record["source_index"]) for record in records}
    sources: dict[int, dict[str, Any]] = {}
    with dataset_path.open(encoding="utf-8") as handle:
        for source_index, line in enumerate(handle):
            if source_index in wanted:
                sources[source_index] = json.loads(line)
    if sources.keys() != wanted:
        raise ValueError(f"missing source rows: {sorted(wanted - sources.keys())}")

    rng = random.Random(seed)
    block_order = list(range(len(records) // 4))
    rng.shuffle(block_order)
    blinded, key_records = [], []
    for public_block_number, source_block_number in enumerate(block_order, start=1):
        source_block = records[source_block_number * 4:(source_block_number + 1) * 4]
        within_order = list(range(4))
        rng.shuffle(within_order)
        for public_position, source_position in enumerate(within_order, start=1):
            selected = source_block[source_position]
            blind_id = f"V2-B{public_block_number:03d}-R{public_position}"
            source_index = int(selected["source_index"])
            blinded.append({
                "blind_id": blind_id,
                "block_id": f"V2-B{public_block_number:03d}",
                "formula": canonical_formula(sources[source_index]),
            })
            key_records.append({
                "blind_id": blind_id,
                "source_block": source_block_number + 1,
                "source_index": source_index,
                "formula_id": selected.get("formula_id"),
                "genre": selected["genre"],
            })

    output_dir.mkdir(parents=True, exist_ok=True)
    blinded_path = output_dir / "blinded_challenge.json"
    key_path = output_dir / "label_key.json"
    blinded_doc = {
        "release": "genre-challenge-v2",
        "status": "frozen_blinded_evaluator_payload",
        "records": blinded,
    }
    key_doc = {
        "release": "genre-challenge-v2",
        "status": "restricted; do not disclose until predictions are frozen",
        "blinding_seed": seed,
        "records": key_records,
    }
    blinded_path.write_text(json.dumps(blinded_doc, indent=2) + "\n", encoding="utf-8")
    key_path.write_text(json.dumps(key_doc, indent=2) + "\n", encoding="utf-8")
    os.chmod(blinded_path, 0o444)
    os.chmod(key_path, 0o400)
    return {
        "challenge_sha256": sha256(challenge_path),
        "source_dataset_sha256": sha256(dataset_path),
        "blinded_payload_sha256": sha256(blinded_path),
        "label_key_sha256": sha256(key_path),
        "records": len(blinded),
        "blocks": len(block_order),
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--challenge", required=True, type=Path)
    parser.add_argument("--dataset", required=True, type=Path)
    parser.add_argument("--output-dir", required=True, type=Path)
    parser.add_argument("--seed", type=int, default=20260715)
    args = parser.parse_args()
    print(json.dumps(archive(args.challenge, args.dataset, args.output_dir, args.seed), indent=2))


if __name__ == "__main__":
    main()