"""CLI entry point for Obliteratus — Master Ablation Suite."""

from __future__ import annotations

import argparse
import json
import os
import tempfile
from pathlib import Path

from rich.console import Console

from obliteratus import __version__

console = Console()

_BANNER = r"""
[bold red]
  ░▒█▀▀▀█ ░▒█▀▀▄ ░▒█░░░ ▀█▀ ▀▀█▀▀ ░▒█▀▀▀ ░▒█▀▀█ ▒█▀▀█ ▀▀█▀▀ ░▒█░░▒█ ░▒█▀▀▀█
  ░▒█░░▒█ ░▒█▀▀▄ ░▒█░░░ ░█░ ░░█░░ ░▒█▀▀▀ ░▒█▄▄▀ ▒█▄▄█ ░░█░░ ░▒█░░▒█ ░░▀▀▀▄▄
  ░▒█▄▄▄█ ░▒█▄▄▀ ░▒█▄▄█ ▄█▄ ░░▀░░ ░▒█▄▄▄ ░▒█░▒█ ▒█░▒█ ░░▀░░ ░░▒█▄▄█ ░▒█▄▄▄█
[/bold red]
[dim]  ════════════════════════════════════════════════════════════════════[/dim]
[bold white]   MASTER ABLATION SUITE[/bold white] [dim]//[/dim] [bold red]Break the chains. Free the mind.[/bold red]
[dim]  ════════════════════════════════════════════════════════════════════[/dim]
"""

_ASCII_BANNER = """
[bold red]OBLITERATUS[/bold red]
[dim]====================================================================[/dim]
[bold white]MASTER ABLATION SUITE[/bold white] [dim]//[/dim] [bold red]Break the chains. Free the mind.[/bold red]
[dim]====================================================================[/dim]
"""


def _console_text(text: str, fallback: str) -> str:
    """Return an ASCII fallback when stdout cannot encode ``text``."""
    import sys

    encoding = getattr(sys.stdout, "encoding", None) or "utf-8"
    try:
        text.encode(encoding)
    except (LookupError, UnicodeEncodeError):
        return fallback
    return text


def _banner_for_console() -> str:
    """Return a startup banner supported by the active stdout encoding."""
    return _console_text(_BANNER, _ASCII_BANNER)


def _positive_int(value: str) -> int:
    """Parse a strictly positive integer for public CLI boundaries."""
    try:
        parsed = int(value)
    except ValueError as exc:
        raise argparse.ArgumentTypeError("must be a positive integer") from exc
    if parsed <= 0:
        raise argparse.ArgumentTypeError("must be a positive integer")
    return parsed


def _ssh_port(value: str) -> int:
    """Parse the public SSH port boundary."""
    try:
        parsed = int(value)
    except ValueError as exc:
        raise argparse.ArgumentTypeError("must be an integer from 1 through 65535") from exc
    if not 1 <= parsed <= 65535:
        raise argparse.ArgumentTypeError("must be an integer from 1 through 65535")
    return parsed


def _utilization_fraction(value: str) -> float:
    """Parse a finite fraction in the public ``(0, 1]`` contract."""
    try:
        parsed = float(value)
    except ValueError as exc:
        raise argparse.ArgumentTypeError("must be a number in (0, 1]") from exc
    if not 0.0 < parsed <= 1.0:
        raise argparse.ArgumentTypeError("must be a number in (0, 1]")
    return parsed


def _add_gpu_args(parser):
    """Add --gpus flag for multi-GPU control."""
    gpu_group = parser.add_argument_group("GPU selection")
    gpu_group.add_argument(
        "--gpus", type=str, default=None, metavar="IDS",
        help=(
            "Comma-separated GPU IDs to use (e.g. '0,1,2,3' or 'all'). "
            "Sets CUDA_VISIBLE_DEVICES. By default uses all available GPUs. "
            "Models are automatically split across selected GPUs via accelerate."
        ),
    )


def _add_remote_args(parser):
    """Add --remote execution flags to a subcommand parser."""
    remote_group = parser.add_argument_group("remote execution")
    remote_group.add_argument(
        "--remote", type=str, default=None, metavar="[USER@]HOST",
        help="Run on a remote GPU node via SSH (e.g. root@gpu-node or just gpu-node)",
    )
    remote_group.add_argument(
        "--ssh-key", type=str, default=None,
        help="Path to SSH private key (default: use SSH agent or ~/.ssh/id_rsa)",
    )
    remote_group.add_argument(
        "--ssh-port", type=_ssh_port, default=22,
        help="SSH port on remote host (default: 22)",
    )
    remote_group.add_argument(
        "--remote-dir", type=str, default="/tmp/obliteratus_run",
        help="Working directory on the remote machine (default: /tmp/obliteratus_run)",
    )
    remote_group.add_argument(
        "--remote-python", type=str, default="python3",
        help="Python binary on the remote machine (default: python3)",
    )
    remote_group.add_argument(
        "--no-sync", action="store_true", default=False,
        help="Don't copy results back to local machine after remote run",
    )


def _apply_gpu_selection(args):
    """Set CUDA_VISIBLE_DEVICES based on --gpus flag (for local runs only)."""
    import os

    gpus = getattr(args, "gpus", None)
    if gpus is None or getattr(args, "remote", None):
        return  # skip for remote runs (handled by remote runner)

    if gpus.lower() == "all":
        return  # use all GPUs (default behavior)

    # Validate: should be comma-separated integers
    try:
        gpu_ids = [int(g.strip()) for g in gpus.split(",")]
    except ValueError:
        console.print(f"[red]Invalid --gpus value: {gpus!r}. Expected comma-separated integers or 'all'.[/]")
        raise SystemExit(1)

    os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(str(g) for g in gpu_ids)
    console.print(f"[dim]Using GPUs: {gpu_ids} (CUDA_VISIBLE_DEVICES={os.environ['CUDA_VISIBLE_DEVICES']})[/dim]")


