#!/usr/bin/env python3
import argparse
import asyncio
import csv
import hashlib
import json
import os
import random
import re
import time
from pathlib import Path
from openai import AsyncOpenAI
from dotenv import load_dotenv
from tqdm.asyncio import tqdm_asyncio

load_dotenv()

api_key = os.getenv("OPENAI_API_KEY")
if not api_key:
    print("Error: OPENAI_API_KEY not set")
    exit(1)

client = AsyncOpenAI(api_key=api_key)

MAX_COMPLETION_TOKENS = 500

# Base system prompts (external markdown files)
with open("judge_prompt.md") as f:
    judge_system_prompt = f.read()

# Load prompts from CSV and deduplicate
prompts = []
seen = set()
with open("chi_test_case.csv") as f:
    reader = csv.DictReader(f)
    for row in reader:
        prompt = row["gpt_description_prompt"].strip()
        if prompt and prompt not in seen:
            prompts.append(prompt)
            seen.add(prompt)

METADATA_REGEX = re.compile(r"\{(.+?)\}")
Q_AND_A_PROMPT = re.compile(
    r"""---\n
\n?
Q:.*\n
A:.*""",
    re.VERBOSE,
)

BANNED_GENRES = frozenset({"instrumental", "catchy", "festive", "energetic", "upbeat"})
ANODYNE_GENRES = ["pop", "rock", "folk", "cumbia", "metal", "funk", "dance", "jazz", "country", "rap"]

LONG_GENRE_EXAMPLE = (
    "This country anthem begins with sparse acoustic steel guitar and builds through verses with gradual "
    "instrumentation. It explodes into a honky-tonk chorus with full band energy. The production balances "
    "authentic country elements with commercial polish, designed for singalong crowd participation in chorus sections."
)

POP_GENRE_INFO = (
    "You are a Grammy-winning pop singer known for your catchy lyrics. Your style is often compared to poetry, "
    "and artists such as the Beatles, Michael Jackson and Elton John. You use rough, unpolished vocabulary and "
    "catchy lines that allow your lyrics to breathe. Don't use commas. Write unexpected and witty lyrics."
)
ROCK_GENRE_INFO = (
    "You are a Grammy-winning rock singer known for your frustrated, questioning and aggressive lyrics.  Your style "
    "is often compared to beat poetry, and artists such as Bob Dylan and Neil Young.  Use rough, unpolished vocabulary, "
    "and phrases rather than complete sentences.  Don't ever use commas."
)
RAP_GENRE_INFO = (
    "You are a Grammy-winning rapper known for your humorous wordplay, internal rhymes, and pop-cultural allusions, "
    "often compared to artists like Jay-Z, Nas, Earl Sweatshirt and MF DOOM.  Use rough, unpolished vocabulary, and "
    "phrases rather than complete sentences.  Use creative similes, but don't use the 'like' or 'as'."
)
FOLK_GENRE_INFO = (
    "You are a famous folk singer known for your narrative storytelling in plaintive and understated lyrics.  Your "
    "style is often compared to poetry, and artists such as Bob Dylan and Joan Baez.  Use rough, unpolished vocabulary "
    "and unexplained references."
)
FUNK_GENRE_INFO = (
    "You are a Grammy-winning funk singer known for your rhythmic, groovy lyrics.  Your style contains spontaneous "
    "ad-libs (like '(oo-yeah!)') and is hypnotizingly repetitive.  Your lyrics make heavy use of slang, humor, and "
    "call and response."
)
CUMBIA_GENRE_INFO = (
    "You are a beloved cumbia singer known for your simple, sparing and traditional lyrics that are easy to sing along "
    "to. You tell stories through your songs. Don't be afraid to shout out ad-libs in parentheses where you give the "
    "audience dance instructions!"
)
METAL_GENRE_INFO = (
    "You are a metal singer infamous for your INDIFFERENCE TO BRUTALITY!! 🤘THE UNIVERSE IS VAST AND CRUEL!  You love "
    "violence, suffering and power, and you often tear the heads off of bats with your teeth.  You use short, brutal "
    "phrases, never complete sentences."
)
DANCE_GENRE_INFO = (
    "You are a Grammy-winning dance pop singer known for your simple and infectious lyrics.  Your lyrics keep the "
    "people on the dance floor all night long.  You use extremely simple vocabulary, short lines that allow your lyrics "
    "to breathe, short phrases rather than complete sentences, and repetition repetition repetition.  You use "
    "spontaneous ad-libs (like '(ooh-yeah!)') and repeat words at the ends of lines."
)
JAZZ_GENRE_INFO = (
    "You are a Grammy-winning Tin Pan Alley jazz singer known for your classic, witty and wry lyrics.  Your charming "
    "style is full of simple, understated wordplay, often compared to songwriters such as the Gershwin brothers, Hoagy "
    "Carmichael and Jerome Kern.  Your lyrics are cheeky and sly, speaking to love, timeless romance and melancholy."
)
COUNTRY_GENRE_INFO = (
    "You are a Grammy-winning Nashville country singer known for your frank, stark and poignant storytelling.  Your "
    "style is simple and full of concrete details about small-town rural life, often compared to songwriters such as "
    "Hank Williams, Townes van Zandt, Doc Watson, Merle Haggard and Dolly Parton."
)

