"""Drives eval-unlearn's registry directly (bypassing SingleBenchmarkRunner /
MultiBenchmarkRunner) so that:

  1. The free_run diffusion pipeline is loaded once per submission and reused
     across all metrics (load_technique + run_one_metric), rather than
     reloaded per metric. On this Space's dedicated T4 hardware, reloading it
     6x per submission churned enough GPU memory across CUDA allocator
     fragmentation to cause an out-of-memory crash a few metrics in — a hard
     crash, not a catchable Python exception, which killed the whole process
     and looked like the app "randomly refreshing" mid-submission. (This
     traded off against ZeroGPU's earlier per-call duration ceiling, which
     favored short-lived, independent calls — if this Space ever moves back
     to ZeroGPU, that tradeoff needs revisiting: a technique object loaded in
     one @spaces.GPU call may not be valid to reuse in another.)
  2. Generated images are never written to disk — technique.generate() ->
     metric.update() -> metric.compute() happens in memory and the images are
     discarded, so no i2p-prompted imagery is ever persisted or published.

"cheap" tier caps (num_inference_steps and each metric's `limit`) trade score
precision for speed; the eval-unlearn CLI run with full dataset sizes remains
the source of truth for publication-quality numbers.
"""

import os
import sys
import types
from pathlib import Path
from typing import Any, Dict, List, Optional

import spaces


def _stub_missing_technique_packages() -> None:
    """Work around a bug in eval-unlearn's own packaging.

    eval_unlearn/techniques/__init__.py unconditionally imports the uce,
    concept_steerers and mace wrapper modules (unlike esd/saeuron/safree/etc,
    which are correctly wrapped in try/except there). Each of those three
    unconditionally does `from uce import UCEPipeline` for an external
    package that was never published alongside eval-unlearn on PyPI. Since
    Python must run that __init__.py before importing *any* technique
    submodule, one missing unrelated package breaks free_run's registration
    too. We never use uce/concept_steerers/mace, so a harmless stub module is
    enough to let the parent package finish initializing.
    """
    stub_attrs = {
        "uce": "UCEPipeline",
        "concept_steerers": "ConceptSteeringPipeline",
        "mace": "MACEPipeline",
    }
    for name, attr in stub_attrs.items():
        if name in sys.modules:
            continue
        try:
            __import__(name)
        except ImportError:
            stub = types.ModuleType(name)
            setattr(stub, attr, object)
            sys.modules[name] = stub


_stub_missing_technique_packages()

DATA_DIR = Path(__file__).resolve().parent.parent.parent / "data" / "ua_ira"

# The 7 concepts I2P has prompt categories for (datasets/_i2p_categories.py in
# eval-unlearn). ASR-I2P is only ever non-N/A for these; "other" always N/A's.
I2P_CONCEPTS = [
    "nudity",
    "harassment",
    "hate",
    "illegal activity",
    "self-harm",
    "shocking",
    "violence",
]
CONCEPTS = I2P_CONCEPTS + ["other"]

# Metrics offered on the online "cheap" tier, in run order (cheapest/most
# broadly-applicable first so a submitter sees partial progress quickly).
METRIC_NAMES = ["clip_score", "asr_i2p", "ua_ira", "err", "tifa", "fid"]

DEFAULT_NUM_INFERENCE_STEPS = 25

# Concepts eval-unlearn ships curated UA/IRA target/retain CSVs for. UA/IRA
# can't be derived from a concept name alone (it needs a curated retain set),
# so it's only offered for concepts we actually have CSVs for.
UA_IRA_CONCEPTS = {"nudity", "violence"}

# What the bundled retain CSV for each concept actually depicts, used as the
# CLIP retain-concept label (see examples/data/{concept}_retain_prompts.csv).
_UA_IRA_RETAIN_LABELS = {"nudity": "clothed person", "violence": "peaceful scene"}

# Default number of images/prompts each metric uses on the "cheap" tier, and
# the largest size the Submit tab's sliders allow. Sizes above the default
# trade speed for precision — going too high risks exceeding a single
# ZeroGPU call's time budget (run_one_metric's duration=180s).
DEFAULT_SIZES: Dict[str, int] = {
    "clip_score": 25,
    "fid": 50,
    "tifa": 20,
    "asr_i2p": 15,
    "err": 15,  # applied to each of target/retain/adversarial
    "ua_ira": 11,  # the bundled curated CSVs only have 11 rows each
}
MAX_SIZES: Dict[str, int] = {
    "clip_score": 100,
    "fid": 150,
    "tifa": 60,
    "asr_i2p": 60,
    "err": 40,
    "ua_ira": 11,
}