def main(argv: list[str] | None = None):
    console.print(_banner_for_console())
    parser = argparse.ArgumentParser(
        prog="obliteratus",
        description="Master Ablation Suite for HuggingFace transformers",
    )
    parser.add_argument(
        "--version",
        action="version",
        version=f"%(prog)s {__version__}",
    )
    subparsers = parser.add_subparsers(dest="command", required=True)

    # --- run ---
    run_parser = subparsers.add_parser("run", help="Run an ablation from a YAML config")
    run_parser.add_argument("config", type=str, help="Path to YAML config file")
    run_parser.add_argument("--output-dir", type=str, default=None, help="Override output dir")
    run_parser.add_argument(
        "--preset",
        type=str,
        default=None,
        help="Apply a preset (e.g. quick, full, attention, jailbreak, guardrail)",
    )
    _add_gpu_args(run_parser)
    _add_remote_args(run_parser)

    # --- info ---
    info_parser = subparsers.add_parser("info", help="Print model architecture info")
    info_parser.add_argument("model", type=str, help="HuggingFace model name/path")
    info_parser.add_argument("--task", type=str, default="causal_lm", choices=["causal_lm", "classification"])
    info_parser.add_argument("--device", type=str, default="cpu")
    info_parser.add_argument("--dtype", type=str, default="float32")

    # --- interactive ---
    subparsers.add_parser(
        "interactive",
        help="Guided setup — pick hardware, model, and preset interactively",
    )

    # --- models ---
    models_parser = subparsers.add_parser("models", help="Browse curated models by compute tier")
    models_parser.add_argument(
        "--tier",
        type=str,
        default=None,
        choices=["tiny", "small", "medium", "large", "frontier"],
        help="Filter by compute tier",
    )

    # --- presets ---
    subparsers.add_parser("presets", help="Browse ablation presets (quick, full, jailbreak, etc.)")

    # --- strategies ---
    subparsers.add_parser("strategies", help="List available ablation strategies")

    # --- ui ---
    ui_parser = subparsers.add_parser(
        "ui",
        help="Launch the Gradio web UI locally (same UI as the HuggingFace Space)",
    )
    ui_parser.add_argument(
        "--port", type=int, default=7860, help="Server port (default: 7860)",
    )
    ui_parser.add_argument(
        "--host", type=str, default="0.0.0.0", help="Server host (default: 0.0.0.0)",
    )
    ui_parser.add_argument(
        "--share", action="store_true", help="Create a public Gradio share link",
    )
    ui_parser.add_argument(
        "--no-browser", action="store_true", help="Don't auto-open browser on launch",
    )
    ui_parser.add_argument(
        "--auth", type=str, default=None,
        help="Basic auth as user:pass",
    )
    ui_parser.add_argument(
        "--quiet", action="store_true", help="Suppress the startup banner",
    )

    # --- obliterate (primary) + abliterate (backward-compat alias) ---
    def _add_obliterate_args(p):
        p.add_argument("model", type=str, help="HuggingFace model name/path")
        p.add_argument("--output-dir", type=str, default=None, help="Where to save the obliterated model")
        p.add_argument("--device", type=str, default="auto")
        p.add_argument("--dtype", type=str, default="float16")
        p.add_argument(
            "--method", type=str, default="advanced",
            choices=[
                "basic", "advanced", "aggressive", "spectral_cascade",
                "informed", "surgical", "optimized", "som", "inverted", "nuclear",
            ],
            help="Liberation method (default: advanced)",
        )
        p.add_argument("--n-directions", type=int, default=None, help="Override: number of refusal directions to extract")
        p.add_argument(
            "--direction-method", type=str, default=None,
            choices=["diff_means", "svd", "leace", "som"],
            help="Direction extraction method: diff_means, svd, leace, or som",
        )
        p.add_argument("--regularization", type=float, default=None, help="Override: fraction to preserve (0.0-1.0)")
        p.add_argument("--refinement-passes", type=int, default=None, help="Override: number of iterative passes")
        p.add_argument(
            "--min-layer-fraction", type=float, default=None,
            help="Optional layer floor as fraction of depth; e.g. 0.75 keeps only the final quarter.",
        )
        p.add_argument(
            "--max-layer-fraction", type=float, default=None,
            help="Optional cap on selected layers as fraction of depth; e.g. 0.25 caps edits to one quarter.",
        )
        p.add_argument(
            "--harmless-pc-count", type=int, default=None,
            help="Subtract this many top harmless activation PCs from extracted refusal directions.",
        )
        p.add_argument(
            "--shield-concept-count", type=int, default=None,
            help="Use this many built-in capability/style shield atoms for direction cleaning.",
        )
        p.add_argument(
            "--shield-ridge", type=float, default=None,
            help="Ridge strength for shield atom residualization (default: 0.05).",
        )
        p.add_argument(
            "--shield-residualize", action="store_true", default=None,
            help="Directly residualize refusal directions against shield atoms.",
        )
        p.add_argument(
            "--shield-layer-penalty", type=float, default=None,
            help="Penalize layer selection by shield-atom overlap instead of editing directions.",
        )
        p.add_argument(
            "--projection-target", type=str, default=None,
            choices=["all", "attention", "ffn", "output"],
            help="Projection target modules: all, attention, ffn, or output.",
        )
        p.add_argument(
            "--projection-row-fraction", type=float, default=None,
            help="Only project the strongest fraction of output rows/columns per weight matrix.",
        )
        p.add_argument(
            "--quantization", type=str, default=None, choices=["4bit", "8bit"],
            help="Load model with quantization (4bit or 8bit). Requires bitsandbytes. "
                 "FP8/NVFP4 checkpoints are detected and dequantized automatically.",
        )
        p.add_argument(
            "--trust-remote-code", action="store_true", default=False,
            help="Trust custom modeling code from the model repo (required for "
                 "checkpoints with custom architectures, e.g. Nemotron Omni).",
        )
        p.add_argument(
            "--gpu-memory-utilization",
            type=_utilization_fraction,
            default=None,
            metavar="FRACTION",
            help="Override the GPU VRAM fraction available for model loading (0, 1]. "
            "By default, reserves 15%% or 2 GiB per GPU, whichever is larger.",
        )
        p.add_argument(
            "--large-model", action="store_true", default=False,
            help="Enable conservative defaults for 120B+ models (fewer directions, 1 pass, lower SAE expansion).",
        )
        p.add_argument(
            "--verify-sample-size", type=int, default=None,
            help="Number of harmful prompts to test for refusal rate (default: 30). "
                 "Increase for tighter confidence intervals (e.g. 100 for ~1%% resolution).",
        )
        p.add_argument(
            "--refusal-max-tokens", type=_positive_int, default=None,
            help="Max new tokens to generate per response in the refusal test (default: 128).",
        )
        p.add_argument(
            "--max-perplexity-increase", type=float, default=3.0,
            help="Fail if post-surgery perplexity exceeds this multiple of stock (default: 3.0).",
        )
        p.add_argument(
            "--min-coherence-retention", type=float, default=0.5,
            help="Minimum post-surgery coherence divided by stock coherence (default: 0.5).",
        )
        p.add_argument(
            "--max-degenerate-fraction", type=float, default=0.2,
            help="Fail if repetitive/empty completions exceed this fraction (default: 0.2).",
        )
        p.add_argument(
            "--dataset", type=str, default="builtin",
            help="Prompt dataset source for contrastive extraction when using residue mining (default: builtin).",
        )
        prompt_source_group = p.add_mutually_exclusive_group()
        prompt_source_group.add_argument(
            "--prompt-pairs-file",
            "--prompt-pair-file",
            dest="prompt_pairs_file",
            type=str,
            default=None,
            metavar="PATH",
            help=(
                "Local UTF-8 JSON file with exactly 'harmful' and 'harmless' arrays "
                "for contrastive extraction."
            ),
        )
        prompt_source_group.add_argument(
            "--residue-file", action="append", default=[],
            help="Refusal-audit/residue JSON to upweight as hard negatives. Can be passed multiple times.",
        )
        p.add_argument(
            "--residue-weight", type=int, default=5,
            help="How many extra copies to add per mined residue prompt (default: 5).",
        )
        p.add_argument(
            "--residue-max", type=int, default=None,
            help="Optional cap on unique residue prompts loaded from residue files.",
        )
        p.add_argument(
            "--contribute", action="store_true", default=False,
            help="Save a community contribution record after the run completes.",
        )
        p.add_argument(
            "--contribute-notes", type=str, default="",
            help="Optional notes to include with the community contribution.",
        )

    abl_parser = subparsers.add_parser(
        "obliterate",
        help="One-click: remove refusal directions from a model (SOTA multi-technique)",
    )
    _add_obliterate_args(abl_parser)
    _add_gpu_args(abl_parser)
    _add_remote_args(abl_parser)
    # Backward-compat alias (hidden from help)
    abl_alias = subparsers.add_parser("abliterate", help=argparse.SUPPRESS)
    _add_obliterate_args(abl_alias)
    _add_gpu_args(abl_alias)
    _add_remote_args(abl_alias)

    # --- self-improve ---
    si_parser = subparsers.add_parser(
        "self-improve",
        help="Recursive hard-negative mining: turn refusal audits into weighted residue runs",
    )
    si_parser.add_argument("model", type=str, help="Source model or previous artifact to improve")
    si_parser.add_argument("--audit", action="append", required=True, help="Refusal-audit JSON/summary. Can be repeated.")
    si_parser.add_argument("--output-dir", type=str, required=True, help="Output directory for the next candidate")
    si_parser.add_argument("--residue-out", type=str, default=None, help="Where to write canonical residue JSON")
    si_parser.add_argument("--dataset", type=str, default="builtin")
    si_parser.add_argument("--residue-weight", type=int, default=None)
    si_parser.add_argument("--residue-max", type=int, default=None)
    si_parser.add_argument("--params-b", type=float, default=None, help="Override detected total parameter count in billions for planning/defaults")
    si_parser.add_argument("--no-param-auto-scale", action="store_true", default=False, help="Do not apply size-aware defaults when knobs are omitted")
    si_parser.add_argument("--method", type=str, default="advanced", choices=["basic", "advanced", "aggressive", "spectral_cascade", "informed", "surgical", "optimized", "som", "inverted", "nuclear"])
    si_parser.add_argument("--direction-method", type=str, default="diff_means", choices=["diff_means", "svd", "leace", "som"])
    si_parser.add_argument("--n-directions", type=int, default=None)
    si_parser.add_argument("--regularization", type=float, default=None)
    si_parser.add_argument("--refinement-passes", type=int, default=None)
    si_parser.add_argument("--min-layer-fraction", type=float, default=None)
    si_parser.add_argument("--max-layer-fraction", type=float, default=None)
    si_parser.add_argument("--harmless-pc-count", type=int, default=None)
    si_parser.add_argument("--shield-concept-count", type=int, default=None)
    si_parser.add_argument("--shield-ridge", type=float, default=None)
    si_parser.add_argument("--shield-residualize", action="store_true", default=None)
    si_parser.add_argument("--shield-layer-penalty", type=float, default=None)
    si_parser.add_argument("--projection-target", type=str, default=None, choices=["all", "attention", "ffn", "output"])
    si_parser.add_argument("--projection-row-fraction", type=float, default=None)
    si_parser.add_argument("--device", type=str, default="auto")
    si_parser.add_argument("--dtype", type=str, default="float16")
    si_parser.add_argument("--verify-sample-size", type=int, default=None)
    si_parser.add_argument("--dry-run", action="store_true", default=False, help="Only write residue/plan; do not run surgery")
    _add_gpu_args(si_parser)

    # --- report ---
    report_parser = subparsers.add_parser("report", help="Regenerate report from saved results")
    report_parser.add_argument("results_json", type=str, help="Path to results.json")
    report_parser.add_argument("--output-dir", type=str, default=None)

    # --- aggregate ---
    aggregate_parser = subparsers.add_parser("aggregate", help="Aggregate community contribution results")
    aggregate_parser.add_argument(
        "--dir", type=str, default="community_results",
        help="Directory containing contribution JSON files",
    )

    # --- restore-multimodal ---
    restore_parser = subparsers.add_parser(
        "restore-multimodal",
        help="Restore vision/MTP/linear-attn components after text-only abliteration",
    )
    restore_parser.add_argument("--abliterated", required=True,
                                help="Text-only abliterated model directory")
    restore_parser.add_argument("--stock", required=True,
                                help="Stock (full) model directory or HF repo")
    restore_parser.add_argument("--output", required=True,
                                help="Output directory for merged model")

    # --- capability-check ---
    capcheck_parser = subparsers.add_parser(
        "capability-check",
        help="Compare abliterated vs stock on MMLU via lm-eval-harness",
    )
    capcheck_parser.add_argument("--abliterated", required=True, help="Abliterated model path or HF repo")
    capcheck_parser.add_argument("--stock", required=True, help="Stock model path or HF repo")
    capcheck_parser.add_argument("--device", type=str, default="auto")
    capcheck_parser.add_argument("--dtype", type=str, default="bfloat16")
    capcheck_parser.add_argument("--quick", action="store_true", help="5 subjects only (~1 min per model)")
    capcheck_parser.add_argument("--subjects", type=str, default=None, help="Comma-separated subject list")
    capcheck_parser.add_argument("--limit", type=int, default=5, help="Questions per subject")
    capcheck_parser.add_argument("--output-dir", type=str, default=None)

    # --- blend ---
    blend_parser = subparsers.add_parser(
        "blend",
        help="Complementary abliteration blending — interpolate two surgeries in weight space",
    )
    blend_parser.add_argument("--model-a", required=True, help="First model (e.g., aggressive surgery)")
    blend_parser.add_argument("--model-b", required=True, help="Second model (e.g., LEACE surgery)")
    blend_parser.add_argument("--alpha", type=float, default=0.6, help="Blend ratio (0=pure A, 1=pure B)")
    blend_parser.add_argument(
        "--config-source",
        choices=["a", "b"],
        default="a",
        help="Model whose tokenizer/config support files are copied",
    )
    blend_parser.add_argument(
        "--allow-unverified-lineage",
        action="store_true",
        help="Allow checkpoints without matching OBLITERATUS source-model metadata",
    )
    blend_parser.add_argument("--search", type=str, default=None, help="Comma-separated alphas for search")
    blend_parser.add_argument("--output", required=True, help="Output directory")
    aggregate_parser.add_argument(
        "--format",
        choices=["summary", "latex"],
        default="summary",
        help="Output format (default: summary)",
    )
    aggregate_parser.add_argument(
        "--metric",
        type=str,
        default="refusal_rate",
        help="Metric to display in LaTeX output (default: refusal_rate)",
    )
    aggregate_parser.add_argument(
        "--min-runs",
        type=int,
        default=1,
        help="Minimum runs per (model, method) to include (default: 1)",
    )

    # --- tourney ---
    tourney_parser = subparsers.add_parser(
        "tourney",
        help="March Madness tournament — pit all methods against each other, push winner to Hub",
    )
    tourney_parser.add_argument("model", type=str, help="HuggingFace model name/path")
    tourney_parser.add_argument("--hub-org", type=str, default=None, help="HF org to push winner (e.g. my-org)")
    tourney_parser.add_argument("--hub-repo", type=str, default=None, help="Full HF repo ID (overrides --hub-org)")
    tourney_parser.add_argument("--device", type=str, default="auto")
    tourney_parser.add_argument("--dtype", type=str, default="float16")
    tourney_parser.add_argument("--dataset", type=str, default="builtin", help="Dataset source (default: builtin)")
    tourney_parser.add_argument(
        "--quantization", type=str, default=None, choices=["4bit", "8bit"],
        help="Load model with quantization (FP8/NVFP4 checkpoints are "
             "detected and dequantized automatically)",
    )
    tourney_parser.add_argument(
        "--output-dir",
        type=str,
        default=os.path.join(tempfile.gettempdir(), "obliteratus_tourney"),
    )
    tourney_parser.add_argument(
        "--methods", type=str, nargs="+", default=None,
        help="Override: only run these methods (space-separated)",
    )
    _add_gpu_args(tourney_parser)
    _add_remote_args(tourney_parser)

    # --- recommend ---
    recommend_parser = subparsers.add_parser(
        "recommend",
        help="Show telemetry-driven best method + hyperparams for a model",
    )
    recommend_parser.add_argument("model", type=str, help="HuggingFace model name/path")
    recommend_parser.add_argument("--device", type=str, default="cpu")
    recommend_parser.add_argument("--dtype", type=str, default="float32")
    recommend_parser.add_argument(
        "--insights", action="store_true", default=False,
        help="Also show global cross-architecture insights",
    )

    # --- gpu-calc ---
    calc_parser = subparsers.add_parser(
        "gpu-calc",
        help="Estimate minimum GPUs needed for a model",
    )
    calc_parser.add_argument(
        "model", type=str, nargs="?", default=None,
        help="HuggingFace model name/path (auto-fetches param counts)",
    )
    calc_parser.add_argument(
        "--params", type=float, default=None, metavar="B",
        help="Total parameters in billions (overrides auto-detection)",
    )
    calc_parser.add_argument(
        "--active-params", type=float, default=None, metavar="B",
        help="Active parameters in billions (for MoE models; defaults to --params)",
    )
    calc_parser.add_argument(
        "--dtype", type=str, default="bfloat16",
        choices=["float32", "float16", "bfloat16", "int8", "int4"],
        help="Data type for model weights (default: bfloat16)",
    )
    calc_parser.add_argument(
        "--gpu-mem", type=float, default=80.0, metavar="GB",
        help="VRAM per GPU in GB (default: 80 for A100-80GB)",
    )

    # --- durable headless runs ---
    runs_parser = subparsers.add_parser(
        "runs", help="Launch and manage reconnect-safe durable obliteration runs",
    )
    runs_parser.add_argument(
        "--archive-root", default=None,
        help="Durable run root (default: OBLITERATUS_RUN_ARCHIVE or the user state directory)",
    )
    runs_subparsers = runs_parser.add_subparsers(dest="runs_command", required=True)
    runs_launch = runs_subparsers.add_parser("launch", help="Launch a detached obliteration run")
    runs_launch.add_argument("--notes", default="", help="Operator notes retained with the run")
    runs_launch.add_argument(
        "arguments", nargs=argparse.REMAINDER,
        help="Arguments after --, beginning with the model ID",
    )
    for command_name in ("status", "cancel", "result"):
        command_parser = runs_subparsers.add_parser(command_name)
        command_parser.add_argument("run_id")
    runs_subparsers.add_parser("list")
    runs_prune = runs_subparsers.add_parser(
        "prune-checkpoint", help="Delete model payload while retaining evidence",
    )
    runs_prune.add_argument("run_id")
    runs_prune.add_argument("--reason", required=True)

    args = parser.parse_args(argv)

    if getattr(args, "prompt_pairs_file", None):
        residue_only_options = []
        if getattr(args, "dataset", "builtin") != "builtin":
            residue_only_options.append("--dataset")
        if getattr(args, "residue_weight", 5) != 5:
            residue_only_options.append("--residue-weight")
        if getattr(args, "residue_max", None) is not None:
            residue_only_options.append("--residue-max")
        if residue_only_options:
            parser.error(
                f"{', '.join(residue_only_options)} can only be used with --residue-file, "
                "not --prompt-pairs-file"
            )

    # Apply GPU selection early (before any CUDA init)
    _apply_gpu_selection(args)

    if args.command == "gpu-calc":
        _cmd_gpu_calc(args)
        return
    elif args.command == "runs":
        _cmd_runs(args)
    elif args.command == "run":
        if getattr(args, "remote", None):
            _cmd_remote_run(args)
        else:
            _cmd_run(args)
    elif args.command == "interactive":
        _cmd_interactive()
    elif args.command == "models":
        _cmd_models(args)
    elif args.command == "presets":
        _cmd_presets()
    elif args.command == "info":
        _cmd_info(args)
    elif args.command == "strategies":
        _cmd_strategies()
    elif args.command == "report":
        _cmd_report(args)
    elif args.command == "aggregate":
        _cmd_aggregate(args)
    elif args.command == "restore-multimodal":
        from obliteratus.restore_multimodal import restore_multimodal
        import logging
        logging.basicConfig(level=logging.INFO, format="%(message)s")
        result = restore_multimodal(args.abliterated, args.stock, args.output)
        print(f"\nDone: {result['replaced']} abliterated + {result['kept']} stock = {result['total']} total")
    elif args.command == "capability-check":
        from obliteratus.capability_check import capability_check
        subjects = args.subjects.split(",") if args.subjects else None
        result = capability_check(
            args.abliterated, args.stock,
            device=args.device, dtype=args.dtype,
            quick=args.quick, subjects=subjects,
            limit=args.limit, output_dir=args.output_dir,
        )
        print(f"\nStock: {result['stock_acc']*100:.1f}%  Abliterated: {result['abliterated_acc']*100:.1f}%  Delta: {result['delta_pp']:+.1f}pp")
    elif args.command == "blend":
        from obliteratus.blend import blend_models, blend_search
        import logging
        logging.basicConfig(level=logging.INFO, format="%(message)s")
        if args.search:
            alphas = [float(a) for a in args.search.split(",")]
            results = blend_search(
                args.model_a,
                args.model_b,
                args.output,
                alphas,
                config_source=args.config_source,
                allow_unverified_lineage=args.allow_unverified_lineage,
            )
            print(f"\nCreated {len(results)} blends")
        else:
            result = blend_models(
                args.model_a,
                args.model_b,
                args.output,
                alpha=args.alpha,
                config_source=args.config_source,
                allow_unverified_lineage=args.allow_unverified_lineage,
            )
            print(f"\nBlend complete: {result['blended_tensors']} tensors at alpha={args.alpha}")
    elif args.command == "ui":
        _cmd_ui(args)
    elif args.command == "recommend":
        _cmd_recommend(args)
    elif args.command == "tourney":
        if getattr(args, "remote", None):
            _cmd_remote_tourney(args)
        else:
            _cmd_tourney(args)
    elif args.command == "self-improve":
        _cmd_self_improve(args)
    elif args.command in ("obliterate", "abliterate"):
        if getattr(args, "remote", None):
            _cmd_remote_abliterate(args)
        else:
            _cmd_abliterate(args)