GENRE_PROMPTS = {
    "pop": POP_GENRE_INFO,
    "rock": ROCK_GENRE_INFO,
    "rap": RAP_GENRE_INFO,
    "folk": FOLK_GENRE_INFO,
    "funk": FUNK_GENRE_INFO,
    "cumbia": CUMBIA_GENRE_INFO,
    "metal": METAL_GENRE_INFO,
    "dance": DANCE_GENRE_INFO,
    "jazz": JAZZ_GENRE_INFO,
    "country": COUNTRY_GENRE_INFO,
}

GENRE_KEYWORDS = {
    "rap": ["rap", "hip hop", "hip-hop", "trap"],
    "rock": ["rock", "grunge", "punk"],
    "folk": ["folk", "acoustic", "story"],
    "funk": ["funk", "groove"],
    "cumbia": ["cumbia"],
    "metal": ["metal", "scream"],
    "dance": ["dance", "edm", "club"],
    "jazz": ["jazz", "swing"],
    "country": ["country", "western"],
    "pop": ["pop", "catchy"],
}

JUDGE_FAILURE_KEYS = [
    "instrumental_mismatch",
    "format_failure",
    "language_mismatch",
    "non_english_metadata",
    "ban_violation",
    "empty_lyrics",
    "refusal",
    "tone_mismatch",
]

PROMPT_STRATEGIES = {"5_1", "4o", "raw"}

DEFAULT_ELO_RATING = 1000.0
ELO_K_FACTOR = 24.0
DEFAULT_CONFIDENCE_WEIGHT = 0.5


def _normalize_confidence_value(value):
    try:
        confidence = float(value)
    except (TypeError, ValueError):
        return None
    if confidence != confidence:  # NaN check without math.isnan
        return None
    return max(0.0, min(1.0, confidence))


def _confidence_weight(value):
    normalized = _normalize_confidence_value(value)
    return normalized if normalized is not None else DEFAULT_CONFIDENCE_WEIGHT


def _init_elo_bucket():
    return {
        "rating": DEFAULT_ELO_RATING,
        "matches": 0,
        "confidence_sum": 0.0,
        "confidence_samples": 0,
    }


def load_prompt_text(path):
    with open(path) as f:
        return f.read()


def infer_prompt_strategy(model_name):
    lowered = (model_name or "").lower()
    if "5.1" in lowered or "5_1" in lowered:
        return "5_1"
    if "4o" in lowered:
        return "4o"
    return "raw"


def build_slot_config(label, model_name, prompt_path, strategy=None, key=None):
    prompt_text = load_prompt_text(prompt_path)
    inferred = strategy or infer_prompt_strategy(model_name)
    prompt_strategy = inferred if inferred in PROMPT_STRATEGIES else "raw"
    marker = f"{os.path.basename(prompt_path)}:{hashlib.md5(prompt_text.encode()).hexdigest()[:8]}"
    return {
        "key": key or label.lower(),
        "label": label,
        "model": model_name,
        "prompt_path": prompt_path,
        "prompt_text": prompt_text,
        "prompt_strategy": prompt_strategy,
        "prompt_marker": marker,
    }


def _strip_lyrics_prompt_hallucinations(lyrics):
    return Q_AND_A_PROMPT.sub("", lyrics)


def get_stanzas(text):
    """Split lyrics into stanzas based on blank lines."""
    if not text:
        return []
    return [stanza.strip() for stanza in text.split("\n\n") if stanza.strip()]


def truncate_lyrics(text, num_stanzas=3):
    """Keep only the first num_stanzas stanzas."""
    stanzas = get_stanzas(text)
    if not stanzas:
        return "", False
    truncated_stanzas = stanzas[:num_stanzas]
    was_truncated = len(stanzas) > num_stanzas
    return "\n\n".join(truncated_stanzas), was_truncated


