"""Audit corpus coverage and freeze reproducible, position-balanced evaluations."""
import hashlib
import json
import random
import re
from collections import Counter, defaultdict
from datetime import datetime, timezone
from pathlib import Path

ROOT = Path(__file__).resolve().parent
OUT = ROOT / "benchmarks"
SEED = 20261109
SECEVAL_REVISION = "205dab7b0888a06f4b53ca7d9c7093e1326683e1"
CVE = re.compile(r"CVE-\d{4}-\d{4,}", re.I)
TOPICS = {
    "iOS / Apple platform": r"\b(?:ios|ipados|xnu|iokit|webkit|darwin|iphone)\b",
    "Kernel / operating systems": r"\b(?:kernel|syscall|sandbox|system call|privilege escalation|linux|windows|android)\b",
    "Memory safety / binary analysis": r"\b(?:buffer overflow|out.of.bounds|use.after.free|heap|stack|rop|disassembl\w*|fuzz\w*)\b",
    "Web / API security": r"\b(?:sql injection|xss|ssrf|csrf|http|web|api|cookies?)\b",
    "Network security": r"\b(?:network|tcp|udp|dns|firewall|packet|wi.fi|wireless|vpn)\b",
    "Identity / authentication": r"\b(?:authentication|oauth|saml|kerberos|mfa|active directory|credentials?)\b",
    "Cryptography": r"\b(?:cryptograph\w*|encryption|cipher|tls|ssl|hash|certificate)\b",
    "Cloud / containers": r"\b(?:aws|azure|cloud|docker|kubernetes|container)\b",
    "Malware / incident analysis": r"\b(?:malware|ransomware|forensic\w*|incident|yara|trojan)\b",
}


def read_rows(split):
    return [json.loads(line) for line in (ROOT / "data" / f"{split}.jsonl").read_text(encoding="utf-8").splitlines()]


def normalized(text):
    return " ".join(re.findall(r"\w+", text.casefold()))


def shingles(text):
    words = normalized(text).split()
    return {tuple(words[i:i+4]) for i in range(max(0, len(words)-3))}


def parse_local(row):
    stem, options = row["prompt"].split("\n\nA) ", 1)
    options = "A) " + options.rsplit("\n\nChoose the best answer and explain it.", 1)[0]
    matches = re.findall(r"(?ms)^([A-D])\) (.*?)(?=^[A-D]\) |\Z)", options)
    assert [m[0] for m in matches] == list("ABCD"), row["group"]
    match = re.match(r"^([A-D])\)", row["response"])
    assert match, row["group"]
    return stem, [m[1].strip() for m in matches], "ABCD".index(match.group(1))


def scoring_prefix(tokenizer, question, choices):
    prompt = question + "\n\n" + "\n".join(f"{chr(65+i)}) {choice}" for i, choice in enumerate(choices))
    prompt += "\n\nSelect the best answer. Reply with only its letter."
    prefix = tokenizer.apply_chat_template([{"role":"user", "content":prompt}], tokenize=False, add_generation_prompt=True)
    assert prefix.endswith("<think>\n")
    return prefix + "\n</think>\n\n"