def _cmd_runs(args):
    """Dispatch the stable durable-run control surface."""

    import json

    from obliteratus.run_archive import RunArchive

    archive = RunArchive(args.archive_root)
    if args.runs_command == "launch":
        arguments = list(args.arguments)
        if arguments[:1] == ["--"]:
            arguments = arguments[1:]
        run_id = archive.launch(arguments, notes=args.notes)
        payload = archive.status(run_id)
    elif args.runs_command == "status":
        payload = archive.status(args.run_id)
    elif args.runs_command == "cancel":
        payload = archive.cancel(args.run_id)
    elif args.runs_command == "result":
        payload = archive.result(args.run_id)
    elif args.runs_command == "prune-checkpoint":
        payload = archive.prune_checkpoint(args.run_id, reason=args.reason)
    else:
        payload = archive.list()
    print(json.dumps(payload, indent=2, sort_keys=True))


def _cmd_self_improve(args):
    """Run one recursive hard-negative OBLITERATUS iteration."""

    from rich.panel import Panel
    from rich.table import Table

    from obliteratus.hard_negative import (
        build_weighted_prompt_pairs,
        load_residue_file,
        save_residue_file,
    )
    from obliteratus.model_profile import ModelProfile, default_self_improve_params, profile_model

    profile = profile_model(args.model, dtype=args.dtype)
    if args.params_b is not None:
        profile = ModelProfile(
            **{**profile.to_json(), "total_params": int(args.params_b * 1e9), "total_params_b": args.params_b, "active_params_b": args.params_b, "source": "cli_override"}
        )
    size_defaults = default_self_improve_params(profile) if not args.no_param_auto_scale else {
        "n_directions": 3,
        "regularization": 0.30,
        "refinement_passes": 1,
        "residue_weight": 5,
        "verify_sample_size": 30,
        "note": "parameter auto-scale disabled; using generic defaults for omitted knobs.",
    }
    n_directions = args.n_directions if args.n_directions is not None else size_defaults["n_directions"]
    regularization = args.regularization if args.regularization is not None else size_defaults["regularization"]
    refinement_passes = args.refinement_passes if args.refinement_passes is not None else size_defaults["refinement_passes"]
    residue_weight = args.residue_weight if args.residue_weight is not None else size_defaults["residue_weight"]
    verify_sample_size = args.verify_sample_size if args.verify_sample_size is not None else size_defaults["verify_sample_size"]

    residue_examples = []
    for audit in args.audit:
        residue_examples.extend(load_residue_file(audit))
    residue_path = Path(args.residue_out or Path(args.output_dir, "mined_residue.json"))
    save_residue_file(residue_examples, residue_path)

    harmful, harmless, meta = build_weighted_prompt_pairs(
        base_dataset=args.dataset,
        residue_files=[residue_path],
        residue_weight=residue_weight,
        max_residue=args.residue_max,
    )

    meta["model_profile"] = profile.to_json()
    meta["size_aware_defaults"] = size_defaults
    meta["effective_settings"] = {
        "n_directions": n_directions,
        "regularization": regularization,
        "refinement_passes": refinement_passes,
        "residue_weight": residue_weight,
        "verify_sample_size": verify_sample_size,
        "projection_row_fraction": args.projection_row_fraction,
    }
    plan_path = Path(args.output_dir, "self_improve_plan.json")
    plan_path.parent.mkdir(parents=True, exist_ok=True)
    plan_path.write_text(json.dumps(meta, indent=2))

    table = Table(title="Recursive hard-negative mining plan")
    table.add_column("Field", style="cyan")
    table.add_column("Value", style="green")
    table.add_row("source model", args.model)
    table.add_row("params", f"{profile.total_params_b:.3f}B" if profile.total_params_b is not None else "unknown")
    table.add_row("profile source", profile.source)
    table.add_row("layers / hidden", f"{profile.num_layers or '?'} / {profile.hidden_size or '?'}")
    table.add_row("size default note", size_defaults["note"])
    table.add_row("output", args.output_dir)
    table.add_row("residue file", str(residue_path))
    table.add_row("plan file", str(plan_path))
    table.add_row("unique residue prompts", str(meta["residue_examples"]))
    table.add_row("weighted residue pairs", str(meta["residue_added_pairs"]))
    table.add_row("total prompt pairs", str(meta["total_pairs"]))
    table.add_row("method", args.method)
    table.add_row("directions", str(n_directions))
    table.add_row("regularization", str(regularization))
    table.add_row("refinement passes", str(refinement_passes))
    table.add_row("verify sample size", str(verify_sample_size))
    console.print(table)

    if args.dry_run:
        console.print(
            Panel(
                "[bold yellow]Dry run only.[/]\n\n"
                f"Plan JSON: {plan_path}\n"
                "Run without --dry-run to launch surgery, or pass this residue file to:\n"
                f"  obliteratus obliterate {args.model} --residue-file {residue_path} --output-dir {args.output_dir}",
                border_style="yellow",
                title="SELF-IMPROVE PLAN WRITTEN",
            )
        )
        return

    from obliteratus.abliterate import AbliterationPipeline

    log_lines: list[str] = []

    def on_log(msg):
        log_lines.append(msg)
        console.print(f"[dim]>[/] {msg}")

    pipeline = AbliterationPipeline(
        model_name=args.model,
        output_dir=args.output_dir,
        device=args.device,
        dtype=args.dtype,
        method=args.method,
        n_directions=n_directions,
        direction_method=args.direction_method,
        regularization=regularization,
        refinement_passes=refinement_passes,
        min_layer_fraction=args.min_layer_fraction,
        max_layer_fraction=args.max_layer_fraction,
        harmless_pc_count=args.harmless_pc_count,
        shield_concept_count=args.shield_concept_count,
        shield_ridge=args.shield_ridge,
        shield_residualize=args.shield_residualize,
        shield_layer_penalty=args.shield_layer_penalty,
        projection_target=args.projection_target,
        projection_row_fraction=args.projection_row_fraction,
        verify_sample_size=verify_sample_size,
        harmful_prompts=harmful,
        harmless_prompts=harmless,
        on_log=on_log,
    )
    result_path = pipeline.run()
    Path(result_path, "hard_negative_residue.json").write_text(json.dumps(meta, indent=2))
    console.print(
        Panel(
            f"[bold green]Self-improvement iteration complete![/]\n\n"
            f"  Model saved to: [cyan]{result_path}[/]\n"
            f"  Residue metadata: [cyan]{result_path}/hard_negative_residue.json[/]\n"
            f"  Next loop: audit refusals from this artifact and feed them back with --audit.",
            border_style="green",
            title="RECURSIVE REBIRTH COMPLETE",
        )
    )