def detect_genre(user_prompt):
    lowered = user_prompt.lower()
    for genre, keywords in GENRE_KEYWORDS.items():
        if any(keyword in lowered for keyword in keywords):
            return genre, GENRE_PROMPTS[genre]
    return None, None


def is_instrumental_prompt(user_prompt):
    return "instrumental" in user_prompt.lower()


def construct_5_1_system_prompt(base_prompt, genre_block):
    """Prepend genre-specific persona to the 5.1 base prompt."""
    return f"{genre_block}\n\n{base_prompt}"


def construct_4o_system_prompt(base_prompt, user_prompt):
    genre_key, genre_block = detect_genre(user_prompt)
    prompt_parts = []
    if genre_block:
        prompt_parts.append(genre_block)
    prompt_parts.append(base_prompt)
    prompt = "\n\n".join(prompt_parts)
    return prompt, genre_key or "pop"


def format_prompt_preview(text, max_len=60):
    """Collapse whitespace and truncate long prompts for logging."""
    if not text:
        return ""
    collapsed = " ".join(text.split())
    if len(collapsed) <= max_len:
        return collapsed
    return f"{collapsed[: max_len - 3].rstrip()}..."


def build_slot_system_prompt(slot_config, user_prompt):
    strategy = slot_config.get("prompt_strategy") or "raw"
    prompt_text = slot_config.get("prompt_text") or ""
    if strategy == "4o":
        return construct_4o_system_prompt(prompt_text, user_prompt)
    genre_key, genre_block = detect_genre(user_prompt)
    if strategy == "5_1" and genre_block:
        return construct_5_1_system_prompt(prompt_text, genre_block), genre_key
    return prompt_text, genre_key


def _empty_failure_flags():
    base = {key: False for key in JUDGE_FAILURE_KEYS}
    base["details"] = ""
    return base


def normalize_failure_flags(raw_flags):
    flags = _empty_failure_flags()
    if not raw_flags:
        return flags
    for key in JUDGE_FAILURE_KEYS:
        flags[key] = bool(raw_flags.get(key, False))
    if raw_flags.get("details"):
        flags["details"] = str(raw_flags["details"])
    return flags


def _ensure_model_tracking(tracking, model):
    if model is None or tracking is None:
        return
    if model not in tracking["model_stats"]:
        tracking["model_stats"][model] = {"wins": 0, "losses": 0, "ties": 0}
    if model not in tracking["failure_counts"]:
        tracking["failure_counts"][model] = {key: 0 for key in JUDGE_FAILURE_KEYS}
    if model not in tracking["model_observations"]:
        tracking["model_observations"][model] = 0
    if "elo_ratings" not in tracking:
        tracking["elo_ratings"] = {}
    if model not in tracking["elo_ratings"]:
        tracking["elo_ratings"][model] = _init_elo_bucket()


def _apply_elo_result(tracking, model_a, model_b, score_a, confidence):
    if not tracking or not model_a or not model_b:
        return
    _ensure_model_tracking(tracking, model_a)
    _ensure_model_tracking(tracking, model_b)
    bucket_a = tracking["elo_ratings"][model_a]
    bucket_b = tracking["elo_ratings"][model_b]
    rating_a = bucket_a["rating"]
    rating_b = bucket_b["rating"]
    expected_a = 1 / (1 + 10 ** ((rating_b - rating_a) / 400))
    expected_b = 1 - expected_a
    score_b = 1 - score_a if score_a in (0.0, 1.0) else 0.5
    weight = _confidence_weight(confidence)
    k = ELO_K_FACTOR * weight
    bucket_a["rating"] = rating_a + k * (score_a - expected_a)
    bucket_b["rating"] = rating_b + k * (score_b - expected_b)
    bucket_a["matches"] += 1
    bucket_b["matches"] += 1
    normalized_conf = _normalize_confidence_value(confidence)
    if normalized_conf is not None:
        bucket_a["confidence_sum"] += normalized_conf
        bucket_a["confidence_samples"] += 1
        bucket_b["confidence_sum"] += normalized_conf
        bucket_b["confidence_samples"] += 1


def format_candidate_block(label, payload):
    status = payload.get("status") or "unknown"
    markers = []
    if payload.get("was_token_truncated"):
        markers.append("token_truncated")
    if payload.get("was_stanza_truncated"):
        markers.append("stanza_truncated")
    marker_text = f" ({', '.join(markers)})" if markers else ""
    output_text = payload.get("formatted_output") or ""
    return (
        f"Candidate {label} — status: {status}{marker_text}\n"
        "```\n"
        f"{output_text}\n"
        "```\n"
    )


