"""Measure backbone calls for real text and photograph requests on the frozen model."""

import argparse
import hashlib
import json
import math
from pathlib import Path

import torch
from train_foundation_head import write

from veyra.data import read_records
from veyra.option_model import OptionModel
from veyra.schema import DecisionRequest


def digest(path):
    return hashlib.sha256(Path(path).read_bytes()).hexdigest()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", type=Path, required=True)
    parser.add_argument("--records", type=Path, default=Path("data/foundation-v11/records.jsonl"))
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    if args.output.exists():
        raise FileExistsError("contract evidence is immutable")
    records = sorted(
        (r for r in read_records(args.records) if r.split == "test"), key=lambda r: r.id
    )
    photo = next(r for r in records if "domain:image_waste" in r.tags)
    text = next(r for r in records if not r.request.state.images)
    questions = {}
    for row in records:
        if row.group_id == photo.group_id:
            for question in row.request.questions.values():
                questions.setdefault(question.type, question)
    if set(questions) != {"choice", "score", "noul"}:
        raise ValueError("photo fixture must cover all three decision types")
    requests = {
        "text": text.request,
        "photograph_three_questions": DecisionRequest(
            state=photo.request.state, questions=questions
        ),
    }
    torch.set_num_threads(4)
    model = OptionModel.load(args.checkpoint, local_files_only=True, merge=True)
    forwards, results = [], {}
    hook = model.encoder.model.register_forward_hook(lambda *_: forwards.append(1))
    try:
        for name, request in requests.items():
            start = len(forwards)
            response = model.predict(request, args.records.parent.resolve())
            count = len(forwards) - start
            valid = all(
                all(math.isfinite(p) and 0 <= p <= 1 for p in answer["probabilities"].values())
                and abs(sum(answer["probabilities"].values()) - 1) <= 1e-5
                for answer in response["answers"].values()
            )
            results[name] = {
                "backbone_forwards": count,
                "generated_answer_tokens": response["usage"]["output_tokens"],
                "valid_probabilities": valid,
                "questions": len(request.questions),
                "types": sorted({q.type for q in request.questions.values()}),
                "input_tokens": response["usage"]["input_tokens"],
            }
    finally:
        hook.remove()
    result = {
        "weights_sha256": digest(args.checkpoint / "head.safetensors"),
        "manifest_sha256": digest(args.checkpoint / "manifest.json"),
        "records_sha256": digest(args.records),
        "source_sha256": digest(Path(__file__)),
        "gpu": torch.cuda.get_device_name(),
        "record_ids": [text.id, photo.id],
        "results": results,
        "passed": all(
            r["backbone_forwards"] == 1
            and r["generated_answer_tokens"] == 0
            and r["valid_probabilities"]
            for r in results.values()
        ),
    }
    write(args.output, result)
    print(json.dumps(result), flush=True)
    if not result["passed"]:
        raise SystemExit(1)


if __name__ == "__main__":
    main()