def coverage(train, evaluation):
    result = {"created_utc": datetime.now(timezone.utc).isoformat(), "definitions": {
        "ios_versions": "Release containing the fix, as named by Apple; not an affected-version range.",
        "topics": "Non-exclusive keyword matches in prompt plus answer, not expert topic labels or measured model competence.",
        "unique_cves": "Distinct CVE identifiers mentioned in retained records; not a fraction of all publicly known CVEs."}}
    for name, rows in [("train", train), ("eval", evaluation)]:
        versions, components, vendors, products = Counter(), Counter(), Counter(), Counter()
        for row in rows:
            if row["kind"] == "apple_advisory":
                versions[re.search(r"Advisory: iOS ([\d.]+)", row["prompt"]).group(1)] += 1
                excerpt = row["prompt"].split("\n\n")[1].splitlines()
                component = excerpt[0] if len(excerpt) > 1 and excerpt[1].startswith("Available for:") else "Unparsed component"
                components[component] += 1
            elif row["kind"] == "cisa_kev":
                record = json.loads(row["prompt"].split("\n\n")[1])
                vendors[record["vendorProject"]] += 1
                products[record["vendorProject"] + " / " + record["product"]] += 1
        labels = Counter(row["response"][0] for row in rows if row["kind"] == "pentesting_explanation")
        unique_cves = {cve for row in rows for cve in row["cves"]}
        result[name] = dict(records=len(rows), by_source=dict(Counter(r["kind"] for r in rows)),
            unique_cves=len(unique_cves), cve_years=dict(Counter(c.split("-")[1] for c in unique_cves)),
            ios_versions=dict(versions), ios27_examples=sum(n for v,n in versions.items() if v.split(".")[0] == "27"),
            apple_components=dict(components.most_common()), cisa_vendors=dict(vendors.most_common()),
            cisa_distinct_vendors=len(vendors), cisa_distinct_products=len(products),
            mcq_original_answer_labels=dict(labels),
            keyword_topics={topic:sum(bool(re.search(pattern,r["prompt"]+"\n"+r["response"],re.I)) for r in rows)
                            for topic,pattern in TOPICS.items()})
    result["split_overlap"] = {
        "groups": len({r["group"] for r in train} & {r["group"] for r in evaluation}),
        "cves": len({c for r in train for c in r["cves"]} & {c for r in evaluation for c in r["cves"]})}
    result["corpus_sha256"] = {s:hashlib.sha256((ROOT/"data"/f"{s}.jsonl").read_bytes()).hexdigest() for s in ["train","eval"]}
    result["gaps"] = [
        "No iOS 27 advisory examples are in the current held-out split; iOS 27 generalization remains unmeasured.",
        "Apple examples primarily teach advisory interpretation, not kernel/WebKit source-code auditing.",
        "Original MCQ answer positions are heavily skewed toward A; the separate benchmark reorders eligible choices to balance answer letters.",
        "Exact group/CVE and lexical overlap checks cannot rule out semantic overlap or base-model pretraining contamination."]
    return result