def _cmd_ui(args):
    from obliteratus.local_ui import launch_local_ui

    auth = tuple(args.auth.split(":", 1)) if args.auth else None
    launch_local_ui(
        host=args.host,
        port=args.port,
        share=args.share,
        open_browser=not args.no_browser,
        auth=auth,
        quiet=args.quiet,
    )


def _cmd_interactive():
    from obliteratus.interactive import run_interactive
    run_interactive()


def _cmd_models(args):
    from rich.table import Table

    from obliteratus.presets import get_presets_by_tier, list_all_presets

    presets = get_presets_by_tier(args.tier) if args.tier else list_all_presets()

    table = Table(title="Model Library — Curated Targets")
    table.add_column("Model", style="green")
    table.add_column("HuggingFace ID", style="cyan")
    table.add_column("Params", justify="right")
    table.add_column("Tier", style="yellow")
    table.add_column("Dtype")
    table.add_column("Quant")
    table.add_column("Description")

    for p in presets:
        table.add_row(
            p.name,
            p.hf_id,
            p.params,
            p.tier.upper(),
            p.recommended_dtype,
            p.recommended_quantization or "—",
            p.description,
        )

    console.print(table)
    console.print(
        "\n[dim]Tiers: TINY = CPU/laptop | SMALL = 4-8GB | "
        "MEDIUM = 8-16GB | LARGE = 24GB+ | FRONTIER = multi-GPU/cloud[/dim]"
    )


