"""Evaluate a completed checkpoint on another source without tuning on it."""
import os

os.environ.setdefault("USE_TF", "0")
os.environ.setdefault("HF_HUB_OFFLINE", "1")
os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")

import argparse
import json
import random
from collections import Counter
from pathlib import Path

import torch

from .fetch import ROOT, sha256
from .finetune import evaluate, probs, tokenize_records
from .metrics import binary_metrics


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", type=Path, default=ROOT / "artifacts/laya-vulnerability")
    parser.add_argument("--base", type=Path, default=ROOT / "models/laya")
    parser.add_argument("--data", type=Path, default=ROOT / "data/vulnerability/processed")
    parser.add_argument("--source", default="devign")
    parser.add_argument("--limit", type=int, default=1024)
    parser.add_argument("--device", default="cpu")
    parser.add_argument("--batch-size", type=int, default=4)
    args = parser.parse_args()
    if not (args.model / "run/metrics.json").exists():
        parser.error("Fine-tuning must finish before external evaluation")
    config = json.loads((args.model / "run/config.json").read_text())
    if sha256(args.base / "model.safetensors") != config["base_sha256"]:
        raise ValueError("External baseline must match this experiment's original checkpoint")
    if (args.data / ".preparing").exists():
        raise ValueError("External data preparation has not finished")
    summary = json.loads((args.data / "summary.json").read_text())
    for split in ("train", "valid", "test"):
        if sha256(args.data / f"{split}.jsonl") != summary["files"][split]["sha256"]:
            raise ValueError("External data does not match its recorded hashes")
    torch.set_num_threads(4)
    seen_ids, seen_shapes = set(), set()
    for path in (args.model / "run").glob("*_selection.jsonl"):
        for line in path.open():
            r = json.loads(line)
            seen_ids.add(r["id"])
            seen_shapes.add(r["structural_hash"])
    rng = random.Random(20261010)
    reservoirs, totals = {0: [], 1: []}, Counter()
    for split in ("train", "valid", "test"):
        for line in (args.data / f"{split}.jsonl").open():
            r = json.loads(line)
            if r["source"] != args.source or r["id"] in seen_ids or r["structural_hash"] in seen_shapes:
                continue
            label = r["label"]
            totals[label] += 1
            if len(reservoirs[label]) < args.limit // 2:
                reservoirs[label].append(r)
            else:
                index = rng.randrange(totals[label])
                if index < args.limit // 2:
                    reservoirs[label][index] = r
    records = reservoirs[0] + reservoirs[1]
    if not reservoirs[0] or not reservoirs[1]:
        raise ValueError("Need both external classes, after excluding known examples")
    import laya
    agent = laya.load(str(args.model.resolve()), device=None if args.device == "auto" else args.device)
    max_len, head_max_len = agent.cfg["max_len"], agent.cfg["head_max_len"]
    items, excluded = tokenize_records(records, agent, max_len, head_max_len)
    if len({r["label"] for r in items}) != 2:
        raise ValueError("External evaluation needs both classes after context filtering")
    temperature = agent.cfg["temperature_by_options"]["choice:2"]
    threshold = agent.cfg["vulnerability_threshold"]
    logits, labels, ids = evaluate(agent.model, items, agent, args.batch_size)
    final = {"at_0.5": binary_metrics(labels, probs(logits, temperature)),
             "at_validation_threshold": binary_metrics(labels, probs(logits, temperature), threshold)}
    del agent
    if args.device == "mps":
        torch.mps.empty_cache()
    baseline = laya.load(str(args.base.resolve()), device=None if args.device == "auto" else args.device)
    bl, by, _ = evaluate(baseline.model, items, baseline, args.batch_size)
    selection = args.model / "run" / f"external-{args.source}-selection.jsonl"
    with selection.open("w") as f:
        for record in records:
            f.write(json.dumps(record) + "\n")
    result = {"source": args.source, "selected": len(records), "used": len(items), "excluded": excluded,
              "seed": 20261010, "requested_limit": args.limit, "selection_sha256": sha256(selection),
              "model_sha256": sha256(args.model / "model.safetensors"), "base_sha256": config["base_sha256"],
              "label_counts": dict(Counter(r["label"] for r in items)), "temperature": temperature, "threshold": threshold,
              "tuning_on_external_data": False, "known_id_and_shape_exclusion": True,
              "baseline_at_0.5": binary_metrics(by, probs(bl)), "fine_tuned": final,
              "dataset_hashes": {s: sha256(args.data / f"{s}.jsonl") for s in ("train", "valid", "test")},
              "limitations": "Balanced, context-filtered function samples; upstream negatives are not proof of complete safety. Not a project-level audit."}
    target = args.model / "run" / f"external-{args.source}.json"
    target.write_text(json.dumps(result, indent=2))
    with target.with_suffix(".jsonl").open("w") as f:
        for index, identity in enumerate(ids):
            f.write(json.dumps({"id": identity, "label": int(labels[index]), "probability": float(probs(logits[index:index+1], temperature)[0])}) + "\n")
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()