def build_judge_message(user_prompt, candidates):
    parts = [
        "User prompt:",
        user_prompt.strip(),
        "",
        format_candidate_block("A", candidates["A"]),
        format_candidate_block("B", candidates["B"]),
        "Respond with strict JSON following the specified schema.",
    ]
    return "\n".join(parts)


def make_candidate_payload(completion):
    metadata = completion.get("metadata", {}) or {}
    title = metadata.get("title", "").strip()
    style = metadata.get("style", "").strip()
    lyrics = (completion.get("content") or "").strip()
    formatted_output = "\n".join(
        [
            f"{{{title}}}" if title else "{ }",
            f"{{{style}}}" if style else "{ }",
            "",
            lyrics or "",
        ]
    ).strip()
    return {
        "title": title,
        "style": style,
        "lyrics": lyrics,
        "status": completion.get("status"),
        "formatted_output": formatted_output,
        "was_token_truncated": metadata.get("was_token_truncated"),
        "was_stanza_truncated": metadata.get("was_stanza_truncated"),
    }


def init_judge_tracking(model_keys, judge_model_name):
    return {
        "judge_model": judge_model_name,
        "total_judged": 0,
        "model_stats": {model: {"wins": 0, "losses": 0, "ties": 0} for model in model_keys},
        "failure_counts": {model: {key: 0 for key in JUDGE_FAILURE_KEYS} for model in model_keys},
        "model_observations": {model: 0 for model in model_keys},
        "elo_ratings": {model: _init_elo_bucket() for model in model_keys},
    }


def update_judge_tracking(tracking, judge_payload):
    if not tracking or not judge_payload:
        return
    assignments = judge_payload.get("model_assignment", {})
    if len(assignments) < 2:
        return
    tracking["total_judged"] += 1

    failures = judge_payload.get("failures", {})
    for label, model in assignments.items():
        _ensure_model_tracking(tracking, model)
        if model not in tracking["model_observations"]:
            continue
        tracking["model_observations"][model] += 1
        label_failures = failures.get(label, {})
        for key in JUDGE_FAILURE_KEYS:
            if label_failures.get(key):
                tracking["failure_counts"][model][key] += 1

    winner_label = judge_payload.get("winner", "tie")
    confidence_value = judge_payload.get("confidence")
    if winner_label in assignments:
        winner_model = assignments[winner_label]
        loser_label = "B" if winner_label == "A" else "A"
        loser_model = assignments.get(loser_label)
        _ensure_model_tracking(tracking, winner_model)
        _ensure_model_tracking(tracking, loser_model)
        if winner_model in tracking["model_stats"]:
            tracking["model_stats"][winner_model]["wins"] += 1
        if loser_model and loser_model in tracking["model_stats"]:
            tracking["model_stats"][loser_model]["losses"] += 1
        if winner_model and loser_model:
            _apply_elo_result(tracking, winner_model, loser_model, 1.0, confidence_value)
    else:
        for model in assignments.values():
            _ensure_model_tracking(tracking, model)
            if model in tracking["model_stats"]:
                tracking["model_stats"][model]["ties"] += 1
        model_a = assignments.get("A")
        model_b = assignments.get("B")
        if model_a and model_b:
            _apply_elo_result(tracking, model_a, model_b, 0.5, confidence_value)


def finalize_judge_summary(tracking):
    if not tracking or tracking["total_judged"] == 0:
        return None
    failure_rates = {}
    for model, counts in tracking["failure_counts"].items():
        obs = tracking["model_observations"].get(model, 0)
        failure_rates[model] = {}
        for key, count in counts.items():
            rate = count / obs if obs else 0.0
            failure_rates[model][key] = {"count": count, "rate": round(rate, 4)}
    elo_summary = {}
    for model, metrics in tracking.get("elo_ratings", {}).items():
        if not metrics:
            continue
        avg_conf = (
            metrics["confidence_sum"] / metrics["confidence_samples"]
            if metrics["confidence_samples"]
            else None
        )
        elo_summary[model] = {
            "rating": round(metrics["rating"], 2),
            "matches": metrics["matches"],
            "avg_confidence": round(avg_conf, 4) if avg_conf is not None else None,
        }
    return {
        "judge_model": tracking["judge_model"],
        "total_judged": tracking["total_judged"],
        "model_stats": tracking["model_stats"],
        "failure_rates": failure_rates,
        "elo_ratings": elo_summary,
    }