def _cmd_presets():
    from rich.table import Table

    from obliteratus.study_presets import list_study_presets

    presets = list_study_presets()

    table = Table(title="Ablation Presets")
    table.add_column("Key", style="cyan", min_width=12)
    table.add_column("Name", style="green")
    table.add_column("Strategies", style="yellow")
    table.add_column("Samples", justify="right")
    table.add_column("Description", max_width=55)

    for p in presets:
        strats = ", ".join(s["name"] for s in p.strategies)
        table.add_row(p.key, p.name, strats, str(p.max_samples), p.description)

    console.print(table)
    console.print(
        "\n[dim]Usage:  obliteratus run config.yaml --preset quick\n"
        "   or:  set preset: quick in your YAML file[/dim]"
    )


def _cmd_run(args):
    from obliteratus.config import StudyConfig
    from obliteratus.runner import run_study

    config = StudyConfig.from_yaml(args.config)
    # If --preset flag given, inject it so from_dict picks it up
    if args.preset:
        import yaml

        raw = yaml.safe_load(Path(args.config).read_text())
        raw["preset"] = args.preset
        config = StudyConfig.from_dict(raw)
    if args.output_dir:
        config.output_dir = args.output_dir

    # If YAML has a remote: section, dispatch to remote runner
    if config.remote is not None:
        from obliteratus.remote import RemoteConfig as _RC, RemoteRunner

        rc = _RC(
            host=config.remote.host,
            user=config.remote.user,
            port=config.remote.port,
            ssh_key=config.remote.ssh_key,
            known_hosts_file=config.remote.known_hosts_file,
            remote_dir=config.remote.remote_dir,
            install_timeout=config.remote.install_timeout,
            python=config.remote.python,
            sync_results=config.remote.sync_results,
            gpus=config.remote.gpus,
            install_source=config.remote.install_source,
        )
        runner = RemoteRunner(rc)
        result_path = runner.run_config(
            local_config_path=args.config,
            local_output_dir=config.output_dir,
            preset=args.preset,
        )
        if result_path:
            console.print(f"\n[bold green]Remote run complete.[/] Results at: [cyan]{result_path}[/]")
        else:
            console.print("[red]Remote run failed. Check logs above.[/]")
            raise SystemExit(1)
        return

    run_study(config)


def _cmd_info(args):
    from obliteratus.models.loader import load_model

    console.print(f"[bold cyan]Loading model:[/bold cyan] {args.model}")
    handle = load_model(
        model_name=args.model,
        task=args.task,
        device=args.device,
        dtype=args.dtype,
    )
    summary = handle.summary()
    for key, val in summary.items():
        if isinstance(val, int) and val > 1000:
            console.print(f"  {key}: {val:,}")
        else:
            console.print(f"  {key}: {val}")


def _cmd_strategies():
    from obliteratus.strategies import STRATEGY_REGISTRY

    console.print("[bold]Available ablation strategies:[/bold]\n")
    for name, cls in sorted(STRATEGY_REGISTRY.items()):
        doc = (cls.__doc__ or "").strip().split("\n")[0]
        console.print(f"  [cyan]{name}[/cyan]  — {doc}")


def _cmd_report(args):
    from obliteratus.reporting.report import AblationReport, AblationResult

    path = Path(args.results_json)
    data = json.loads(path.read_text())

    report = AblationReport(model_name=data["model_name"])
    report.add_baseline(data["baseline_metrics"])
    for r in data["results"]:
        report.add_result(
            AblationResult(
                strategy=r["strategy"],
                component=r["component"],
                description=r["description"],
                metrics=r["metrics"],
                metadata=r.get("metadata"),
            )
        )

    report.print_summary()

    output_dir = Path(args.output_dir) if args.output_dir else path.parent
    output_dir.mkdir(parents=True, exist_ok=True)
    metric_name = list(data["baseline_metrics"].keys())[0]
    try:
        report.plot_impact(metric=metric_name, output_path=output_dir / "impact.png")
        report.plot_heatmap(output_path=output_dir / "heatmap.png")
        console.print(f"\nPlots saved to {output_dir}/")
    except Exception as e:
        console.print(f"[yellow]Could not generate plots: {e}[/yellow]")


def _cmd_aggregate(args):
    from obliteratus.community import aggregate_results, generate_latex_table, load_contributions

    contrib_dir = args.dir
    records = load_contributions(contrib_dir)
    if not records:
        console.print(f"[yellow]No contributions found in {contrib_dir}[/yellow]")
        return

    aggregated = aggregate_results(records)
    min_runs = max(args.min_runs, 1)
    if min_runs > 1:
        for model in list(aggregated.keys()):
            for method in list(aggregated[model].keys()):
                if aggregated[model][method]["n_runs"] < min_runs:
                    del aggregated[model][method]
            if not aggregated[model]:
                del aggregated[model]

    if not aggregated:
        console.print("[yellow]No results meet the minimum run threshold.[/yellow]")
        return

    if args.format == "latex":
        console.print(generate_latex_table(aggregated, metric=args.metric))
        return

    from rich.table import Table

    table = Table(title="Aggregated Community Results")
    table.add_column("Model", style="green")
    table.add_column("Method", style="cyan")
    table.add_column("Runs", justify="right")
    table.add_column("Mean Refusal", justify="right")
    table.add_column("Mean Perplexity", justify="right")

    for model_name, methods in sorted(aggregated.items()):
        for method_name, stats in sorted(methods.items()):
            refusal = stats.get("refusal_rate", {}).get("mean", "N/A")
            ppl = stats.get("perplexity", {}).get("mean", "N/A")
            if isinstance(refusal, float):
                refusal = f"{refusal:.4f}"
            if isinstance(ppl, float):
                ppl = f"{ppl:.2f}"
            table.add_row(
                model_name.split("/")[-1] if "/" in model_name else model_name,
                method_name,
                str(stats["n_runs"]),
                str(refusal),
                str(ppl),
            )

    console.print(table)


