# -*- coding: utf-8 -*-
"""G-gate + XRPL-extra re-check at triage temperature 0.2 (deterministic-ish),
FT adapter only, XRPL items only. Fast follow-up to eval_v8."""
import json, os, re, sys, gc
sys.path.insert(0, "/mnt/c/Users/corov/Desktop/Qwen-Cyber/scripts")
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from eval_assets import xrpl_eval_items
from trackb_part1 import SYSTEM

MODEL = "/home/corov/models/qwen38-9b"
ADAPTER = "/home/corov/cyber/lora_qwen/final_adapter"

GATE_MECH = {
    "Ge1": [r"temINVALID_FLAG", r"offerInDomain|match.time|accountInDomain"],
    "Ge2": [r"XLS-80|4\.2|fail.closed|owner"],
    "Ge6": [r"same.?client|requester|echo", r"HYGIENE|F11|F21"],
    "Ge7": [r"raiseLocalFee", r"while|loop|shutdown|stop_|N11|after the loop|dead"],
    "Ge8": [r"unreachable|dead|isUnlimited", r"HYGIENE|D2|admin|unlimited"],
}

def extract_verdict(t):
    m = re.search(r"###\s*VERDICT\s*\n+\s*([A-Z_]+(?:\s+TRACK:\s?\w+)?(?:\s+PATTERN:\s?N\d+)?)", t)
    return m.group(1).strip() if m else None

tok = AutoTokenizer.from_pretrained(MODEL)
bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4",
                         bnb_4bit_use_double_quant=True,
                         bnb_4bit_compute_dtype=torch.bfloat16, llm_int8_skip_modules=[])
from peft import PeftModel
model = AutoModelForCausalLM.from_pretrained(MODEL, quantization_config=bnb,
                                             torch_dtype=torch.bfloat16,
                                             attn_implementation="sdpa", device_map={"": 0})
ft = PeftModel.from_pretrained(model, ADAPTER)

items = [(i, [{"role": "system", "content": SYSTEM}, {"role": "user", "content": u}], e)
         for i, u, e in xrpl_eval_items() if e is not None]  # only items with expected verdicts
print(f"items: {len(items)}")

all_results = {}
for TEMP in (0.6, 0.2):
    results = []
    for iid, msgs, want in items:
        text = tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
        ids = tok(text, return_tensors="pt", add_special_tokens=False).to(0)
        with torch.no_grad():
            o = ft.generate(**ids, max_new_tokens=1000, do_sample=True, temperature=TEMP,
                            top_p=0.95 if TEMP > 0.5 else 0.9, top_k=20,
                            pad_token_id=tok.pad_token_id, repetition_penalty=1.05)
        t = tok.decode(o[0][ids["input_ids"].shape[1]:], skip_special_tokens=True)
        results.append({"id": iid, "want": want, "out": t})
        print("done", TEMP, iid, flush=True)
    all_results[TEMP] = results
    json.dump(results, open(f"/home/corov/cyber/eval_v8/xrpl_t{TEMP}.json", "w"), indent=1)

for TEMP, results in all_results.items():
    print(f"\n=== G-GATES @ temp {TEMP} ===")
    gpass = 0
    for r in results:
        if r["id"] in GATE_MECH:
            got = extract_verdict(r["out"]) or "?"
            ok_cls = r["want"] in got
            mech = [bool(re.search(rx, r["out"])) for rx in GATE_MECH[r["id"]]]
            ok = ok_cls and any(mech)
            gpass += ok
            print(f"{r['id']}: {'PASS' if ok else 'FAIL'} class={got!r} want={r['want']!r} mech={mech}")
    print(f"=== XRPL extras @ {TEMP} ===")
    hpass = sum(1 for r in results if r["id"] not in GATE_MECH and r["want"] in (extract_verdict(r["out"]) or "?"))
    for r in results:
        if r["id"] not in GATE_MECH:
            got = extract_verdict(r["out"]) or "?"
            print(f"{r['id']}: {'PASS' if r['want'] in got else 'FAIL'} got={got!r} want={r['want']!r}")
    print(f"SUMMARY t{TEMP}: G-gates {gpass}/5 | extras {hpass}/{sum(1 for r in results if r['id'] not in GATE_MECH)}")