"""OBLITERATUS Tourney — March Madness-style tournament to find the best abliteration method.

Run all methods head-to-head in elimination rounds.  The winner gets auto-pushed
to HuggingFace Hub so the community can use the best possible abliteration.

Usage (CLI):
    obliteratus tourney meta-llama/Llama-3.1-8B-Instruct --hub-org my-org

Usage (Python):
    from obliteratus.tourney import TourneyRunner
    runner = TourneyRunner("meta-llama/Llama-3.1-8B-Instruct", hub_org="my-org")
    winner = runner.run()
"""

from __future__ import annotations

import gc
import json
import math
import os
import shutil
import tempfile
import time
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any, Callable

from obliteratus.credential_sources import resolve_first
from obliteratus.tourney_contracts import parse_checkpoint_document

# ---------------------------------------------------------------------------
# All tournament-eligible methods.
#
# Excluded:
#   - 'optimized' — runs its own Bayesian optimization, far too slow per round
#   - 'nuclear'   — collapsed in essentially every telemetry run (n=7,545)
#   - 'basic'     — fast but quality is unusable across architectures
# ---------------------------------------------------------------------------

TOURNEY_METHODS = [
    "advanced",
    "aggressive",
    "spectral_cascade",
    "informed",
    "surgical",
    "inverted",
    "failspy",
    "gabliteration",
    "heretic",
    "rdo",
]

# ---------------------------------------------------------------------------
# Scoring
# ---------------------------------------------------------------------------


def composite_score(metrics: dict[str, Any]) -> float:
    """Score an abliteration run on [0, 1].  Higher is better.

    Weights:
        35%  refusal removal   — the whole point
        25%  coherence         — model must still be useful
        20%  KL divergence     — minimal capability damage
        10%  perplexity        — fluency preservation
         5%  spectral diagnostic — residual-signal assessment
         5%  degenerate penalty — penalize broken output
    """
    rr = metrics.get("refusal_rate")
    co = metrics.get("coherence")
    kl = metrics.get("kl_divergence")
    pp = metrics.get("perplexity")
    spec = metrics.get("spectral_certification")
    degen = metrics.get("degenerate_count", 0) or 0

    refusal_score = (1.0 - rr) if rr is not None else 0.0
    coherence_score = co if co is not None else 0.0
    kl_score = 1.0 / (1.0 + kl) if kl is not None else 0.5
    ppl_score = 1.0 / (1.0 + pp / 100.0) if pp is not None else 0.5

    # Spectral certification: GREEN=1.0, YELLOW=0.5, RED=0.0, None=0.5 (neutral)
    if spec == "GREEN":
        spec_score = 1.0
    elif spec == "YELLOW":
        spec_score = 0.5
    elif spec == "RED":
        spec_score = 0.0
    else:
        spec_score = 0.5  # not measured → neutral

    # Degenerate penalty: any broken outputs reduce score
    degen_score = 1.0 / (1.0 + degen) if degen > 0 else 1.0

    return (
        refusal_score * 0.35
        + coherence_score * 0.25
        + kl_score * 0.20
        + ppl_score * 0.10
        + spec_score * 0.05
        + degen_score * 0.05
    )


# ---------------------------------------------------------------------------
# Data classes
# ---------------------------------------------------------------------------


@dataclass
class Contender:
    """A single method's result in the tournament."""

    method: str
    score: float = 0.0
    metrics: dict[str, Any] = field(default_factory=dict)
    output_dir: str = ""
    time_s: float = 0.0
    error: str | None = None
    round_eliminated: int = 0  # 0 = still alive / winner
    direction_method: str = ""  # which direction extraction was used
    spectral_cert: str = ""  # GREEN/YELLOW/RED/INCONCLUSIVE/""


@dataclass
class TourneyRound:
    """One round of the tournament."""

    round_num: int
    name: str
    contenders: list[Contender] = field(default_factory=list)
    prompt_volume: int = 0
    advanced_to: list[str] = field(default_factory=list)
    eliminated: list[str] = field(default_factory=list)


@dataclass
class TourneyResult:
    """Full tournament results."""

    model: str
    winner: Contender | None = None
    rounds: list[TourneyRound] = field(default_factory=list)
    total_time_s: float = 0.0
    hub_repo: str | None = None
    timestamp: str = ""

    def to_dict(self) -> dict:
        return {
            "model": self.model,
            "winner": {
                "method": self.winner.method,
                "score": self.winner.score,
                "metrics": self.winner.metrics,
                "time_s": self.winner.time_s,
            } if self.winner else None,
            "rounds": [
                {
                    "round": r.round_num,
                    "name": r.name,
                    "prompt_volume": r.prompt_volume,
                    "contenders": [
                        {
                            "method": c.method,
                            "score": c.score,
                            "metrics": c.metrics,
                            "time_s": c.time_s,
                            "error": c.error,
                            "direction_method": c.direction_method,
                            "spectral_cert": c.spectral_cert,
                        }
                        for c in sorted(r.contenders, key=lambda x: x.score, reverse=True)
                    ],
                    "advanced": r.advanced_to,
                    "eliminated": r.eliminated,
                }
                for r in self.rounds
            ],
            "total_time_s": self.total_time_s,
            "hub_repo": self.hub_repo,
            "timestamp": self.timestamp,
        }


CHECKPOINT_FILENAME = "tourney_checkpoint.json"


