"""Unified segmentation backends — classical baseline or a learned model.

Every segmentation in the toolkit (hematoma, tumour, vessels, ventricles, brain,
skull) has a fast, dependency-free **classical** method (intensity/HU thresholds,
vesselness) and can instead use a **learned** model. This module gives one entry
point — ``segment(volume, task, method=...)`` — plus a registry so a user can
plug in their own trained model (e.g. an nnU-Net / MONAI bundle) without changing
call sites:

    from sulcora import seg_backends as SB
    SB.register_model("tumor_enhancing", my_nnunet_fn)   # fn(Volume, **kw) -> mask Volume
    mask = SB.segment(vol, "tumor_enhancing", method="model")

The classical method is always available and is the default; ``method="model"``
raises :class:`BackendUnavailable` with guidance until a model is registered.
Learned models are more robust on real, noisy clinical data; the classical
baselines are the reproducible fallback.
"""
from __future__ import annotations
from .segmentation import BackendUnavailable   # reuse the one exception type

# task -> classical segmenter fn(volume, **kw) -> mask Volume
_CLASSICAL = {}
# task -> registered learned model fn(volume, **kw) -> mask Volume
_MODELS = {}


def _register_classical():
    """Wire the built-in classical segmenters (lazy imports avoid cycles)."""
    from . import hematoma, tumor, vessels, ventricle, ct
    _CLASSICAL.update({
        "hematoma": lambda v, **k: hematoma.segment_hematoma(v, **k),
        "tumor_enhancing": lambda v, **k: tumor.segment_enhancing(v, **k),
        "tumor_flair": lambda v, **k: tumor.segment_flair(v, **k),
        "vessels": lambda v, **k: vessels.vessel_mask(v, **k),
        "ventricles": lambda v, **k: ventricle.segment_ventricles(v, **k),
        "brain": lambda v, **k: ct.brain_mask(v, **k),
        "skull": lambda v, **k: ct.skull_mask(v, **k),
    })


def tasks() -> list:
    if not _CLASSICAL:
        _register_classical()
    return sorted(_CLASSICAL)


def register_model(task: str, fn) -> None:
    """Register a learned segmenter for `task` — ``fn(volume, **kw) -> mask Volume``.

    This is the integration point for a trained model (nnU-Net, MONAI bundle, a
    TotalSegmentator task, or any callable). Once registered, call
    ``segment(volume, task, method="model")``.
    """
    _MODELS[task] = fn


def has_model(task: str) -> bool:
    return task in _MODELS


def use_totalsegmentator(task="brain", totalseg_task="total", roi_subset=None):
    """Register TotalSegmentator as the learned model for `task` (needs the
    ``[seg]`` extra + downloaded weights). Registers a fn that writes the volume
    to NIfTI, runs TotalSegmentator, and returns the (merged) label mask.

    Raises :class:`BackendUnavailable` at call time if TotalSegmentator is not
    installed — the classical ``method='threshold'`` path always remains.
    """
    def _run(volume, **kw):
        try:
            from totalsegmentator.python_api import totalsegmentator
        except Exception as e:  # pragma: no cover - needs the seg extra
            raise BackendUnavailable(
                "TotalSegmentator not installed — `pip install \".[seg]\"` and download "
                f"weights, or use method='threshold'. ({e})")
        import tempfile, os, numpy as np, nibabel as nib
        from .volume import Volume
        with tempfile.TemporaryDirectory() as d:  # pragma: no cover
            ip = os.path.join(d, "in.nii.gz")
            nib.save(nib.Nifti1Image(np.asarray(volume.data), volume.affine), ip)
            seg = totalsegmentator(ip, os.path.join(d, "out"), task=totalseg_task,
                                   roi_subset=roi_subset, quiet=True)
            arr = np.asanyarray(seg.dataobj)
            return Volume((arr > 0).astype(np.uint8), volume.affine.copy())
    register_model(task, _run)
    return task


def segment(volume, task: str, method="threshold", **kw):
    """Segment `task` with the classical baseline (default) or a learned model."""
    if not _CLASSICAL:
        _register_classical()
    if task not in _CLASSICAL:
        raise ValueError(f"Unknown segmentation task {task!r}; choose from {sorted(_CLASSICAL)}.")
    if method == "threshold":
        return _CLASSICAL[task](volume, **kw)
    if method == "model":
        if task in _MODELS:
            return _MODELS[task](volume, **kw)
        raise BackendUnavailable(
            f"No learned model registered for task {task!r}. Register one with "
            f"seg_backends.register_model({task!r}, fn) — fn(volume) -> mask Volume — "
            "or use method='threshold' (the classical baseline).")
    raise ValueError(f"Unknown method {method!r}; use 'threshold' or 'model'.")
