"""
eval/fix_batch2_labels.py

One-off fix for a batch2 data bug: prepare_golden_sample_batch2.py wrote
raw numeric HF dataset labels ('1'/'0') for the two new rows (row_0060,
row_0061) instead of the string convention ('phishing'/'legitimate') used
by the original 60 rows. score_golden_eval.py only recognizes the exact
string "phishing", so the numeric labels were silently mis-binned as
"legitimate" -- making a correct model prediction (row_0060, correctly
flagged high-risk) look like a false positive in the confusion matrix.

Confirmed mapping (verified against actual row content on Day 6):
  '1' -> phishing    (row_0060: obfuscated software-spam)
  '0' -> legitimate  (row_0061: energy-scheduling log)

This script fixes BOTH files that hold a copy of true_label:
  - eval/golden_sample.csv        (label column)
  - eval/golden_eval_results.jsonl (true_label field)

Safe to re-run: only touches rows whose label is exactly '1' or '0',
so already-fixed rows (now 'phishing'/'legitimate') are left untouched.
"""

import csv
import json
from pathlib import Path

LABEL_MAP = {"1": "phishing", "0": "legitimate"}

CSV_PATH = Path("eval/golden_sample.csv")
JSONL_PATH = Path("eval/golden_eval_results.jsonl")


def fix_csv():
    rows = list(csv.DictReader(open(CSV_PATH, encoding="utf-8")))
    fieldnames = list(rows[0].keys())
    changed = 0
    for r in rows:
        if r.get("label") in LABEL_MAP:
            r["label"] = LABEL_MAP[r["label"]]
            changed += 1
    with open(CSV_PATH, "w", encoding="utf-8", newline="") as f:
        w = csv.DictWriter(f, fieldnames=fieldnames)
        w.writeheader()
        w.writerows(rows)
    print(f"[csv] fixed {changed} row(s) in {CSV_PATH}")


def fix_jsonl():
    rows = []
    with open(JSONL_PATH, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                rows.append(json.loads(line))
    changed = 0
    for r in rows:
        if r.get("true_label") in LABEL_MAP:
            r["true_label"] = LABEL_MAP[r["true_label"]]
            changed += 1
    with open(JSONL_PATH, "w", encoding="utf-8") as f:
        for r in rows:
            f.write(json.dumps(r) + "\n")
    print(f"[jsonl] fixed {changed} row(s) in {JSONL_PATH}")


if __name__ == "__main__":
    fix_csv()
    fix_jsonl()
    print("done")