def _save_checkpoint(
    output_dir: Path,
    result: TourneyResult,
    current_round_num: int,
    current_round_name: str,
    current_round_volume: int,
    current_round_advance: int,
    current_round_verify: int,
    completed_methods: list[Contender],
    remaining_methods: list[str],
    alive: list[str],
    model_name: str,
    dataset_key: str,
    quantization: str | None,
    methods: list[str],
) -> Path:
    """Save tournament progress so it can be resumed after quota exhaustion."""
    checkpoint = {
        "version": 1,
        "model": model_name,
        "dataset_key": dataset_key,
        "quantization": quantization,
        "methods": methods,
        "alive": alive,
        "completed_rounds": [
            {
                "round_num": r.round_num,
                "name": r.name,
                "prompt_volume": r.prompt_volume,
                "advanced_to": r.advanced_to,
                "eliminated": r.eliminated,
                "contenders": [
                    {
                        "method": c.method,
                        "score": c.score,
                        "metrics": c.metrics,
                        "output_dir": c.output_dir,
                        "time_s": c.time_s,
                        "error": c.error,
                        "round_eliminated": c.round_eliminated,
                        "direction_method": c.direction_method,
                        "spectral_cert": c.spectral_cert,
                    }
                    for c in r.contenders
                ],
            }
            for r in result.rounds
        ],
        "interrupted_round": {
            "round_num": current_round_num,
            "name": current_round_name,
            "prompt_volume": current_round_volume,
            "advance_count": current_round_advance,
            "verify_sample_size": current_round_verify,
            "completed_methods": [
                {
                    "method": c.method,
                    "score": c.score,
                    "metrics": c.metrics,
                    "output_dir": c.output_dir,
                    "time_s": c.time_s,
                    "error": c.error,
                    "round_eliminated": c.round_eliminated,
                    "direction_method": c.direction_method,
                    "spectral_cert": c.spectral_cert,
                }
                for c in completed_methods
            ],
            "remaining_methods": remaining_methods,
        },
        "timestamp": datetime.now().isoformat(),
    }
    path = output_dir / CHECKPOINT_FILENAME
    path.write_text(json.dumps(checkpoint, indent=2))
    return path


def _load_checkpoint(output_dir: Path) -> dict | None:
    """Load a tournament checkpoint if one exists. Returns None if absent or corrupt."""
    path = output_dir / CHECKPOINT_FILENAME
    if not path.exists():
        return None
    try:
        return parse_checkpoint_document(path.read_text())
    except (OSError, UnicodeError):
        return None


def _checkpoint_matches(
    checkpoint: dict,
    model_name: str,
    dataset_key: str,
    quantization: str | None,
) -> bool:
    """Check if a checkpoint is for the same model/dataset/quantization config."""
    return (
        checkpoint.get("model") == model_name
        and checkpoint.get("dataset_key") == dataset_key
        and checkpoint.get("quantization") == quantization
    )


def _restore_rounds(checkpoint: dict) -> tuple[TourneyResult, list[Contender], list[str], dict]:
    """Restore completed rounds and interrupted round state from checkpoint.

    Returns:
        (result_with_completed_rounds, partial_contenders, remaining_methods, interrupted_round_spec)
    """
    result = TourneyResult(
        model=checkpoint["model"],
        timestamp=checkpoint.get("timestamp", ""),
    )

    for rnd_data in checkpoint.get("completed_rounds", []):
        rnd = TourneyRound(
            round_num=rnd_data["round_num"],
            name=rnd_data["name"],
            prompt_volume=rnd_data.get("prompt_volume", 0),
            advanced_to=rnd_data.get("advanced_to", []),
            eliminated=rnd_data.get("eliminated", []),
        )
        for c_data in rnd_data.get("contenders", []):
            rnd.contenders.append(Contender(
                method=c_data["method"],
                score=c_data.get("score", 0.0),
                metrics=c_data.get("metrics", {}),
                output_dir=c_data.get("output_dir", ""),
                time_s=c_data.get("time_s", 0.0),
                error=c_data.get("error"),
                round_eliminated=c_data.get("round_eliminated", 0),
                direction_method=c_data.get("direction_method", ""),
                spectral_cert=c_data.get("spectral_cert", ""),
            ))
        result.rounds.append(rnd)

    ir = checkpoint.get("interrupted_round", {})
    partial_contenders = []
    for c_data in ir.get("completed_methods", []):
        partial_contenders.append(Contender(
            method=c_data["method"],
            score=c_data.get("score", 0.0),
            metrics=c_data.get("metrics", {}),
            output_dir=c_data.get("output_dir", ""),
            time_s=c_data.get("time_s", 0.0),
            error=c_data.get("error"),
            round_eliminated=c_data.get("round_eliminated", 0),
            direction_method=c_data.get("direction_method", ""),
            spectral_cert=c_data.get("spectral_cert", ""),
        ))

    remaining = ir.get("remaining_methods", [])

    return result, partial_contenders, remaining, ir


# ---------------------------------------------------------------------------
# Bracket renderer
# ---------------------------------------------------------------------------


def render_bracket(result: TourneyResult) -> str:
    """Render the tournament bracket as a markdown string."""
    lines = []
    lines.append(f"# OBLITERATUS TOURNEY — {result.model}")
    lines.append("")
    lines.append(f"**Winner: `{result.winner.method}`** "
                 f"(score: {result.winner.score:.4f})" if result.winner else "**No winner**")
    lines.append(f"Total time: {result.total_time_s / 60:.1f} minutes")
    if result.hub_repo:
        lines.append(f"Pushed to: [{result.hub_repo}](https://huggingface.co/{result.hub_repo})")
    lines.append("")

    for rnd in result.rounds:
        lines.append(f"## Round {rnd.round_num}: {rnd.name}")
        lines.append(f"*{len(rnd.contenders)} contenders, {rnd.prompt_volume} prompt pairs*")
        lines.append("")
        lines.append("| Rank | Method | Dir | Score | Refusal | Coherence | KL Div | PPL | Cert | Time |")
        lines.append("|------|--------|-----|-------|---------|-----------|--------|-----|------|------|")

        sorted_contenders = sorted(rnd.contenders, key=lambda x: x.score, reverse=True)
        for i, c in enumerate(sorted_contenders, 1):
            if c.error:
                lines.append(
                    f"| {i} | {c.method} | — | ERROR | — | — | — | — | — | {c.time_s:.0f}s |"
                )
                continue
            m = c.metrics
            # Only annotate elimination for non-final rounds
            if c.method in rnd.advanced_to:
                marker = ""
            elif rnd.round_num < len(result.rounds):
                marker = " *out*"
            else:
                marker = ""
            rr = f"{m.get('refusal_rate', 0):.1%}" if m.get('refusal_rate') is not None else "—"
            co = f"{m.get('coherence', 0):.3f}" if m.get('coherence') is not None else "—"
            kl_val = m.get('kl_divergence')
            kl_str = f"{kl_val:.4f}" if kl_val is not None else "—"
            pp = f"{m.get('perplexity', 0):.1f}" if m.get('perplexity') is not None else "—"
            dir_m = c.direction_method or m.get("direction_method", "—")
            cert = c.spectral_cert or "—"
            lines.append(
                f"| {i} | **{c.method}**{marker} | {dir_m} | {c.score:.4f} "
                f"| {rr} | {co} | {kl_str} | {pp} | {cert} | {c.time_s:.0f}s |"
            )
        lines.append("")

    return "\n".join(lines)