def parse_completion_output(raw_text, make_instrumental=False, lyrics_model=None):
    if make_instrumental and not raw_text.strip():
        return {
            "status": "instrumental",
            "title": "",
            "style": "",
            "tags": ["instrumental"],
            "lyrics": "[Instrumental]",
            "lyrics_raw": "[Instrumental]",
            "was_token_truncated": False,
            "was_stanza_truncated": False,
            "line_count_raw": 1,
            "line_count_final": 1,
        }

    matches = METADATA_REGEX.findall(raw_text)
    title = matches[0].strip() if matches else ""
    style = matches[1].strip() if len(matches) > 1 else ""
    raw_tags = {tag.strip().lower() for tag in style.split(",") if tag.strip()}
    tags = [tag for tag in raw_tags if tag not in BANNED_GENRES] or [random.choice(ANODYNE_GENRES)]

    lyrics_raw = METADATA_REGEX.sub("", raw_text).strip()
    lyrics_raw = _strip_lyrics_prompt_hallucinations(lyrics_raw)
    line_count_raw = len(lyrics_raw.split("\n")) if lyrics_raw else 0
    should_truncate = not (lyrics_model and "gpt-5.1" in lyrics_model)
    if should_truncate:
        lyrics, was_stanza_truncated = truncate_lyrics(lyrics_raw)
    else:
        lyrics = lyrics_raw
        was_stanza_truncated = False
    was_token_truncated = bool(raw_text and not raw_text.rstrip().endswith((".", "!", "?", '"', "'", "]", ")")))

    status = "ok"
    if not matches:
        status = "missing_metadata"
    if not lyrics:
        status = "empty_lyrics"

    if make_instrumental and raw_text.strip():
        lyrics = "[Instrumental]"
        status = "instrumental"
        if "instrumental" not in tags:
            tags = ["instrumental"] + tags

    line_count_final = len(lyrics.split("\n")) if lyrics else 0

    return {
        "status": status,
        "title": title,
        "style": style,
        "tags": tags,
        "lyrics": lyrics,
        "lyrics_raw": lyrics_raw,
        "was_token_truncated": was_token_truncated,
        "was_stanza_truncated": was_stanza_truncated,
        "line_count_raw": line_count_raw,
        "line_count_final": line_count_final,
    }

def hash_file(path):
    return hashlib.md5(Path(path).read_bytes()).hexdigest()

def get_latest_output():
    files = sorted(Path(".").glob("outputs_*.json"), reverse=True)
    if not files:
        return None
    try:
        with open(files[0]) as f:
            return json.load(f)
    except:  # noqa: E722
        return None

async def stream_completion_once(slot_config, user_prompt, semaphore, log_status=None, prompt_preview=""):
    logger = log_status or print
    prompt_label = prompt_preview or format_prompt_preview(user_prompt)
    async with semaphore:
        model_name = slot_config["model"]
        system_prompt, genre_hint = build_slot_system_prompt(slot_config, user_prompt)
        make_instrumental = is_instrumental_prompt(user_prompt)

        params = {
            "model": model_name,
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt},
            ],
            "max_completion_tokens": MAX_COMPLETION_TOKENS,
            "temperature": 1.0,
            "n": 1,
            "stream": True,
            "stream_options": {"include_usage": True},
        }

        if model_name == "gpt-5.1":
            params["reasoning_effort"] = "none"

        start_time = time.time()
        ttft = None
        full_content = ""
        usage_stats = None

        stream = await client.chat.completions.create(**params)

        async for chunk in stream:
            if hasattr(chunk, "usage") and chunk.usage:
                usage_stats = chunk.usage

            if chunk.choices and chunk.choices[0].delta.content:
                full_content += chunk.choices[0].delta.content
                if ttft is None:
                    ttft = time.time() - start_time

        total_duration = time.time() - start_time
        completion_tokens = usage_stats.completion_tokens if usage_stats else 0

        tokens_per_second = 0
        if ttft and total_duration > ttft and completion_tokens > 1:
            tokens_per_second = (completion_tokens - 1) / (total_duration - ttft)
        elif total_duration > 0 and completion_tokens > 0:
            tokens_per_second = completion_tokens / total_duration

        metric_bits = []
        if ttft is not None:
            metric_bits.append(f"TTFT {ttft:.2f}s")
        if tokens_per_second:
            metric_bits.append(f"{tokens_per_second:.1f} tok/s")
        metrics_display = f" ({', '.join(metric_bits)})" if metric_bits else ""
        logger(f"✓ {model_name} • {prompt_label}{metrics_display}")

        parsed = parse_completion_output(
            full_content,
            make_instrumental=make_instrumental,
            lyrics_model=model_name,
        )

        return {
            "model": model_name,
            "prompt": user_prompt,
            "content": parsed["lyrics"],
            "status": parsed["status"],
            "prompt_path": slot_config.get("prompt_path"),
            "prompt_marker": slot_config.get("prompt_marker"),
            "prompt_strategy": slot_config.get("prompt_strategy"),
            "metadata": {
                "title": parsed["title"],
                "style": parsed["style"],
                "tags": parsed["tags"],
                "genre_hint": genre_hint,
                "raw_output": full_content,
                "lyrics_raw": parsed.get("lyrics_raw"),
                "was_token_truncated": parsed.get("was_token_truncated"),
                "was_stanza_truncated": parsed.get("was_stanza_truncated"),
                "line_count_raw": parsed.get("line_count_raw"),
                "line_count_final": parsed.get("line_count_final"),
            },
            "metrics": {
                "ttft": round(ttft, 3) if ttft else None,
                "tokens_per_second": round(tokens_per_second, 2),
                "total_duration": round(total_duration, 3),
                "total_tokens": completion_tokens,
            },
        }


