backend/app/models/base_model.py

"""Abstract base detector interface.

Every detector returns a ``DetectorOutput`` regardless of paradigm. The router
then maps that into the public ``ModelResult`` schema.
"""
from __future__ import annotations

import abc
from dataclasses import dataclass, field
from typing import Any, Dict, Optional

import torch


@dataclass
class DetectorOutput:
    prediction: str                  # 'real' | 'fake'
    confidence: float                # in [0, 1] for the predicted label
    inference_time_ms: float
    features: Dict[str, Any] = field(default_factory=dict)
    notes: Optional[str] = None
    fallback_to: Optional[str] = None  # model_id that actually produced the result, if fallback


class FeatureCache:
    """Per-request shared feature cache. Avoids recomputing XLS-R / mel
    when multiple detectors need the same backbone features.

    Detectors that benefit:
        * Nes2Net + SONAR  → both consume XLS-R 1024-dim features
        * BiCrossMamba-ST + VoiceRadar  → both consume mel spectrogram
    """

    def __init__(self) -> None:
        self._xlsr = None
        self._mel = None

    def get_xlsr(self, waveform, sample_rate: int = 16000):
        if self._xlsr is None:
            from app.features.ssl_extractor import XLSRExtractor

            self._xlsr = XLSRExtractor.get().extract(waveform, sample_rate=sample_rate)
        return self._xlsr

    def get_mel(self, waveform, sample_rate: int = 16000):
        if self._mel is None:
            from app.features.spectrogram import mel_spectrogram

            self._mel = mel_spectrogram(waveform, sample_rate=sample_rate)
        return self._mel


class BaseDetector(abc.ABC):
    """Common interface for all detectors."""

    model_id: str = "abstract"
    display_name: str = "Abstract Detector"
    paradigm: str = "abstract"
    family: str = "abstract"        # 'production' | 'paper_architecture'
    description: str = ""
    eer_reference: Dict[str, Optional[float]] = {}
    params_k: Optional[int] = None
    backend_params_k: Optional[int] = None

    def __init__(self) -> None:
        self._loaded = False
        self._status = "not_loaded"

    # ---- lifecycle ----
    def warm_up(self) -> None:
        """Idempotent: load weights & set ``_loaded`` / ``_status``."""
        if self._loaded:
            return
        try:
            self._load()
            self._loaded = True
            self._status = "live"
        except Exception as exc:  # noqa: BLE001
            self._status = f"error: {type(exc).__name__}"
            raise

    @abc.abstractmethod
    def _load(self) -> None: ...

    @abc.abstractmethod
    def predict(self, waveform: torch.Tensor, sample_rate: int = 16000,
                cache: Optional["FeatureCache"] = None) -> DetectorOutput: ...

    # ---- introspection ----
    @property
    def status(self) -> str:
        return self._status

    @property
    def is_loaded(self) -> bool:
        return self._loaded