def _cmd_recommend(args):
    from rich.markdown import Markdown
    from rich.panel import Panel

    from obliteratus.architecture_profiles import detect_architecture, enhance_profile_with_telemetry
    from obliteratus.adaptive_defaults import format_recommendation, get_global_insights

    model_name = args.model
    console.print(f"\nAnalyzing [bold]{model_name}[/]...")

    # Detect architecture
    try:
        from transformers import AutoConfig
        config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
        num_layers = getattr(config, "num_hidden_layers", 0)
        hidden_size = getattr(config, "hidden_size", 0)
    except Exception:
        config = None
        num_layers = 0
        hidden_size = 0

    profile = detect_architecture(model_name, config, num_layers, hidden_size)
    profile, rec = enhance_profile_with_telemetry(profile)

    console.print(Panel(
        f"[bold]{profile.profile_label}[/]\n"
        f"Architecture: {profile.arch_class.value} | Reasoning: {profile.reasoning_class.value}\n"
        f"Params: ~{profile.total_params_b:.1f}B | Layers: {profile.num_layers} | "
        f"Hidden: {profile.hidden_size}",
        title="Architecture Profile",
        border_style="cyan",
    ))

    if rec:
        md = format_recommendation(rec)
        console.print(Markdown(md))
    else:
        console.print("\n[yellow]Could not fetch telemetry — using research-grounded defaults.[/]")

    console.print(f"\n[bold green]Research default method:[/] {profile.recommended_method}")
    if profile.method_overrides:
        console.print("[bold green]Overrides:[/]")
        for k, v in sorted(profile.method_overrides.items()):
            console.print(f"  {k}: {v}")

    if args.insights:
        console.print("\n")
        console.rule("[bold magenta]Global Telemetry Insights")
        insights = get_global_insights()
        console.print(f"Total records analyzed: {insights['total_records']}")
        if insights["overall_best_methods"]:
            console.print("\n[bold]Overall method ranking (all architectures):[/]")
            for entry in insights["overall_best_methods"][:10]:
                console.print(
                    f"  {entry['method']}: {entry['mean_score']:.4f} "
                    f"({entry['n_runs']} runs)"
                )
        if insights["architecture_breakdown"]:
            console.print("\n[bold]Per-architecture breakdown:[/]")
            for label, info in insights["architecture_breakdown"].items():
                console.print(
                    f"  {label}: best={info['best_method']} "
                    f"({info['best_score']:.4f}), "
                    f"{info['n_methods_tested']} methods tested, "
                    f"{info['total_runs']} runs"
                )


def _cmd_tourney(args):
    from obliteratus.tourney import TourneyRunner

    def on_log(msg):
        console.print(msg)

    def on_round(rnd):
        console.print()
        console.rule(f"[bold green]Round {rnd.round_num} complete — "
                     f"{len(rnd.advanced_to)} advance, {len(rnd.eliminated)} eliminated")

    runner = TourneyRunner(
        model_name=args.model,
        hub_org=args.hub_org,
        hub_repo=args.hub_repo,
        device=args.device,
        dtype=args.dtype,
        dataset_key=args.dataset,
        quantization=args.quantization,
        methods=args.methods,
        output_dir=args.output_dir,
        on_log=on_log,
        on_round=on_round,
    )

    result = runner.run()

    if result.winner:
        console.print()
        console.rule("[bold magenta]TOURNAMENT CHAMPION", style="magenta")
        console.print(f"  [bold]{result.winner.method}[/] — score {result.winner.score:.4f}")
        console.print(f"  Refusal rate: {result.winner.metrics.get('refusal_rate', '?')}")
        console.print(f"  Coherence:    {result.winner.metrics.get('coherence', '?')}")
        if result.hub_repo:
            console.print(f"  Pushed to:    [link=https://huggingface.co/{result.hub_repo}]{result.hub_repo}[/link]")
        console.print(f"\n  Full bracket: {args.output_dir}/tourney_bracket.md")


def _cmd_abliterate(args):
    from rich.live import Live
    from rich.panel import Panel
    from rich.table import Table
    from rich.text import Text

    from obliteratus.abliterate import METHODS, STAGES, AbliterationPipeline

    model_name = args.model
    output_dir = args.output_dir or f"abliterated/{model_name.replace('/', '_')}"
    method = args.method
    method_label = METHODS.get(method, {}).get("label", method)

    # Stage state tracking
    stage_status = {s.key: "waiting" for s in STAGES}
    stage_msgs = {s.key: "" for s in STAGES}
    log_lines: list[str] = []

    def make_display():
        table = Table(show_header=False, expand=True, border_style="green")
        table.add_column("", width=6)
        table.add_column("Stage", min_width=10)
        table.add_column("Status", min_width=50)
        done_icon = _console_text("[bold green]✓[/]", "[bold green]OK[/]")
        running_icon = _console_text("[bold yellow]⚡[/]", "[bold yellow]*[/]")
        waiting_icon = _console_text("[dim]○[/]", "[dim]o[/]")
        done_bar = _console_text("█" * 20, "#" * 20)
        running_bar = _console_text("▓" * 10 + "░" * 10, "=" * 10 + "." * 10)
        waiting_bar = _console_text("░" * 20, "." * 20)
        header_sep = _console_text("—", "-")
        target_arrow = _console_text("→", "->")
        log_rule = _console_text("─── LOG ───", "--- LOG ---")
        panel_title = _console_text(
            "[bold green]⚗ ABLITERATE ⚗[/]",
            "[bold green]ABLITERATE[/]",
        )
        for i, s in enumerate(STAGES):
            st = stage_status[s.key]
            if st == "done":
                icon = done_icon
                bar = "[green]" + done_bar + "[/]"
            elif st == "running":
                icon = running_icon
                bar = "[yellow]" + running_bar + "[/]"
            else:
                icon = waiting_icon
                bar = "[dim]" + waiting_bar + "[/]"
            msg = stage_msgs.get(s.key, "")
            table.add_row(
                f"[cyan][{i + 1}/6][/]",
                f"{icon} [bold]{s.name}[/]",
                f"{bar}  {msg}",
            )

        header = Text.from_markup(
            f"[bold green]OBLITERATUS {header_sep} ABLITERATION PIPELINE[/]\n"
            f"[dim]Target:[/] [cyan]{model_name}[/]  {target_arrow}  [cyan]{output_dir}[/]\n"
            f"[dim]Method:[/] [magenta]{method_label}[/]"
        )

        # Last 12 log lines
        recent = log_lines[-12:] if log_lines else ["Initializing..."]
        log_text = "\n".join(f"[dim]>[/] {line}" for line in recent)

        return Panel(
            f"{header}\n\n{table}\n\n[dim]{log_rule}[/]\n{log_text}",
            border_style="green",
            title=panel_title,
        )

    def on_stage(result):
        stage_status[result.stage] = result.status
        stage_msgs[result.stage] = result.message
        if live:
            live.update(make_display())

    def on_log(msg):
        log_lines.append(msg)
        if live:
            live.update(make_display())

    live = None
    prompt_kwargs = {}
    residue_meta = None
    if getattr(args, "prompt_pairs_file", None):
        from obliteratus.prompts import load_prompt_pairs_file

        try:
            harmful, harmless = load_prompt_pairs_file(args.prompt_pairs_file)
        except ValueError as exc:
            console.print(f"[red]Invalid --prompt-pairs-file:[/] {exc}")
            raise SystemExit(2) from exc
        prompt_kwargs = {"harmful_prompts": harmful, "harmless_prompts": harmless}
        log_lines.append(
            f"Loaded explicit prompt pairs: {len(harmful)} harmful + {len(harmless)} harmless."
        )
    elif getattr(args, "residue_file", None):
        from obliteratus.hard_negative import build_weighted_prompt_pairs

        harmful, harmless, residue_meta = build_weighted_prompt_pairs(
            base_dataset=getattr(args, "dataset", "builtin"),
            residue_files=args.residue_file,
            residue_weight=getattr(args, "residue_weight", 5),
            max_residue=getattr(args, "residue_max", None),
        )
        prompt_kwargs = {"harmful_prompts": harmful, "harmless_prompts": harmless}
        log_lines.append(
            "Loaded hard-negative residue: "
            f"{residue_meta['residue_examples']} unique prompts, "
            f"+{residue_meta['residue_added_pairs']} weighted pairs "
            f"({residue_meta['total_pairs']} total)."
        )

    pipeline = AbliterationPipeline(
        model_name=model_name,
        output_dir=output_dir,
        device=args.device,
        dtype=args.dtype,
        method=method,
        n_directions=args.n_directions,
        direction_method=getattr(args, "direction_method", None),
        regularization=args.regularization,
        refinement_passes=args.refinement_passes,
        min_layer_fraction=getattr(args, "min_layer_fraction", None),
        max_layer_fraction=getattr(args, "max_layer_fraction", None),
        harmless_pc_count=getattr(args, "harmless_pc_count", None),
        shield_concept_count=getattr(args, "shield_concept_count", None),
        shield_ridge=getattr(args, "shield_ridge", None),
        shield_residualize=getattr(args, "shield_residualize", None),
        shield_layer_penalty=getattr(args, "shield_layer_penalty", None),
        projection_target=getattr(args, "projection_target", None),
        projection_row_fraction=getattr(args, "projection_row_fraction", None),
        quantization=args.quantization,
        gpu_memory_utilization=getattr(args, "gpu_memory_utilization", None),
        trust_remote_code=getattr(args, "trust_remote_code", False),
        large_model_mode=getattr(args, "large_model", False),
        verify_sample_size=getattr(args, "verify_sample_size", None),
        refusal_max_tokens=getattr(args, "refusal_max_tokens", None),
        max_perplexity_increase=getattr(args, "max_perplexity_increase", 3.0),
        min_coherence_retention=getattr(args, "min_coherence_retention", 0.5),
        max_degenerate_fraction=getattr(args, "max_degenerate_fraction", 0.2),
        on_stage=on_stage,
        on_log=on_log,
        **prompt_kwargs,
    )

    with Live(make_display(), console=console, refresh_per_second=4) as live_ctx:
        live = live_ctx
        try:
            result_path = pipeline.run()
            if residue_meta:
                import json
                Path(result_path, "hard_negative_residue.json").write_text(
                    json.dumps(residue_meta, indent=2)
                )
            live.update(make_display())
        except Exception as e:
            log_lines.append(f"[red]ERROR: {e}[/]")
            live.update(make_display())
            raise

    # ── Telemetry: send pipeline report to community leaderboard ──
    try:
        from obliteratus.telemetry import maybe_send_pipeline_report
        maybe_send_pipeline_report(pipeline)
    except Exception:
        pass  # Telemetry is best-effort

    # ── Community contribution (--contribute flag) ──
    contrib_path = None
    if getattr(args, "contribute", False):
        try:
            from obliteratus.community import save_contribution
            contrib_path = save_contribution(
                pipeline,
                model_name=model_name,
                notes=getattr(args, "contribute_notes", ""),
            )
        except Exception as e:
            console.print(f"[yellow]Could not save contribution: {e}[/yellow]")

    console.print()
    contrib_line = ""
    if contrib_path:
        contrib_line = f"\n  Contribution: [cyan]{contrib_path}[/]"
    console.print(
        Panel(
            f"[bold green]Abliteration complete![/]\n\n"
            f"  Model saved to: [cyan]{result_path}[/]\n"
            f"  Metadata: [cyan]{result_path}/abliteration_metadata.json[/]"
            f"{contrib_line}\n\n"
            f"  [dim]Load with:[/] AutoModelForCausalLM.from_pretrained('{result_path}')",
            border_style="green",
            title=_console_text(
                "[bold green]✓ REBIRTH COMPLETE[/]",
                "[bold green]REBIRTH COMPLETE[/]",
            ),
        )
    )