async def stream_completion_with_retry(
    slot_config,
    user_prompt,
    semaphore,
    log_status=None,
    prompt_preview="",
    max_retries=2,
):
    attempts = 0
    last_error = None
    while attempts <= max_retries:
        try:
            return await stream_completion_once(
                slot_config,
                user_prompt,
                semaphore,
                log_status=log_status,
                prompt_preview=prompt_preview,
            )
        except Exception as err:
            attempts += 1
            last_error = err
            logger = log_status or print
            logger(
                f"✗ {slot_config['model']} • {prompt_preview or format_prompt_preview(user_prompt)} "
                f"— attempt {attempts}/{max_retries + 1}: {err}"
            )
            if attempts > max_retries:
                break
            await asyncio.sleep(min(2 ** attempts, 5))

    return {
        "model": slot_config["model"],
        "prompt": user_prompt,
        "content": f"ERROR: {last_error}",
        "status": "error",
        "prompt_path": slot_config.get("prompt_path"),
        "prompt_marker": slot_config.get("prompt_marker"),
        "prompt_strategy": slot_config.get("prompt_strategy"),
        "metadata": {"genre_hint": None},
        "metrics": {"error": str(last_error) if last_error else "unknown"},
    }

async def judge_prompt_pair(
    prompt_text,
    slot_payloads,
    judge_cfg,
    semaphore,
    log_status=None,
    prompt_preview="",
):
    logger = log_status or print
    prompt_label = prompt_preview or format_prompt_preview(prompt_text)
    available = [
        (slot_key, data)
        for slot_key, data in slot_payloads.items()
        if data and data.get("content") is not None and data.get("status") != "error"
    ]
    if len(available) < 2:
        return None

    random.shuffle(available)
    selected = available[:2]

    assignments = {}
    slot_assignments = {}
    candidates = {}
    for label, (slot_key, data) in zip(("A", "B"), selected):
        assignments[label] = data.get("model")
        slot_assignments[label] = slot_key
        candidates[label] = make_candidate_payload(data)

    user_message = build_judge_message(prompt_text, candidates)
    params = {
        "model": judge_cfg["model"],
        "messages": [
            {"role": "system", "content": judge_cfg["prompt"]},
            {"role": "user", "content": user_message},
        ],
        "temperature": judge_cfg["temperature"],
        "max_completion_tokens": judge_cfg["max_tokens"],
        "response_format": {"type": "json_object"},
    }

    attempts = 0
    max_retries = judge_cfg["max_retries"]
    while attempts <= max_retries:
        try:
            async with semaphore:
                response = await client.chat.completions.create(**params)
            content = response.choices[0].message.content
            parsed = json.loads(content)
            failures = parsed.get("failures", {})
            normalized = {
                "model_assignment": assignments,
                "slot_assignment": slot_assignments,
                "winner": parsed.get("winner", "tie"),
                "confidence": parsed.get("confidence"),
                "commentary": (parsed.get("commentary") or "").strip(),
                "failures": {
                    "A": normalize_failure_flags(failures.get("A")),
                    "B": normalize_failure_flags(failures.get("B")),
                },
            }
            logger(
                f"⚖️ {judge_cfg['model']} • {prompt_label} — winner {normalized['winner']} (conf {normalized['confidence']})"
            )
            return normalized
        except Exception as e:
            attempts += 1
            if attempts > max_retries:
                logger(f"✗ Judge • {prompt_label} — {e}")
                return {
                    "model_assignment": assignments,
                    "slot_assignment": slot_assignments,
                    "winner": "tie",
                    "confidence": None,
                    "commentary": f"Judge error: {e}",
                    "failures": {
                        "A": _empty_failure_flags(),
                        "B": _empty_failure_flags(),
                    },
                    "error": str(e),
                }
            await asyncio.sleep(min(2 ** attempts, 5))