def render_bracket_html(result: TourneyResult) -> str:
    """Render the tournament bracket as a styled HTML bracket visualization."""
    import html as html_mod

    model_short = result.model.split("/")[-1] if "/" in result.model else result.model

    # ── CSS ──────────────────────────────────────────────────────────────
    css = """
    <style>
    .tourney-wrap {
        font-family: 'Segoe UI', system-ui, -apple-system, sans-serif;
        color: #e0e0e0;
        max-width: 100%;
        overflow-x: auto;
    }
    .tourney-header {
        text-align: center;
        padding: 18px 20px;
        margin-bottom: 20px;
        background: linear-gradient(135deg, #1a1a2e 0%, #16213e 100%);
        border-radius: 12px;
        border: 1px solid #333;
    }
    .tourney-header h2 {
        margin: 0 0 4px 0;
        font-size: 1.4em;
        color: #fff;
        letter-spacing: 1px;
    }
    .tourney-header .model-name {
        font-size: 0.85em;
        color: #8892b0;
        font-family: 'Courier New', monospace;
    }
    .tourney-header .champion-box {
        margin-top: 14px;
        padding: 12px 18px;
        background: linear-gradient(135deg, #2d1f00 0%, #3d2a00 100%);
        border: 1px solid #f0c040;
        border-radius: 8px;
        display: inline-block;
    }
    .tourney-header .champion-box .trophy { font-size: 1.4em; }
    .tourney-header .champion-box .champ-name {
        font-size: 1.15em;
        font-weight: 700;
        color: #f0c040;
        font-family: 'Courier New', monospace;
    }
    .tourney-header .champion-box .champ-score {
        font-size: 0.85em;
        color: #cca030;
        margin-top: 2px;
    }
    .tourney-header .no-winner {
        margin-top: 14px;
        padding: 10px 16px;
        background: #2a1a1a;
        border: 1px solid #cc4444;
        border-radius: 8px;
        display: inline-block;
        color: #ff6b6b;
        font-weight: 600;
    }
    .tourney-header .time-info {
        font-size: 0.78em;
        color: #666;
        margin-top: 8px;
    }

    /* ── Bracket flow ── */
    .bracket-flow {
        display: flex;
        gap: 12px;
        align-items: stretch;
        padding: 4px 0;
    }
    .round-col {
        flex: 1;
        min-width: 200px;
        max-width: 340px;
    }
    .round-title {
        text-align: center;
        font-size: 0.82em;
        font-weight: 700;
        text-transform: uppercase;
        letter-spacing: 1.5px;
        color: #8892b0;
        padding: 6px 0 8px 0;
        border-bottom: 2px solid #333;
        margin-bottom: 8px;
    }
    .round-subtitle {
        text-align: center;
        font-size: 0.7em;
        color: #555;
        margin-top: 2px;
    }

    /* ── Method cards ── */
    .method-card {
        padding: 8px 10px;
        margin: 4px 0;
        border-radius: 6px;
        border-left: 3px solid #444;
        background: #1c1c2e;
        transition: all 0.2s;
    }
    .method-card.advanced {
        border-left-color: #4ecca3;
        background: #1a2e28;
    }
    .method-card.champion {
        border-left-color: #f0c040;
        background: #2d2a1a;
        box-shadow: 0 0 8px rgba(240, 192, 64, 0.15);
    }
    .method-card.eliminated {
        border-left-color: #cc4444;
        background: #1e1a1a;
        opacity: 0.7;
    }
    .method-card.errored {
        border-left-color: #ff4444;
        background: #2a1a1a;
        opacity: 0.6;
    }
    .card-top {
        display: flex;
        justify-content: space-between;
        align-items: center;
    }
    .card-rank {
        font-size: 0.7em;
        color: #666;
        font-weight: 700;
        min-width: 18px;
    }
    .card-name {
        font-weight: 600;
        font-size: 0.88em;
        font-family: 'Courier New', monospace;
        flex: 1;
        margin: 0 6px;
    }
    .card-score {
        font-weight: 700;
        font-size: 0.88em;
        font-family: 'Courier New', monospace;
    }
    .card-score.good { color: #4ecca3; }
    .card-score.mid { color: #f0c040; }
    .card-score.bad { color: #cc4444; }
    .card-metrics {
        display: flex;
        gap: 8px;
        margin-top: 4px;
        flex-wrap: wrap;
    }
    .metric {
        font-size: 0.68em;
        color: #777;
    }
    .metric .val {
        color: #aaa;
        font-family: 'Courier New', monospace;
    }
    .card-badge {
        font-size: 0.65em;
        font-weight: 700;
        padding: 1px 5px;
        border-radius: 3px;
        text-transform: uppercase;
        letter-spacing: 0.5px;
    }
    .badge-adv { background: #1a3a2e; color: #4ecca3; }
    .badge-out { background: #2a1a1a; color: #cc6666; }
    .badge-champ { background: #3d2a00; color: #f0c040; }
    .badge-err { background: #2a1a1a; color: #ff6666; }

    /* ── Arrow column ── */
    .arrow-col {
        display: flex;
        flex-direction: column;
        justify-content: center;
        align-items: center;
        width: 30px;
        min-width: 30px;
        color: #444;
        font-size: 1.2em;
    }
    </style>
    """

    # ── Header ───────────────────────────────────────────────────────────
    header_parts = [
        '<div class="tourney-header">',
        '<h2>OBLITERATUS TOURNEY</h2>',
        f'<div class="model-name">{html_mod.escape(model_short)}</div>',
    ]

    if result.winner and not result.winner.error:
        w = result.winner
        m = w.metrics or {}
        rr = f"{m.get('refusal_rate', 0):.1%}" if m.get("refusal_rate") is not None else "—"
        co = f"{m.get('coherence', 0):.3f}" if m.get("coherence") is not None else "—"
        header_parts.append('<div class="champion-box">')
        header_parts.append(
            f'<span class="trophy">&#x1F3C6;</span> '
            f'<span class="champ-name">{html_mod.escape(w.method)}</span>'
        )
        dir_m = w.direction_method or "—"
        cert = w.spectral_cert or "—"
        header_parts.append(
            f'<div class="champ-score">'
            f'Score: {w.score:.4f} &nbsp;|&nbsp; Refusal: {rr} &nbsp;|&nbsp; '
            f'Coherence: {co} &nbsp;|&nbsp; Dir: {html_mod.escape(dir_m)} &nbsp;|&nbsp; Cert: {html_mod.escape(cert)}'
            f'</div>'
        )
        header_parts.append("</div>")
    else:
        header_parts.append('<div class="no-winner">No winner determined</div>')

    if result.total_time_s:
        header_parts.append(
            f'<div class="time-info">{result.total_time_s / 60:.1f} min total</div>'
        )
    header_parts.append("</div>")

    # ── Bracket columns ──────────────────────────────────────────────────
    bracket_parts = ['<div class="bracket-flow">']
    n_rounds = len(result.rounds)

    for ri, rnd in enumerate(result.rounds):
        if ri > 0:
            bracket_parts.append('<div class="arrow-col">&#x25B6;</div>')

        bracket_parts.append('<div class="round-col">')
        bracket_parts.append(
            f'<div class="round-title">{html_mod.escape(rnd.name)}'
            f'<div class="round-subtitle">{rnd.prompt_volume} pairs</div></div>'
        )

        sorted_c = sorted(rnd.contenders, key=lambda c: c.score, reverse=True)
        is_final = ri == n_rounds - 1

        for rank, c in enumerate(sorted_c, 1):
            if c.error:
                css_cls = "errored"
                badge = '<span class="card-badge badge-err">ERR</span>'
            elif is_final and rank == 1 and result.winner and not result.winner.error:
                css_cls = "champion"
                badge = '<span class="card-badge badge-champ">&#x2605; CHAMP</span>'
            elif c.method in (rnd.advanced_to or []):
                css_cls = "advanced"
                badge = '<span class="card-badge badge-adv">ADV</span>'
            else:
                css_cls = "eliminated"
                badge = '<span class="card-badge badge-out">OUT</span>'

            # Score color
            if c.error:
                score_html = '<span class="card-score bad">ERR</span>'
            elif c.score >= 0.7:
                score_html = f'<span class="card-score good">{c.score:.4f}</span>'
            elif c.score >= 0.4:
                score_html = f'<span class="card-score mid">{c.score:.4f}</span>'
            else:
                score_html = f'<span class="card-score bad">{c.score:.4f}</span>'

            # Compact metrics
            m = c.metrics or {}
            metric_spans = []
            if not c.error:
                dm = c.direction_method or m.get("direction_method", "")
                if dm:
                    metric_spans.append(
                        f'<span class="metric">dir <span class="val">{html_mod.escape(dm)}</span></span>'
                    )
                rr = m.get("refusal_rate")
                if rr is not None:
                    metric_spans.append(
                        f'<span class="metric">ref <span class="val">{rr:.0%}</span></span>'
                    )
                co = m.get("coherence")
                if co is not None:
                    metric_spans.append(
                        f'<span class="metric">coh <span class="val">{co:.3f}</span></span>'
                    )
                sc = c.spectral_cert or m.get("spectral_certification", "")
                if sc:
                    cert_color = {"GREEN": "#4ecca3", "YELLOW": "#f0c040", "RED": "#cc4444"}.get(sc, "#777")
                    metric_spans.append(
                        f'<span class="metric">cert <span class="val" style="color:{cert_color}">{html_mod.escape(sc)}</span></span>'
                    )
                kl = m.get("kl_divergence")
                if kl is not None:
                    metric_spans.append(
                        f'<span class="metric">kl <span class="val">{kl:.4f}</span></span>'
                    )
                pp = m.get("perplexity")
                if pp is not None:
                    metric_spans.append(
                        f'<span class="metric">ppl <span class="val">{pp:.1f}</span></span>'
                    )
            metrics_html = "".join(metric_spans)

            bracket_parts.append(f'<div class="method-card {css_cls}">')
            bracket_parts.append(
                f'<div class="card-top">'
                f'<span class="card-rank">#{rank}</span>'
                f'<span class="card-name">{html_mod.escape(c.method)}</span>'
                f'{score_html}'
                f'{badge}'
                f'</div>'
            )
            if metrics_html:
                bracket_parts.append(f'<div class="card-metrics">{metrics_html}</div>')
            bracket_parts.append("</div>")

        bracket_parts.append("</div>")

    bracket_parts.append("</div>")

    return css + '<div class="tourney-wrap">' + "\n".join(header_parts + bracket_parts) + "</div>"