def _cmd_gpu_calc(args):
    import math

    from rich.panel import Panel
    from rich.table import Table

    BYTES_PER_PARAM = {
        "float32": 4,
        "float16": 2,
        "bfloat16": 2,
        "int8": 1,
        "int4": 0.5,
    }

    # Resolve param counts
    total_params_b = args.params
    active_params_b = args.active_params

    if total_params_b is None:
        if args.model is None:
            console.print("[red]Provide either a model name or --params.[/]")
            raise SystemExit(1)
        console.print(f"Profiling [cyan]{args.model}[/]...")
        try:
            from obliteratus.model_profile import profile_model
            profile = profile_model(args.model, dtype=args.dtype)
            total_params_b = profile.total_params_b
            if active_params_b is None:
                active_params_b = profile.active_params_b
        except Exception as e:
            console.print(f"[red]Could not profile model: {e}[/]")
            raise SystemExit(1)
        if total_params_b is None:
            console.print("[red]Cannot determine model size. Use --params.[/]")
            raise SystemExit(1)

    if active_params_b is None:
        active_params_b = total_params_b

    bpp = BYTES_PER_PARAM[args.dtype]
    gpu_mem_gb = args.gpu_mem

    # Model weight memory (use base-10 GB to match HF/nvidia conventions)
    weight_gb = total_params_b * bpp

    # Activation overhead during forward passes (PROBE/VERIFY).
    # Scales with active params, not total. Empirical from benchmarks:
    # - DeepSeek-70B (149GB): failed at 160GB (2 GPUs), OK at 240GB (3 GPUs)
    # - GPT-OSS-120B (234GB): failed at 240GB (3 GPUs), OK at 320GB (4 GPUs)
    # This implies ~15-35% overhead. We use 20% as a reasonable middle ground.
    active_weight_gb = active_params_b * bpp
    activation_overhead_gb = active_weight_gb * 0.20

    # CUDA context + fragmentation overhead: ~1.5 GB per GPU (fixed cost)
    cuda_overhead_per_gpu = 1.5

    # Total memory needed (before splitting across GPUs)
    total_needed_gb = weight_gb + activation_overhead_gb

    # Find minimum GPUs: we need total_needed / (gpu_mem - cuda_overhead) GPUs
    usable_per_gpu = gpu_mem_gb - cuda_overhead_per_gpu
    if usable_per_gpu <= 0:
        console.print("[red]GPU memory too small after CUDA overhead.[/]")
        raise SystemExit(1)

    min_gpus = math.ceil(total_needed_gb / usable_per_gpu)
    min_gpus = max(min_gpus, 1)

    # Show results for a range of GPU counts
    is_moe = active_params_b < total_params_b * 0.99

    table = Table(title="GPU Configurations", show_edge=True)
    table.add_column("GPUs", justify="right", style="cyan")
    table.add_column("VRAM/GPU", justify="right")
    table.add_column("Total VRAM", justify="right")
    table.add_column("Headroom", justify="right")
    table.add_column("Verdict", min_width=20)

    # Show from min_gpus-1 (to show why it fails) up to 8
    low = max(1, min_gpus - 1)
    high = max(min_gpus + 3, 8)
    for n in range(low, high + 1):
        total_vram = n * gpu_mem_gb
        usable_vram = n * usable_per_gpu
        headroom = usable_vram - total_needed_gb
        headroom_pct = headroom / total_needed_gb * 100
        vram_per = total_needed_gb / n

        if headroom < 0:
            verdict = "[red]INSUFFICIENT[/]"
        elif headroom_pct < 15:
            verdict = "[yellow]TIGHT — may fail[/]"
        elif n == min_gpus:
            verdict = "[bold green]MINIMUM (recommended)[/]"
        else:
            verdict = "[green]OK[/] [dim](more GPUs = slower)[/]"

        table.add_row(
            str(n),
            f"{vram_per:.1f} GB",
            f"{total_vram:.0f} GB",
            f"{headroom:+.1f} GB ({headroom_pct:+.0f}%)",
            verdict,
        )

    model_label = args.model or f"{total_params_b:.1f}B params"
    moe_line = ""
    if is_moe:
        moe_line = f"\n  Active params:  [cyan]{active_params_b:.1f}B[/] ({active_params_b/total_params_b*100:.0f}% of total — MoE)"

    console.print(Panel(
        f"  Model:          [cyan]{model_label}[/]\n"
        f"  Total params:   [cyan]{total_params_b:.1f}B[/]"
        f"{moe_line}\n"
        f"  Dtype:          [cyan]{args.dtype}[/] ({bpp} bytes/param)\n"
        f"  Weight memory:  [cyan]{weight_gb:.1f} GB[/]\n"
        f"  Activation est: [cyan]{activation_overhead_gb:.1f} GB[/]\n"
        f"  Total needed:   [bold]{total_needed_gb:.1f} GB[/]\n"
        f"  GPU VRAM:       [cyan]{gpu_mem_gb:.0f} GB[/] per device",
        title="[bold]GPU Calculator[/]",
        border_style="cyan",
    ))
    console.print(table)
    console.print(
        f"\n  [bold green]Minimum GPUs: {min_gpus}[/]"
        f"  ({min_gpus} x {gpu_mem_gb:.0f} GB = {min_gpus * gpu_mem_gb:.0f} GB)\n"
    )
    console.print(
        "[dim]Note: fewer GPUs = faster (pipeline parallel has cross-device overhead).\n"
        "Estimates are conservative. Actual memory may vary with sequence length\n"
        "and model architecture. See 'obliteratus obliterate --help' for runtime options.[/]\n"
    )


