#!/usr/bin/env python3
"""Recount, hash, and consolidate generated corpora without trusting filenames."""

from __future__ import annotations

import hashlib
import json
import re
from pathlib import Path

from transformers import AutoTokenizer

from generate_corpora import CONFIG, combined_question, exact_serialization, token_counts


ROOT = Path(__file__).resolve().parent
SOURCES = (ROOT / "generated", ROOT / "generated-near1m")


def main() -> int:
    tokenizers = {
        name: AutoTokenizer.from_pretrained(
            spec["repository"], revision=spec["revision"], trust_remote_code=True
        )
        for name, spec in CONFIG["tokenizers"].items()
    }
    tier_lookup = {row["name"]: row for row in CONFIG["tiers"]}
    cases: dict[str, Path] = {}
    for source in SOURCES:
        if not source.exists():
            continue
        for candidate in source.iterdir():
            if candidate.is_dir() and re.fullmatch(r"(?:32k|128k|256k|512k|near-1m)-trial-[123]", candidate.name):
                cases[candidate.name] = candidate
    expected_ids = {
        f"{tier['name']}-trial-{trial}"
        for tier in CONFIG["tiers"]
        for trial in range(1, CONFIG["trials_per_tier"] + 1)
    }
    if set(cases) != expected_ids:
        raise RuntimeError(f"case mismatch: missing={sorted(expected_ids-set(cases))}, extra={sorted(set(cases)-expected_ids)}")

    rows = []
    for case_id in sorted(cases, key=lambda value: (list(tier_lookup).index(value.rsplit("-trial-", 1)[0]), value)):
        case_dir = cases[case_id]
        tier_name, trial_text = case_id.rsplit("-trial-", 1)
        trial = int(trial_text)
        corpus = (case_dir / "corpus.txt").read_text(encoding="utf-8")
        expected = json.loads((case_dir / "expected.json").read_text(encoding="utf-8"))
        question = (case_dir / "question.txt").read_text(encoding="utf-8").strip()
        if question != combined_question(expected):
            raise RuntimeError(f"question drift in {case_id}")
        evidence_ids = set(re.findall(r"EVID-[A-Z]+-[0-9]+-[A-Z0-9-]+", corpus))
        expected_evidence = set(re.findall(r"EVID-[A-Z]+-[0-9]+-[A-Z0-9-]+", json.dumps(expected)))
        if not expected_evidence.issubset(evidence_ids):
            raise RuntimeError(f"missing evidence in {case_id}: {sorted(expected_evidence-evidence_ids)}")
        corpus_counts = token_counts(corpus, tokenizers)
        serialized_counts = token_counts(exact_serialization(corpus, question), tokenizers)
        tier = tier_lookup[tier_name]
        if max(corpus_counts.values()) > tier["target_corpus_tokens"]:
            raise RuntimeError(f"{case_id} exceeds corpus target: {corpus_counts}")
        rows.append({
            "case_id": case_id,
            "source_directory": str(case_dir.relative_to(ROOT)),
            "tier": tier_name,
            "trial": trial,
            "utf8_bytes": len(corpus.encode("utf-8")),
            "sha256": hashlib.sha256(corpus.encode("utf-8")).hexdigest(),
            "document_record_count": len(re.findall(r"^DOC-[0-9]{7} \|", corpus, re.MULTILINE)),
            "evidence_id_count": len(evidence_ids),
            "corpus_token_counts": corpus_counts,
            "transparent_serialized_request_token_counts": serialized_counts,
            "nominal_context_tokens": tier["nominal_context_tokens"],
            "target_corpus_tokens": tier["target_corpus_tokens"],
            "max_completion_tokens": CONFIG["max_completion_tokens"],
            "api_accepted_prompt_tokens": {"kimi": None, "deepseek": None},
            "status": "generated_not_submitted"
        })
    manifest = {
        "schema_version": "1.0",
        "generated_from": "kingy-context-reality-v1",
        "verified_case_count": len(rows),
        "token_count_semantics": (
            "Exact offline counts from pinned released tokenizers. Not API-accepted counts; successful endpoint usage is authoritative."
        ),
        "tokenizers": CONFIG["tokenizers"],
        "corpora": rows,
    }
    output = ROOT / "generated" / "manifest.json"
    output.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    print(json.dumps({
        "manifest": str(output),
        "cases": len(rows),
        "counts_by_tier": {
            tier: [
                {"trial": row["trial"], **row["corpus_token_counts"]}
                for row in rows if row["tier"] == tier
            ]
            for tier in tier_lookup
        }
    }, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