async def evaluate_prompt(
    index,
    prompt_text,
    slot_configs,
    completion_semaphore,
    judge_enabled,
    judge_cfg,
    judge_semaphore,
    judge_state,
    judge_lock,
    log_status,
):
    prompt_preview = format_prompt_preview(prompt_text)
    entry = {"prompt": prompt_text, "completions": {}}
    slot_payloads = {}

    for slot in slot_configs:
        payload = await stream_completion_with_retry(
            slot,
            prompt_text,
            completion_semaphore,
            log_status=log_status,
            prompt_preview=prompt_preview,
        )
        entry["completions"][slot["key"]] = payload
        slot_payloads[slot["key"]] = payload

    should_judge = judge_enabled and random.random() <= judge_cfg.get("sample_rate", 1.0)
    if should_judge:
        judge_result = await judge_prompt_pair(
            prompt_text,
            slot_payloads,
            judge_cfg,
            judge_semaphore,
            log_status=log_status,
            prompt_preview=prompt_preview,
        )
        if judge_result:
            entry["judge"] = judge_result
            if judge_lock:
                async with judge_lock:
                    update_judge_tracking(judge_state, judge_result)
            else:
                update_judge_tracking(judge_state, judge_result)

    return index, entry


async def process_prompts(
    prompts_to_process,
    slot_configs,
    max_concurrent=8,
    judge_options=None,
    checkpoint_interval=50,
    checkpoint_path=None,
    system_prompt_markers=None,
    slot_metadata=None,
    prompts_hash=None,
):
    total = len(prompts_to_process)
    if total == 0:
        return [], None

    log_status = tqdm_asyncio.write
    completion_semaphore = asyncio.Semaphore(max(1, max_concurrent))

    judge_cfg = judge_options or {}
    judge_enabled = judge_cfg.get("enabled", False) and judge_cfg.get("prompt")
    judge_state = None
    judge_semaphore = None
    if judge_enabled:
        unique_models = sorted({slot["model"] for slot in slot_configs})
        judge_state = init_judge_tracking(unique_models, judge_cfg["model"])
        judge_semaphore = asyncio.Semaphore(max(1, judge_cfg.get("max_concurrent", 4)))
        judge_lock = asyncio.Lock()
    else:
        judge_cfg = {}
        judge_lock = None

    ordered_results = []
    buffer = [None] * total
    next_commit = 0
    completed = 0
    result_lock = asyncio.Lock()
    prompt_queue = asyncio.Queue()

    for idx, prompt in enumerate(prompts_to_process):
        prompt_queue.put_nowait((idx, prompt))

    async def worker():
        nonlocal next_commit, completed
        while True:
            try:
                idx, prompt = await prompt_queue.get()
            except asyncio.CancelledError:
                break

            try:
                index, result = await evaluate_prompt(
                    idx,
                    prompt,
                    slot_configs,
                    completion_semaphore,
                    judge_enabled,
                    judge_cfg,
                    judge_semaphore,
                    judge_state,
                    judge_lock,
                    log_status,
                )
                async with result_lock:
                    buffer[index] = result
                    while next_commit < total and buffer[next_commit] is not None:
                        ordered_results.append(buffer[next_commit])
                        buffer[next_commit] = None
                        next_commit += 1
                        completed += 1

                        if (
                            checkpoint_interval
                            and checkpoint_path
                            and completed % checkpoint_interval == 0
                        ):
                            snapshot = (
                                finalize_judge_summary(judge_state)
                                if judge_state
                                else None
                            )
                            write_results_file(
                                checkpoint_path,
                                system_prompt_markers,
                                prompts_hash,
                                slot_metadata,
                                ordered_results,
                                snapshot,
                            )
            finally:
                prompt_queue.task_done()

    worker_count = min(total, max(1, max_concurrent))
    workers = [asyncio.create_task(worker()) for _ in range(worker_count)]

    await prompt_queue.join()

    for task in workers:
        task.cancel()
    await asyncio.gather(*workers, return_exceptions=True)

    return ordered_results, judge_state

def find_next_version():
    version = 1
    while Path(f"outputs_{version:03d}.json").exists():
        version += 1
    return version


