"""
Solution: scoring prompt versions the way Lab 01's 06_prompt_techniques.py does.
"""
import json
import re

CATEGORIES = {"billing", "bug", "account", "feature_request", "security"}
PRIORITIES = {"P1", "P2", "P3"}


def parse_reply(text: str) -> dict | None:
    """Read a reply the way routing code would: JSON alone, or JSON inside <answer>…</answer>."""
    m = re.search(r"<answer>(.*?)</answer>", text, re.S)
    body = (m.group(1) if m else text).strip()
    try:
        data = json.loads(body)
    except ValueError:
        return None
    if not isinstance(data, dict):
        return None
    if data.get("category") not in CATEGORIES or data.get("priority") not in PRIORITIES:
        return None
    return {"category": data["category"], "priority": data["priority"]}


def is_right(reply: str, label: dict) -> bool:
    got = parse_reply(reply)
    return got is not None and got == label


def score(replies: dict, labels: dict) -> tuple[int, int]:
    """(right, total): right needs both category and priority; an unreadable reply is wrong."""
    right = sum(is_right(replies[t], labels[t]) for t in labels)
    return right, len(labels)


def changed(before: dict, after: dict, labels: dict) -> dict:
    """Which tickets a new version fixed, and which it broke."""
    fixed = [t for t in labels if not is_right(before[t], labels[t]) and is_right(after[t], labels[t])]
    broke = [t for t in labels if is_right(before[t], labels[t]) and not is_right(after[t], labels[t])]
    return {"fixed": fixed, "broke": broke}


def pick(versions: list[dict]) -> str:
    """The version with the best held-back score; among ties, the one with the fewest tokens."""
    best = max(v["holdout"] for v in versions)
    return min((v for v in versions if v["holdout"] == best), key=lambda v: v["tokens"])["name"]