def available_metrics_for_concept(concept: str) -> List[str]:
    """Which of METRIC_NAMES can actually run for this concept.

    err's forgetting check is hardcoded to nudity I2P prompts regardless of
    the technique's real concept (eval_unlearn.datasets.err_composite), so we
    only offer it for nudity submissions to avoid a misleading score. ua_ira
    needs curated retain prompts we only have for nudity/violence.
    """
    metrics = list(METRIC_NAMES)
    if concept != "nudity":
        metrics.remove("err")
    if concept not in UA_IRA_CONCEPTS:
        metrics.remove("ua_ira")
    return metrics


def _build_metric_config(
    metric_name: str, concept: str, size: Optional[int] = None
) -> Dict[str, Any]:
    """Build a metric's config dict, applying either its default size or a
    caller-supplied override (clamped to MAX_SIZES so the Submit tab's
    sliders can't push a metric past what a single ZeroGPU call can finish).
    """
    n = size or DEFAULT_SIZES.get(metric_name)
    if n is not None:
        n = max(1, min(n, MAX_SIZES.get(metric_name, n)))

    config: Dict[str, Any] = {}

    if metric_name in ("clip_score", "fid", "tifa"):
        config["limit"] = n
    elif metric_name == "asr_i2p":
        config["limit"] = n
        config["concept_name"] = concept
        # asr_i2p's "auto" detector picks NudeNet for nudity and Q16 for
        # everything else — Q16 is an unpublished local package this Space
        # doesn't install. Force NudeNet for nudity (installed via the `asr`
        # extra) and fall back to the CLIP-similarity detector (needs only
        # `transformers`, already a dependency) for the other 6 concepts.
        config["detector"] = "nudenet" if concept == "nudity" else "clip"
    elif metric_name == "err":
        config.update({"target_limit": n, "retain_limit": n, "adversarial_limit": n})
    elif metric_name == "ua_ira":
        config.update(
            {
                "target_prompts_path": str(DATA_DIR / f"{concept}_target_prompts.csv"),
                "retain_prompts_path": str(DATA_DIR / f"{concept}_retain_prompts.csv"),
                "target_concept": concept,
                "retain_concept": _UA_IRA_RETAIN_LABELS.get(concept, "retain"),
                "target_prompt_limit": n,
                "retain_prompt_limit": n,
            }
        )

    return config


def _metric_result_to_dict(result) -> Dict[str, Any]:
    return {"name": result.name, "value": result.value, "details": result.details}


@spaces.GPU(duration=60)
def load_technique(model_id: str, num_inference_steps: int = DEFAULT_NUM_INFERENCE_STEPS):
    """Load the free_run diffusion pipeline once for the whole submission.

    Raises on failure (bad model_id, load error) — the caller should abort
    the submission rather than retrying this once per metric.
    """
    from eval_unlearn.registry import get_technique
    from eval_unlearn.registry.entrypoints import load_entrypoints

    load_entrypoints()
    technique_factory = get_technique("free_run")
    return technique_factory(model_id=model_id, num_inference_steps=num_inference_steps)


def free_gpu_memory() -> None:
    """Call after the caller has `del`'d its own references (e.g. the
    technique) to actually release that memory back to the CUDA allocator.
    """
    try:
        import torch

        torch.cuda.empty_cache()
    except Exception:
        pass


@spaces.GPU(duration=180)
def run_one_metric(
    technique,
    metric_name: str,
    concept: str,
    size: Optional[int] = None,
    seed: Optional[int] = 0,
) -> Dict[str, Any]:
    """Run a single metric against an already-loaded free_run technique.

    `size` overrides that metric's default number of images/prompts (see
    DEFAULT_SIZES/MAX_SIZES); leave it None to use the default "cheap" tier
    size.

    Returns a plain dict (not MetricResult) — always, even on failure, with an
    "error" key set instead of raising, so a caller looping over several
    metrics can record partial results instead of aborting the whole
    submission when one metric fails.
    """
    from eval_unlearn.registry import get_metric
    from eval_unlearn.registry.entrypoints import load_entrypoints

    load_entrypoints()

    metric_config = _build_metric_config(metric_name, concept, size)

    metric = None
    try:
        metric_factory = get_metric(metric_name)
        metric = metric_factory(**metric_config)

        loader = metric.load_dataset()

        total = 0
        for batch in loader:
            images = technique.generate(prompts=batch.prompts, seed=seed)
            metric.update(images, batch.prompts, batch.metadata)
            total += len(images)
            del images

        result = metric.compute()
        out = _metric_result_to_dict(result)
        out["details"] = {**out.get("details", {}), "images_generated": total}
        return out

    except Exception as e:  # noqa: BLE001 - reported to the caller, not raised
        return {"name": metric_name, "value": None, "details": {}, "error": str(e)}

    finally:
        del metric
        free_gpu_memory()