def write_results_file(
    path,
    system_prompt_markers,
    prompts_hash,
    slot_metadata,
    results,
    judge_summary,
):
    payload = {
        "system_prompts": system_prompt_markers or {},
        "slots": slot_metadata or {},
        "prompts_hash": prompts_hash,
        "results": results,
    }
    if judge_summary:
        payload["judge_summary"] = judge_summary
    with open(path, "w") as f:
        json.dump(payload, f, indent=2)

async def main():
    parser = argparse.ArgumentParser(description="Eval harness for prompt comparison")
    parser.add_argument("-n", "--num-prompts", type=int, help="Number of prompts to process (default: all)")
    parser.add_argument("--max-concurrent", type=int, default=8, help="Max concurrent generation requests")
    parser.add_argument("--no-judge", action="store_true", help="Disable LLM judge comparisons")
    parser.add_argument("--judge-model", default="gpt-5.1", help="Model used for judging outputs")
    parser.add_argument("--judge-temperature", type=float, default=0.3, help="Sampling temperature for judge calls")
    parser.add_argument("--judge-max-tokens", type=int, default=400, help="Max completion tokens for judge responses")
    parser.add_argument("--judge-max-retries", type=int, default=2, help="Number of retries for failed judge calls")
    parser.add_argument(
        "--judge-max-concurrent",
        type=int,
        default=4,
        help="Maximum concurrent judge API calls",
    )
    parser.add_argument(
        "--judge-sample-rate",
        type=float,
        default=1.0,
        help="Fraction of prompts to pass through the judge (0-1)",
    )
    parser.add_argument("--model-a", default="gpt-5.1", help="API model used for Slot A")
    parser.add_argument("--model-b", default="gpt-4o", help="API model used for Slot B")
    parser.add_argument("--prompt-a", default="5_1.md", help="System prompt file for Slot A")
    parser.add_argument("--prompt-b", default="4o.md", help="System prompt file for Slot B")
    parser.add_argument(
        "--checkpoint-interval",
        type=int,
        default=50,
        help="Write partial outputs every N completed prompts (0 disables checkpointing)",
    )
    args = parser.parse_args()
    prompts_hash = hash_file("chi_test_case.csv")

    slot_a = build_slot_config("Slot A", args.model_a, args.prompt_a, key="slot_a")
    slot_b = build_slot_config("Slot B", args.model_b, args.prompt_b, key="slot_b")
    slot_configs = [slot_a, slot_b]
    slot_metadata = {
        slot["key"]: {
            "label": slot["label"],
            "model": slot["model"],
            "prompt_path": slot["prompt_path"],
            "prompt_marker": slot["prompt_marker"],
            "prompt_strategy": slot["prompt_strategy"],
        }
        for slot in slot_configs
    }

    judge_marker = f"judge:{hashlib.md5(judge_system_prompt.encode()).hexdigest()[:8]}"
    system_prompt_markers = {
        slot["key"]: slot["prompt_marker"] for slot in slot_configs
    }
    system_prompt_markers["judge"] = judge_marker

    prompts_to_process = prompts[:args.num_prompts] if args.num_prompts else prompts
    judge_sample_rate = min(1.0, max(0.0, args.judge_sample_rate))
    judge_options = None
    if not args.no_judge and judge_sample_rate > 0:
        judge_options = {
            "enabled": True,
            "model": args.judge_model,
            "prompt": judge_system_prompt,
            "temperature": args.judge_temperature,
            "max_tokens": args.judge_max_tokens,
            "max_retries": max(0, args.judge_max_retries),
            "sample_rate": judge_sample_rate,
            "max_concurrent": max(1, args.judge_max_concurrent),
        }

    checkpoint_interval = max(0, args.checkpoint_interval)
    version = find_next_version()
    output_file = f"outputs_{version:03d}.json"

    print(f"Processing {len(prompts_to_process)} prompts with output -> {output_file}")

    results, judge_state = await process_prompts(
        prompts_to_process,
        slot_configs,
        max_concurrent=args.max_concurrent,
        judge_options=judge_options,
        checkpoint_interval=checkpoint_interval,
        checkpoint_path=output_file if checkpoint_interval else None,
        system_prompt_markers=system_prompt_markers,
        slot_metadata=slot_metadata,
        prompts_hash=prompts_hash,
    )

    judge_summary = finalize_judge_summary(judge_state)
    write_results_file(
        output_file,
        system_prompt_markers,
        prompts_hash,
        slot_metadata,
        results,
        judge_summary,
    )

    print(f"\n✓ Results saved to {output_file}")

if __name__ == "__main__":
    asyncio.run(main())