"""
eval/prepare_golden_sample_batch2.py

Samples a fresh, non-overlapping batch of rows from ealvaradob/phishing-dataset's
texts.json and appends them to the existing eval/golden_sample.csv, continuing the
id sequence. Purpose: get enough additional report_text-bearing rows for
LLM-as-Judge scoring (Day 6, batch 2) while also modestly extending the
golden-dataset accuracy sample size as a side benefit.

Safe to run once. Dedup is by exact text match against the existing sample.
"""

import csv
import json
import random
from pathlib import Path
import re

from huggingface_hub import hf_hub_download

GOLDEN_SAMPLE_PATH = Path(__file__).resolve().parent / "golden_sample.csv"
BATCH2_SIZE = 16          # ~8 legit / 8 phishing
RANDOM_SEED = 20260705    # distinct from Day 6's original seed, on purpose

ID_PATTERN = re.compile(r"^(?P<prefix>.*?)(?P<num>\d+)$")


def load_existing():
    if not GOLDEN_SAMPLE_PATH.exists():
        raise SystemExit(f"{GOLDEN_SAMPLE_PATH} not found — run prepare_golden_sample.py first.")
    with open(GOLDEN_SAMPLE_PATH, "r", encoding="utf-8", newline="") as f:
        return list(csv.DictReader(f))


def parse_id(id_str: str):
    """Splits e.g. 'row_0000' into ('row_', 0, width=4)."""
    m = ID_PATTERN.match(id_str)
    if not m:
        raise ValueError(f"Unrecognized id format: {id_str!r}")
    prefix = m.group("prefix")
    num_str = m.group("num")
    return prefix, int(num_str), len(num_str)


def format_id(prefix: str, num: int, width: int) -> str:
    return f"{prefix}{str(num).zfill(width)}"

def main():
    existing = load_existing()
    existing_texts = {r["text"] for r in existing}

    parsed_ids = [parse_id(r["id"]) for r in existing]
    prefix, _, width = parsed_ids[0]           # assume consistent format across all rows
    max_num = max(num for _, num, _ in parsed_ids)
    print(f"Existing sample: {len(existing)} rows, max id = {format_id(prefix, max_num, width)}")

    print("Downloading texts.json from ealvaradob/phishing-dataset...")
    path = hf_hub_download(
        repo_id="ealvaradob/phishing-dataset",
        filename="texts.json",
        repo_type="dataset",
    )
    with open(path, "r", encoding="utf-8") as f:
        data = json.load(f)

    legit = [d for d in data if int(d["label"]) == 0 and d["text"] not in existing_texts]
    phish = [d for d in data if int(d["label"]) == 1 and d["text"] not in existing_texts]

    rng = random.Random(RANDOM_SEED)
    rng.shuffle(legit)
    rng.shuffle(phish)

    half = BATCH2_SIZE // 2
    picked = legit[:half] + phish[:BATCH2_SIZE - half]
    rng.shuffle(picked)

    print(f"Picked {len(picked)} new non-overlapping rows "
          f"({half} legit / {BATCH2_SIZE - half} phishing).")

    with open(GOLDEN_SAMPLE_PATH, "a", encoding="utf-8", newline="") as f:
        writer = csv.writer(f)
        next_num = max_num + 1
        first_id, last_id = None, None
        for item in picked:
            new_id = format_id(prefix, next_num, width)
            if first_id is None:
                first_id = new_id
            last_id = new_id
            writer.writerow([new_id, item["text"], int(item["label"])])
            next_num += 1

    print(f"Appended rows {first_id}..{last_id} to {GOLDEN_SAMPLE_PATH}")
    print("Next: rerun `python eval/run_golden_eval.py` — it will skip ids "
          f"already done and process only the new ones.")


if __name__ == "__main__":
    main()