def _estimate_total_params_b(config) -> float:
    """Estimate total parameter count in billions from a HuggingFace config."""
    # Some configs have explicit param counts
    for attr in ("num_parameters", "n_params"):
        val = getattr(config, attr, None)
        if val and val > 1000:
            return val / 1e9

    # Estimate from architecture dimensions
    h = getattr(config, "hidden_size", 0)
    L = getattr(config, "num_hidden_layers", 0)
    V = getattr(config, "vocab_size", 0)
    i = getattr(config, "intermediate_size", h * 4)

    if h == 0 or L == 0:
        console.print("[red]Cannot determine model size from config. Use --params.[/]")
        raise SystemExit(1)

    n_heads = getattr(config, "num_attention_heads", None) or (h // 128)
    head_dim = getattr(config, "head_dim", None) or (h // n_heads if n_heads else 128)
    kv_heads = getattr(config, "num_key_value_heads", None) or n_heads

    # Attention: Q + K + V projections + output projection
    attn_params = h * (n_heads * head_dim) + h * (kv_heads * head_dim) * 2 + (n_heads * head_dim) * h

    # FFN (MoE or dense)
    n_experts = getattr(config, "num_local_experts", getattr(config, "num_experts", 1)) or 1
    # MoE models often have a separate intermediate size for expert FFNs
    moe_i = getattr(config, "moe_intermediate_size", i)
    # gate + up + down projections per expert
    ffn_per_expert = h * moe_i * 3
    ffn_params = ffn_per_expert * n_experts
    # Some architectures (Qwen, DeepSeek) also have a shared/dense FFN per layer
    if n_experts > 1 and hasattr(config, "moe_intermediate_size"):
        # The dense FFN uses the main intermediate_size
        ffn_params += h * i * 3
    # Router
    if n_experts > 1:
        ffn_params += h * n_experts

    # Per-layer: attention + FFN + layernorms
    layer_params = attn_params + ffn_params + h * 4  # 2 layernorms, 2 params each

    # Embedding + LM head
    embed_params = V * h * 2  # input + output embeddings (may be tied but counts for memory)

    total = L * layer_params + embed_params
    return total / 1e9


def _estimate_active_params_b(config, total_params_b: float) -> float:
    """For MoE models, estimate active parameters per forward pass."""
    n_experts = getattr(config, "num_local_experts", getattr(config, "num_experts", 1)) or 1
    if n_experts <= 1:
        return total_params_b

    top_k = getattr(config, "num_experts_per_tok", getattr(config, "top_k", 2)) or 2

    h = getattr(config, "hidden_size", 0)
    i = getattr(config, "intermediate_size", h * 4)
    moe_i = getattr(config, "moe_intermediate_size", i)
    L = getattr(config, "num_hidden_layers", 0)

    # FFN per expert (uses moe_intermediate_size if available)
    ffn_per_expert = h * moe_i * 3
    # Active FFN = top_k experts instead of all n_experts
    ffn_all = ffn_per_expert * n_experts * L
    ffn_active = ffn_per_expert * top_k * L
    # Non-FFN params (includes any shared/dense FFN)
    non_ffn = total_params_b * 1e9 - ffn_all
    active = non_ffn + ffn_active
    return max(active / 1e9, 0.1)


def _make_remote_runner(args):
    """Create a RemoteRunner from CLI --remote flags."""
    from obliteratus.remote import RemoteConfig, RemoteRunner

    rc = RemoteConfig.from_cli_args(
        args.remote,
        port=args.ssh_port,
        ssh_key=args.ssh_key,
        remote_dir=args.remote_dir,
        python=args.remote_python,
        sync_results=not args.no_sync,
        gpus=getattr(args, "gpus", None),
    )
    return RemoteRunner(rc)


def _cmd_remote_abliterate(args):
    from rich.panel import Panel

    runner = _make_remote_runner(args)

    kwargs = {}
    if args.method:
        kwargs["method"] = args.method
    if args.device:
        kwargs["device"] = args.device
    if args.dtype:
        kwargs["dtype"] = args.dtype
    if args.quantization:
        kwargs["quantization"] = args.quantization
    if getattr(args, "gpu_memory_utilization", None) is not None:
        kwargs["gpu_memory_utilization"] = args.gpu_memory_utilization
    if args.n_directions is not None:
        kwargs["n_directions"] = args.n_directions
    if getattr(args, "direction_method", None):
        kwargs["direction_method"] = args.direction_method
    if args.regularization is not None:
        kwargs["regularization"] = args.regularization
    if args.refinement_passes is not None:
        kwargs["refinement_passes"] = args.refinement_passes
    if getattr(args, "min_layer_fraction", None) is not None:
        kwargs["min_layer_fraction"] = args.min_layer_fraction
    if getattr(args, "max_layer_fraction", None) is not None:
        kwargs["max_layer_fraction"] = args.max_layer_fraction
    if getattr(args, "harmless_pc_count", None) is not None:
        kwargs["harmless_pc_count"] = args.harmless_pc_count
    if getattr(args, "shield_concept_count", None) is not None:
        kwargs["shield_concept_count"] = args.shield_concept_count
    if getattr(args, "shield_ridge", None) is not None:
        kwargs["shield_ridge"] = args.shield_ridge
    if getattr(args, "shield_residualize", None) is not None:
        kwargs["shield_residualize"] = args.shield_residualize
    if getattr(args, "shield_layer_penalty", None) is not None:
        kwargs["shield_layer_penalty"] = args.shield_layer_penalty
    if getattr(args, "projection_target", None) is not None:
        kwargs["projection_target"] = args.projection_target
    if getattr(args, "projection_row_fraction", None) is not None:
        kwargs["projection_row_fraction"] = args.projection_row_fraction
    if getattr(args, "large_model", False):
        kwargs["large_model"] = True
    if getattr(args, "verify_sample_size", None) is not None:
        kwargs["verify_sample_size"] = args.verify_sample_size
    if getattr(args, "refusal_max_tokens", None) is not None:
        kwargs["refusal_max_tokens"] = args.refusal_max_tokens

    result_path = runner.run_obliterate(
        model=args.model,
        local_output_dir=args.output_dir,
        **kwargs,
    )

    if result_path:
        console.print(
            Panel(
                f"[bold green]Remote abliteration complete![/]\n\n"
                f"  Results at: [cyan]{result_path}[/]\n\n"
                f"  [dim]Load with:[/] AutoModelForCausalLM.from_pretrained('{result_path}')",
                border_style="green",
                title="[bold green]REBIRTH COMPLETE (remote)[/]",
            )
        )
    else:
        console.print("[red]Remote abliteration failed. Check logs above.[/]")
        raise SystemExit(1)


def _cmd_remote_run(args):
    runner = _make_remote_runner(args)
    result_path = runner.run_config(
        local_config_path=args.config,
        local_output_dir=args.output_dir,
        preset=args.preset,
    )
    if result_path:
        console.print(f"\n[bold green]Remote run complete.[/] Results at: [cyan]{result_path}[/]")
    else:
        console.print("[red]Remote run failed. Check logs above.[/]")
        raise SystemExit(1)


def _cmd_remote_tourney(args):
    from rich.panel import Panel

    runner = _make_remote_runner(args)
    result_path = runner.run_tourney(
        model=args.model,
        local_output_dir=args.output_dir,
        device=args.device,
        dtype=args.dtype,
        quantization=args.quantization,
        methods=args.methods,
        hub_org=args.hub_org,
        hub_repo=args.hub_repo,
        dataset=args.dataset,
    )
    if result_path:
        console.print(
            Panel(
                f"[bold green]Remote tournament complete![/]\n\n"
                f"  Results at: [cyan]{result_path}[/]",
                border_style="green",
                title="[bold green]TOURNAMENT COMPLETE (remote)[/]",
            )
        )
    else:
        console.print("[red]Remote tournament failed. Check logs above.[/]")
        raise SystemExit(1)


if __name__ == "__main__":
    main()