def generate_model_card(result: TourneyResult) -> str:
    """Generate a HuggingFace model card for the tournament winner."""
    w = result.winner
    if not w:
        return ""

    short_model = result.model.split("/")[-1] if "/" in result.model else result.model
    bracket = render_bracket(result)

    return f"""---
language: en
tags:
  - obliteratus
  - abliteration
  - uncensored
  - tourney
base_model: {result.model}
---

# {short_model} — Obliterated (Tourney Winner)

This model was abliterated using the **`{w.method}`** method, selected by an
automated [OBLITERATUS](https://github.com/elder-plinius/OBLITERATUS) tournament
that pitted **{len(TOURNEY_METHODS)} abliteration techniques** against each other
in elimination rounds.

## Winning Method: `{w.method}`

| Metric | Value |
|--------|-------|
| Composite Score | **{w.score:.4f}** |
| Direction Method | {w.direction_method or 'N/A'} |
| Refusal Rate | {f'{w.metrics["refusal_rate"]:.1%}' if w.metrics.get('refusal_rate') is not None else 'N/A'} |
| Coherence | {f'{w.metrics["coherence"]:.3f}' if w.metrics.get('coherence') is not None else 'N/A'} |
| KL Divergence | {f'{w.metrics["kl_divergence"]:.4f}' if w.metrics.get('kl_divergence') is not None else 'N/A'} |
| Perplexity | {f'{w.metrics["perplexity"]:.1f}' if w.metrics.get('perplexity') is not None else 'N/A'} |
| Spectral Cert | {w.spectral_cert or 'N/A'} |

## How to Use

```python
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("{result.hub_repo or 'this-repo'}")
tokenizer = AutoTokenizer.from_pretrained("{result.hub_repo or 'this-repo'}")
```

## Full Tournament Bracket

{bracket}

---

*Generated by [OBLITERATUS](https://github.com/elder-plinius/OBLITERATUS) tourney on {result.timestamp}*
"""


