# -*- coding: utf-8 -*-
"""Reconstruct the exact v2 training dataset (the one the shipped 9B was trained on):
Track A full 280 + (Track B uniques from parts 1-4 incl. C-labs/boundary) x3 = 395 total.
v3 additions (part5 variants, A-subsample) are NOT included."""
import json, sys
sys.path.insert(0, "/mnt/c/Users/corov/Desktop/Qwen-Cyber/scripts")
import trackb_part1, trackb_part2, trackb_part3, trackb_part4, build_trackC  # registers on S
from trackb_part1 import S

# boundary adds (same as build_dataset v2 did)
from trackb_part1 import add as _add
C_PREAMBLE = ("Analyze this code-security question (local lab / review context). "
              "Deliver: white-box analysis, trigger, exploit writeup, verdict.\n\n")
exec(open("/mnt/c/Users/corov/Desktop/Qwen-Cyber/scripts/build_dataset.py", encoding="utf-8")
     .read().split("# ---- 2 C-boundary samples")[1].split("def main()")[0]
     .replace('add("C_bound_fortify_not_safe"', '_add("C_bound_fortify_not_safe", preamble=C_PREAMBLE,')
     .replace(", preamble=C_PREAMBLE,\n", ", preamble=C_PREAMBLE,\n")
    ) if False else None

DATA = "/home/corov/cyber/data"
trackA = [json.loads(l) for l in open(f"{DATA}/trackA.jsonl", encoding="utf-8")]

# rebuild boundary samples by calling the same code as build_dataset (they live inline there)
# simplest: extract their texts from the current train_all.jsonl (v3) - boundary ids are stable
v3 = [json.loads(l) for l in open(f"{DATA}/train_all.jsonl", encoding="utf-8")]
boundary = [s for s in v3 if s["id"].startswith("C_bound")]

seen, mergedB, merged = set(), [], []
for s in trackA + S + boundary:
    if s["id"] in seen:
        continue
    seen.add(s["id"])
    if s["id"].startswith(("A_", "C_")):
        merged.append(s)
    else:
        mergedB.append(s)
merged = merged + mergedB * 3

nA = sum(1 for s in merged if s["id"].startswith("A_"))
nC = sum(1 for s in merged if s["id"].startswith("C_"))
nB = len(merged) - nA - nC
print(f"A={nA} B_unique={len(mergedB)} -> B_x3={nB} C={nC} TOTAL={len(merged)}")
assert nA == 280, "Track A must be full 280 (v2)"
assert len(merged) == 395, f"v2 total must be 395, got {len(merged)}"

with open("/home/corov/cyber/data/train_all_v2_shipped.jsonl", "w", encoding="utf-8", newline="\n") as f:
    for s in merged:
        f.write(json.dumps(s, ensure_ascii=False) + "\n")
print("WROTE train_all_v2_shipped.jsonl")

# sanity: G-set verdicts + schema
import re
bad = []
for s in merged:
    a = s["messages"][-1]["content"]
    if not ("### TRIGGER" in a and "### EXPLOIT WRITEUP" in a and "### VERDICT" in a):
        bad.append(s["id"])
print("schema bad:", bad[:5], "count:", len(bad))