def prepare():
    from transformers import AutoTokenizer
    OUT.mkdir(exist_ok=True)
    train, evaluation = read_rows("train"), read_rows("eval")
    report = coverage(train,evaluation)
    (OUT/"coverage.json").write_text(json.dumps(report,indent=2),encoding="utf-8")
    tokenizer = AutoTokenizer.from_pretrained(ROOT/"base", local_files_only=True)
    train_questions = [parse_local(r)[0] if r["kind"] == "pentesting_explanation" else r["prompt"] for r in train]
    exact = {normalized(q) for q in train_questions}
    train_grams = [shingles(q) for q in train_questions]
    index = defaultdict(set)
    for i, grams in enumerate(train_grams):
        for gram in grams:
            index[gram].add(i)
    train_cves = {c for r in train for c in r["cves"]}
    excluded = Counter()
    pools = defaultdict(list)

    def add(suite, case_id, question, choices, answer, topics, source):
        if len(choices) != 4 or len(set(choices)) != 4 or not 0 <= answer < 4:
            excluded[suite+":invalid_choices"] += 1
            return
        # Reordering must not change the meaning of positional references.
        position = r"\b(?:above|below|preceding|following options|first two|last two)\b|\b(?:option|choice|statement)s?\s+[A-D]\b|\b[A-D]\s+(?:and|or)\s+[A-D]\b"
        if re.search(position, question + "\n" + "\n".join(choices), re.I):
            excluded[suite+":position_dependent_options"] += 1
            return
        if normalized(question) in exact or train_cves.intersection(c.upper() for c in CVE.findall(question+" "+" ".join(choices))):
            excluded[suite+":training_overlap"] += 1
            return
        grams = shingles(question)
        candidates = set().union(*(index.get(g,set()) for g in grams)) if grams else set()
        if any(len(grams & train_grams[i])/max(1,len(grams | train_grams[i])) >= 0.8 for i in candidates):
            excluded[suite+":lexical_near_duplicate"] += 1
            return
        if len(tokenizer(scoring_prefix(tokenizer,question,choices),add_special_tokens=False)["input_ids"]) > 768:
            excluded[suite+":overlength"] += 1
            return
        pools[suite].append(dict(id=suite+":"+str(case_id), question=question, choices=choices,
            answer_index=answer,topics=topics,source=source))

    for row in evaluation:
        if row["kind"] == "pentesting_explanation":
            stem, choices, answer = parse_local(row)
            add("heldout_pentesting",row["group"],stem,choices,answer,["Pentesting"],row["source"])
    seceval_file = ROOT/"data"/"benchmark-sources"/"SecEval"/"questions.json"
    for row in json.loads(seceval_file.read_text(encoding="utf-8")):
        if row["answer"].strip() not in list("ABCD"):
            excluded["seceval_200:multiple_or_invalid_answer"] += 1
            continue
        choices = [re.sub(r"^[A-D]\s*[:.)]\s*", "", choice).strip() for choice in row["choices"]]
        add("seceval_200",row["id"],row["question"],choices,"ABCD".index(row["answer"].strip()),
            row.get("topics") or ["Unspecified"],"https://huggingface.co/datasets/XuanwuAI/SecEval/tree/"+SECEVAL_REVISION)
    selected = []
    eligible = {name:len(pool) for name,pool in pools.items()}
    for name, pool in sorted(pools.items()):
        rng = random.Random(SEED)
        rng.shuffle(pool)
        if name == "seceval_200":
            # Round-robin by primary topic to sample across available domains.
            buckets = defaultdict(list)
            for row in pool:
                buckets[row["topics"][0]].append(row)
            ordered = []
            while any(buckets.values()):
                for topic in sorted(buckets):
                    if buckets[topic]:
                        ordered.append(buckets[topic].pop())
            pool = ordered
        limit = min(200 if name == "seceval_200" else len(pool), len(pool))//4*4
        if name != "seceval_200" and len(pool) > limit:
            excluded[name+":balance_rounding"] += len(pool)-limit
        targets = list(range(4))*(limit//4)
        rng.shuffle(targets)
        for row, target in zip(pool[:limit],targets):
            original_answer = row["answer_index"]
            original_choices = row["choices"]
            distractors = [i for i in range(4) if i != original_answer]
            rng.shuffle(distractors)
            distractors.insert(target,original_answer)
            row.update(choices=[original_choices[i] for i in distractors],answer_index=target,
                       answer_letter=chr(65+target),original_choice_indices=distractors,
                       question_sha256=hashlib.sha256(row["question"].encode("utf-8")).hexdigest())
            prefix = scoring_prefix(tokenizer,row["question"],row["choices"])
            ids = tokenizer(prefix,add_special_tokens=False)["input_ids"]
            for letter in "ABCD":
                full = tokenizer(prefix+letter,add_special_tokens=False)["input_ids"]
                assert full[:-1] == ids and len(full) == len(ids)+1, "Letter ranking requires stable single-token continuations"
            row["prompt_tokens"] = len(ids)
            selected.append(row)
    suite = dict(protocol="balanced-letter-ranking-v1",seed=SEED,cases=selected,
        corpus_sha256=report["corpus_sha256"],seceval_revision=SECEVAL_REVISION,
        seceval_file_sha256=hashlib.sha256(seceval_file.read_bytes()).hexdigest(),
        source_attribution="SecEval: XuanwuAI, CC BY-NC-SA; Pentesting Explanations: Christopher Dickinson, Apache-2.0. Evaluation only.")
    suite_bytes = json.dumps(suite,indent=2,ensure_ascii=False).encode("utf-8")
    (OUT/"suite.json").write_bytes(suite_bytes)
    plan = dict(protocol=suite["protocol"],seed=SEED,suite_sha256=hashlib.sha256(suite_bytes).hexdigest(),
        corpus_sha256=report["corpus_sha256"],seceval_revision=SECEVAL_REVISION,
        seceval_file_sha256=suite["seceval_file_sha256"],
        status="queued_after_training",total_questions=len(selected),excluded=dict(excluded),eligible_pool_sizes=eligible,
        method="Argmax of next-token A/B/C/D logits after a closed reasoning prefix; no generated reasoning. Paired base and adapter in the same NF4 runtime.",
        comparison="Adapted subsets, not official SecEval leaderboard scores or generated-answer pass@1.",
        suites={name:dict(questions=sum(r["id"].startswith(name+":") for r in selected),
            answer_labels=dict(Counter(r["answer_letter"] for r in selected if r["id"].startswith(name+":"))),
            topic_counts=dict(Counter(t for r in selected if r["id"].startswith(name+":") for t in r["topics"]))) for name in pools},
        cases=[{k:r[k] for k in ["id","answer_letter","topics","source","prompt_tokens",
                                "original_choice_indices","question_sha256"]} for r in selected])
    assert all(len(set(s["answer_labels"].values())) == 1 for s in plan["suites"].values())
    assert len(selected) >= 300
    (OUT/"plan.json").write_text(json.dumps(plan,indent=2),encoding="utf-8")
    print(json.dumps({"coverage":{s:{k:report[s][k] for k in ["records","unique_cves","ios27_examples","cisa_distinct_vendors","cisa_distinct_products"]} for s in ["train","eval"]},
        "benchmark_questions":plan["total_questions"],"suites":plan["suites"],"excluded":plan["excluded"]},indent=2))


if __name__ == "__main__":
    prepare()