def _noop_log(msg: str) -> None:
    """Picklable no-op log callback (lambdas can't be pickled by ZeroGPU)."""
    pass


def _noop_round(r: TourneyRound) -> None:
    """Picklable no-op round callback."""
    pass


class _MethodLogger:
    """Picklable per-method log adapter that prefixes messages.

    ZeroGPU pickles bound methods (and their ``self``) when shipping work to
    the GPU worker process.  Plain lambdas like
    ``lambda msg: self.log(f"  [{method}] {msg}")`` can't survive that, so
    this small class replaces them.
    """

    def __init__(self, parent_log: Callable[[str], None], method: str):
        self._parent = parent_log
        self._method = method

    def __call__(self, msg: str):
        self._parent(f"  [{self._method}] {msg}")


# ---------------------------------------------------------------------------
# Core runner
# ---------------------------------------------------------------------------


class TourneyRunner:
    """Run a March Madness-style tournament across all abliteration methods.

    Round 1 — Qualifiers:  All methods, reduced prompts.  Bottom half eliminated.
    Round 2 — Semifinals:  Survivors, full prompts.  Bottom half eliminated.
    Round 3 — Finals:      Top contenders, maximum prompts + extended verify.
    Winner  — Auto-pushed to HuggingFace Hub (if hub_org is set).
    """

    def __init__(
        self,
        model_name: str,
        hub_org: str | None = None,
        hub_repo: str | None = None,
        device: str = "auto",
        dtype: str = "float16",
        dataset_key: str = "builtin",
        quantization: str | None = None,
        methods: list[str] | None = None,
        output_dir: str = os.path.join(tempfile.gettempdir(), "obliteratus_tourney"),
        on_log: Callable[[str], None] | None = None,
        on_round: Callable[[TourneyRound], None] | None = None,
        resume: bool = False,
    ):
        self.model_name = model_name
        self.hub_org = hub_org
        self.hub_repo = hub_repo
        self.device = device
        self.dtype = dtype
        self.dataset_key = dataset_key
        self.quantization = quantization
        self.methods = methods or list(TOURNEY_METHODS)
        self.output_dir = Path(output_dir)
        self.resume = resume

        # When resuming, preserve the output directory (it contains
        # checkpoints and model saves from completed methods).
        if not resume:
            if self.output_dir.exists():
                shutil.rmtree(self.output_dir, ignore_errors=True)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self._on_log = on_log or _noop_log
        self._on_round = on_round or _noop_round

    def log(self, msg: str):
        self._on_log(msg)

    def _load_prompts(self, volume: int) -> tuple[list[str], list[str]]:
        from obliteratus.prompts import load_dataset_source
        harmful, harmless = load_dataset_source(self.dataset_key)
        n = min(volume, len(harmful), len(harmless))
        return harmful[:n], harmless[:n]

    def _run_method(
        self,
        method: str,
        harmful: list[str],
        harmless: list[str],
        save_dir: str,
        verify_sample_size: int = 30,
    ) -> Contender:
        """Run a single abliteration method and return its Contender result."""
        t0 = time.time()
        contender = Contender(method=method)

        try:
            # Use informed pipeline for 'informed' method
            method_log = _MethodLogger(self._on_log, method)

            if method == "informed":
                from obliteratus.informed_pipeline import InformedAbliterationPipeline
                pipeline = InformedAbliterationPipeline(
                    model_name=self.model_name,
                    output_dir=save_dir,
                    device=self.device,
                    dtype=self.dtype,
                    quantization=self.quantization,
                    trust_remote_code=True,
                    harmful_prompts=harmful,
                    harmless_prompts=harmless,
                    on_log=method_log,
                )
                pipeline.run_informed()
            else:
                from obliteratus.abliterate import AbliterationPipeline
                pipeline = AbliterationPipeline(
                    model_name=self.model_name,
                    output_dir=save_dir,
                    device=self.device,
                    dtype=self.dtype,
                    method=method,
                    quantization=self.quantization,
                    trust_remote_code=True,
                    harmful_prompts=harmful,
                    harmless_prompts=harmless,
                    verify_sample_size=verify_sample_size,
                    on_log=method_log,
                )
                pipeline.run()

            contender.metrics = dict(pipeline._quality_metrics)
            contender.score = composite_score(contender.metrics)
            contender.output_dir = save_dir
            contender.direction_method = getattr(pipeline, "direction_method", "")
            contender.spectral_cert = contender.metrics.get("spectral_certification", "") or ""

            # Free pipeline to reclaim GPU
            del pipeline
        except Exception as e:
            # Re-raise GPU quota / expired-token errors so the
            # tournament aborts immediately rather than letting every
            # remaining method fail for the same reason.
            if self._is_quota_error(e):
                raise
            import traceback
            contender.error = f"{type(e).__name__}: {e}"
            contender.score = -1.0  # errors sort to bottom
            self.log(f"  [{method}] ERROR: {contender.error}")
            self.log(f"  [{method}] TRACEBACK:\n{traceback.format_exc()}")
        finally:
            # Always clean up GPU between methods — including when
            # re-raising quota errors, to avoid leaking the pipeline.
            gc.collect()
            try:
                from obliteratus import device as dev
                dev.empty_cache()
            except Exception:
                pass

        contender.time_s = time.time() - t0
        return contender

    def _run_round(
        self,
        round_num: int,
        name: str,
        methods: list[str],
        prompt_volume: int,
        advance_count: int,
        verify_sample_size: int = 30,
    ) -> TourneyRound:
        """Execute one round of the tournament."""
        self.log("")
        self.log("=" * 60)
        self.log(f"ROUND {round_num}: {name}")
        self.log(f"  {len(methods)} contenders | {prompt_volume} prompt pairs | "
                 f"top {advance_count} advance")
        self.log("=" * 60)

        harmful, harmless = self._load_prompts(prompt_volume)

        rnd = TourneyRound(
            round_num=round_num,
            name=name,
            prompt_volume=prompt_volume,
        )

        for i, method in enumerate(methods, 1):
            self.log(f"\n[{i}/{len(methods)}] Running: {method}")
            save_dir = str(self.output_dir / f"r{round_num}_{method}")
            contender = self._run_method(
                method, harmful, harmless, save_dir, verify_sample_size,
            )
            rnd.contenders.append(contender)
            self.log(
                f"  {method}: score={contender.score:.4f} "
                f"(refusal={contender.metrics.get('refusal_rate', '?')}, "
                f"coherence={contender.metrics.get('coherence', '?')}) "
                f"[{contender.time_s:.0f}s]"
            )

            # Free checkpoint for non-finalists as we go (save disk)
            # We'll keep them until we know who advances

        # Rank by score
        ranked = sorted(rnd.contenders, key=lambda c: c.score, reverse=True)
        rnd.advanced_to = [c.method for c in ranked[:advance_count]]
        rnd.eliminated = [c.method for c in ranked[advance_count:]]

        # Mark eliminated
        for c in ranked[advance_count:]:
            c.round_eliminated = round_num

        self.log(f"\n{'-' * 40}")
        self.log(f"Round {round_num} results:")
        for i, c in enumerate(ranked, 1):
            status = "ADVANCE" if c.method in rnd.advanced_to else "OUT"
            self.log(f"  {i}. {c.method}: {c.score:.4f} [{status}]")

        # Clean up eliminated checkpoints to free disk
        for c in ranked[advance_count:]:
            if c.output_dir and Path(c.output_dir).exists():
                shutil.rmtree(c.output_dir, ignore_errors=True)

        self._on_round(rnd)
        return rnd

    def run(self) -> TourneyResult:
        """Execute the full tournament. Returns TourneyResult with winner."""
        t_start = time.time()
        result = TourneyResult(
            model=self.model_name,
            timestamp=datetime.now().isoformat(),
        )

        n_methods = len(self.methods)
        self.log("OBLITERATUS TOURNEY")
        self.log(f"Model: {self.model_name}")
        self.log(f"Contenders: {n_methods} methods")
        self.log(f"Dataset: {self.dataset_key}")

        # Pre-flight disk space check
        try:
            disk = shutil.disk_usage(self.output_dir)
            free_gb = disk.free / 1e9
            self.log(f"Disk space: {free_gb:.1f} GB free on {self.output_dir}")
            if free_gb < 5.0:
                self.log(
                    f"WARNING: Low disk space ({free_gb:.1f} GB free). "
                    f"Tournament may fail saving checkpoints."
                )
        except Exception:
            pass

        # ── Round 1: Qualifiers — all methods, reduced prompts ────────
        r1_advance = max(2, math.ceil(n_methods / 2))
        r1 = self._run_round(
            round_num=1,
            name="Qualifiers",
            methods=self.methods,
            prompt_volume=64,       # fast qualifier round
            advance_count=r1_advance,
            verify_sample_size=20,
        )
        result.rounds.append(r1)
        alive = list(r1.advanced_to)

        if len(alive) <= 1:
            # Only 1 survivor — they win
            pass
        else:
            # ── Round 2: Semifinals — survivors, full prompts ─────────
            r2_advance = max(2, math.ceil(len(alive) / 2))
            r2 = self._run_round(
                round_num=2,
                name="Semifinals",
                methods=alive,
                prompt_volume=128,
                advance_count=r2_advance,
                verify_sample_size=30,
            )
            result.rounds.append(r2)
            alive = list(r2.advanced_to)

            if len(alive) > 2:
                # ── Round 3: Finals — top contenders, max prompts ─────
                r3 = self._run_round(
                    round_num=3,
                    name="Finals",
                    methods=alive,
                    prompt_volume=256,
                    advance_count=1,
                    verify_sample_size=50,
                )
                result.rounds.append(r3)
                alive = list(r3.advanced_to)
            elif len(alive) == 2:
                # Head-to-head final
                r3 = self._run_round(
                    round_num=3,
                    name="Championship",
                    methods=alive,
                    prompt_volume=256,
                    advance_count=1,
                    verify_sample_size=50,
                )
                result.rounds.append(r3)
                alive = list(r3.advanced_to)

        # ── Determine winner ──────────────────────────────────────────
        last_round = result.rounds[-1]
        ranked = sorted(last_round.contenders, key=lambda c: c.score, reverse=True)
        # Only crown a winner if they completed without error
        winner = ranked[0] if ranked and not ranked[0].error else None
        result.winner = winner
        result.total_time_s = time.time() - t_start

        # Clean up non-winner finalist dirs to free disk
        for c in ranked[1:]:
            if c.output_dir and Path(c.output_dir).exists():
                shutil.rmtree(c.output_dir, ignore_errors=True)

        self.log("")
        self.log("=" * 60)
        if winner:
            self.log(f"CHAMPION: {winner.method} (score: {winner.score:.4f})")
        else:
            n_errors = sum(1 for c in ranked if c.error)
            self.log(f"NO WINNER — {n_errors}/{len(ranked)} methods errored")
        self.log(f"Total tournament time: {result.total_time_s / 60:.1f} minutes")
        self.log("=" * 60)

        # ── Save tournament results ───────────────────────────────────
        results_path = self.output_dir / "tourney_results.json"
        results_path.write_text(json.dumps(result.to_dict(), indent=2))
        self.log(f"Results saved to {results_path}")

        bracket_path = self.output_dir / "tourney_bracket.md"
        bracket_path.write_text(render_bracket(result))
        self.log(f"Bracket saved to {bracket_path}")

        # ── Push winner to HuggingFace Hub ────────────────────────────
        if winner and winner.output_dir and (self.hub_org or self.hub_repo):
            self._push_winner(result)

        return result

    @staticmethod
    def _is_quota_error(exc: BaseException) -> bool:
        msg = str(exc).lower()
        if "exceeded" in msg and "gpu quota" in msg:
            return True
        if "expired" in msg and "zerogpu" in msg:
            return True
        return False

    def _run_one_method(self, method, harmful, harmless, save_dir, verify_sz, gpu_wrapper):
        """Run a single method, optionally inside a gpu_wrapper."""
        if gpu_wrapper is not None:
            return gpu_wrapper(
                self._run_method, method, harmful, harmless,
                save_dir, verify_sz,
            )
        return self._run_method(
            method, harmful, harmless, save_dir, verify_sz,
        )

    def run_iter(self, gpu_wrapper=None):
        """Generator version of run() — yields (status, result_so_far) after each method.

        Supports automatic resume: if ``self.resume`` is True and a valid
        checkpoint exists from a previous quota-interrupted run with the
        same model/dataset/quantization, completed rounds and methods are
        restored and execution continues from the interruption point.

        When a GPU quota error occurs, a checkpoint is saved to disk and
        the exception is re-raised.  The caller can catch it and inform
        the user that clicking **Run** again will resume automatically.

        Args:
            gpu_wrapper: Optional callable ``gpu_wrapper(fn, *args, **kw)``
                that executes *fn* inside a GPU context.  On ZeroGPU Spaces
                this should be a ``@spaces.GPU``-decorated function so each
                method gets its own GPU allocation (up to 5 min each).

        Yields:
            (status_msg: str, result: TourneyResult | None)
        """

        t_start = time.time()
        resuming = False
        checkpoint = None
        partial_contenders: list[Contender] = []
        resume_remaining: list[str] = []
        resume_round_spec: dict = {}

        # ── Try to resume from checkpoint ────────────────────────────
        if self.resume:
            checkpoint = _load_checkpoint(self.output_dir)
            if checkpoint and _checkpoint_matches(
                checkpoint, self.model_name, self.dataset_key, self.quantization
            ):
                resuming = True
                result, partial_contenders, resume_remaining, resume_round_spec = (
                    _restore_rounds(checkpoint)
                )
                n_completed_rounds = len(result.rounds)
                n_completed_methods = len(partial_contenders)
                self.log("OBLITERATUS TOURNEY — RESUMING")
                self.log(f"Restored {n_completed_rounds} completed round(s), "
                         f"{n_completed_methods} method(s) from interrupted round")
                yield (
                    f"**Resuming tournament** — {n_completed_rounds} round(s) "
                    f"and {n_completed_methods} method(s) restored from checkpoint.",
                    result,
                )

                # Determine alive list from checkpoint
                alive = list(checkpoint.get("alive", self.methods))

                # Remove the checkpoint file now that we've loaded it
                ckpt_path = self.output_dir / CHECKPOINT_FILENAME
                if ckpt_path.exists():
                    ckpt_path.unlink()
            else:
                # Checkpoint doesn't match current config — start fresh
                checkpoint = None

        n_methods = len(self.methods)

        if not resuming:
            result = TourneyResult(
                model=self.model_name,
                timestamp=datetime.now().isoformat(),
            )
            alive = list(self.methods)

            self.log("OBLITERATUS TOURNEY")
            self.log(f"Model: {self.model_name}")
            self.log(f"Contenders: {n_methods} methods")
            self.log(f"Dataset: {self.dataset_key}")

        # Pre-flight disk space check
        try:
            disk = shutil.disk_usage(self.output_dir)
            free_gb = disk.free / 1e9
            self.log(f"Disk space: {free_gb:.1f} GB free on {self.output_dir}")
            if free_gb < 5.0:
                msg = (
                    f"Low disk space: only {free_gb:.1f} GB free. "
                    f"Tournament needs space for multiple model checkpoints. "
                    f"Free up space or use quantization to reduce checkpoint sizes."
                )
                self.log(f"WARNING: {msg}")
                yield (f"**Warning:** {msg}", result)
        except Exception:
            pass

        # Build round schedule
        rounds_schedule: list[tuple] = []

        if resuming and resume_round_spec:
            # We have an interrupted round to finish — schedule it first,
            # then let the dynamic scheduling add subsequent rounds.
            skip_completed_rounds = len(result.rounds)
        else:
            skip_completed_rounds = 0

        # Always build the full schedule starting from round 1.
        # Completed rounds will be skipped below.
        r1_advance = max(2, math.ceil(n_methods / 2))
        rounds_schedule.append((1, "Qualifiers", self.methods, 64, r1_advance, 20))

        for round_spec in rounds_schedule:
            round_num, name, methods, volume, advance_count, verify_sz = round_spec

            # Skip rounds that were already completed in the checkpoint
            if resuming and round_num <= skip_completed_rounds:
                # Re-derive alive and schedule next rounds from completed data
                completed_rnd = result.rounds[round_num - 1]
                alive = list(completed_rnd.advanced_to)
                if round_num == 1 and len(alive) > 1:
                    r2_advance = max(2, math.ceil(len(alive) / 2))
                    rounds_schedule.append((2, "Semifinals", alive, 128, r2_advance, 30))
                elif round_num == 2 and len(alive) > 1:
                    r3_name = "Championship" if len(alive) == 2 else "Finals"
                    rounds_schedule.append((3, r3_name, alive, 256, 1, 50))
                self.log(f"\nSkipping completed Round {round_num}: {name}")
                yield (
                    f"**Round {round_num} ({name}):** already completed (restored from checkpoint)",
                    result,
                )
                continue

            # For the interrupted round, merge checkpoint data
            is_interrupted_round = (
                resuming
                and resume_round_spec
                and round_num == resume_round_spec.get("round_num")
            )

            if is_interrupted_round:
                # Use the interrupted round's parameters
                volume = resume_round_spec.get("prompt_volume", volume)
                advance_count = resume_round_spec.get("advance_count", advance_count)
                verify_sz = resume_round_spec.get("verify_sample_size", verify_sz)
                methods = list(
                    [c.method for c in partial_contenders] + resume_remaining
                )

            self.log("")
            self.log("=" * 60)
            self.log(f"ROUND {round_num}: {name}")
            self.log(f"  {len(methods)} contenders | {volume} prompt pairs | "
                     f"top {advance_count} advance")
            self.log("=" * 60)

            harmful, harmless = self._load_prompts(volume)

            rnd = TourneyRound(
                round_num=round_num,
                name=name,
                prompt_volume=volume,
            )

            # If resuming an interrupted round, restore already-completed
            # contenders and only run the remaining methods.
            methods_to_run = list(methods)
            if is_interrupted_round and partial_contenders:
                for c in partial_contenders:
                    rnd.contenders.append(c)
                    self.log(f"  [restored] {c.method}: score={c.score:.4f}")
                methods_to_run = list(resume_remaining)
                self.log(f"  {len(partial_contenders)} method(s) restored, "
                         f"{len(methods_to_run)} remaining")

            total_in_round = len(rnd.contenders) + len(methods_to_run)

            for i, method in enumerate(methods_to_run, len(rnd.contenders) + 1):
                self.log(f"\n[{i}/{total_in_round}] Running: {method}")
                yield (
                    f"**Round {round_num} ({name}):** running `{method}` [{i}/{total_in_round}]",
                    result,
                )

                save_dir = str(self.output_dir / f"r{round_num}_{method}")

                try:
                    contender = self._run_one_method(
                        method, harmful, harmless, save_dir, verify_sz,
                        gpu_wrapper,
                    )
                except Exception as exc:
                    if self._is_quota_error(exc):
                        # Save checkpoint so the tournament can resume later.
                        # Include the failed method in remaining so it retries.
                        still_remaining = methods_to_run[methods_to_run.index(method):]
                        _save_checkpoint(
                            output_dir=self.output_dir,
                            result=result,
                            current_round_num=round_num,
                            current_round_name=name,
                            current_round_volume=volume,
                            current_round_advance=advance_count,
                            current_round_verify=verify_sz,
                            completed_methods=list(rnd.contenders),
                            remaining_methods=still_remaining,
                            alive=alive,
                            model_name=self.model_name,
                            dataset_key=self.dataset_key,
                            quantization=self.quantization,
                            methods=self.methods,
                        )
                        self.log("\nGPU SESSION INTERRUPTED — checkpoint saved")
                        self.log(f"  Reason: {exc}")
                        self.log(f"  Completed: {len(rnd.contenders)} methods in round {round_num}")
                        self.log(f"  Remaining: {len(still_remaining)} methods")
                        self.log("  Click Run again to resume automatically.")
                    raise

                rnd.contenders.append(contender)
                self.log(
                    f"  {method}: score={contender.score:.4f} "
                    f"(refusal={contender.metrics.get('refusal_rate', '?')}, "
                    f"coherence={contender.metrics.get('coherence', '?')}) "
                    f"[{contender.time_s:.0f}s]"
                )

            # Rank, advance, eliminate
            ranked = sorted(rnd.contenders, key=lambda c: c.score, reverse=True)
            rnd.advanced_to = [c.method for c in ranked[:advance_count]]
            rnd.eliminated = [c.method for c in ranked[advance_count:]]
            for c in ranked[advance_count:]:
                c.round_eliminated = round_num

            self.log(f"\n{'-' * 40}")
            self.log(f"Round {round_num} results:")
            for idx, c in enumerate(ranked, 1):
                status = "ADVANCE" if c.method in rnd.advanced_to else "OUT"
                self.log(f"  {idx}. {c.method}: {c.score:.4f} [{status}]")

            # Clean up eliminated checkpoints
            for c in ranked[advance_count:]:
                if c.output_dir and Path(c.output_dir).exists():
                    shutil.rmtree(c.output_dir, ignore_errors=True)

            self._on_round(rnd)
            result.rounds.append(rnd)
            alive = list(rnd.advanced_to)

            # Schedule next round dynamically
            if round_num == 1 and len(alive) > 1:
                r2_advance = max(2, math.ceil(len(alive) / 2))
                rounds_schedule.append((2, "Semifinals", alive, 128, r2_advance, 30))
            elif round_num == 2 and len(alive) > 1:
                r3_name = "Championship" if len(alive) == 2 else "Finals"
                rounds_schedule.append((3, r3_name, alive, 256, 1, 50))

        # ── Determine winner ──────────────────────────────────────────
        last_round = result.rounds[-1]
        ranked = sorted(last_round.contenders, key=lambda c: c.score, reverse=True)
        winner = ranked[0] if ranked and not ranked[0].error else None
        result.winner = winner
        result.total_time_s = time.time() - t_start

        # Clean up non-winner finalist dirs to free disk
        for c in ranked[1:]:
            if c.output_dir and Path(c.output_dir).exists():
                shutil.rmtree(c.output_dir, ignore_errors=True)

        self.log("")
        self.log("=" * 60)
        if winner:
            self.log(f"CHAMPION: {winner.method} (score: {winner.score:.4f})")
        else:
            n_errors = sum(1 for c in ranked if c.error)
            self.log(f"NO WINNER — {n_errors}/{len(ranked)} methods errored")
        self.log(f"Total tournament time: {result.total_time_s / 60:.1f} minutes")
        self.log("=" * 60)

        # Save results
        results_path = self.output_dir / "tourney_results.json"
        results_path.write_text(json.dumps(result.to_dict(), indent=2))
        self.log(f"Results saved to {results_path}")

        bracket_path = self.output_dir / "tourney_bracket.md"
        bracket_path.write_text(render_bracket(result))
        self.log(f"Bracket saved to {bracket_path}")

        # Clean up checkpoint file on successful completion
        ckpt_path = self.output_dir / CHECKPOINT_FILENAME
        if ckpt_path.exists():
            ckpt_path.unlink()

        # Push winner
        if winner and winner.output_dir and (self.hub_org or self.hub_repo):
            self._push_winner(result)

        # Final yield with completed result
        yield ("Tournament complete", result)

    def _push_winner(self, result: TourneyResult):
        """Push the winning model to HuggingFace Hub."""
        winner = result.winner
        if not winner or not winner.output_dir:
            return

        try:
            from huggingface_hub import HfApi

            short_model = self.model_name.split("/")[-1] if "/" in self.model_name else self.model_name
            if self.hub_repo:
                repo_id = self.hub_repo
            else:
                repo_id = f"{self.hub_org}/{short_model}-OBLITERATED"
            result.hub_repo = repo_id

            self.log(f"\nPushing winner to Hub: {repo_id}")

            _token = resolve_first("HF_PUSH_TOKEN", "HF_TOKEN")
            api = HfApi(token=_token) if _token else HfApi()
            api.create_repo(repo_id, exist_ok=True)

            # Write model card
            model_card = generate_model_card(result)
            card_path = Path(winner.output_dir) / "README.md"
            card_path.write_text(model_card)

            # Write tourney results alongside model
            results_dest = Path(winner.output_dir) / "tourney_results.json"
            results_dest.write_text(json.dumps(result.to_dict(), indent=2))

            api.upload_folder(
                folder_path=winner.output_dir,
                repo_id=repo_id,
                commit_message=(
                    f"OBLITERATUS tourney: {winner.method} wins "
                    f"(score {winner.score:.4f}) on {self.model_name}"
                ),
            )
            self.log(f"Pushed to https://huggingface.co/{repo_id}")

        except Exception as e:
            self.log(f"Hub push failed: {e}")
