from __future__ import annotations

import logging
import threading
from dataclasses import replace

import numpy as np

from .config import _SEG_MODELS, override_settings
from .pipeline import MattingPipeline, MattingResult
from .segmentation import make_segmenter
from .settings import PipelineSettings
from .vitmatte_infer import ViTMatteRunner

logger = logging.getLogger("bgfilter.service")

# HTTP-facing seg_model values, in preference order (first = default).
SEG_MODELS: tuple[str, ...] = tuple(_SEG_MODELS)


class PipelineManager:
    """Long-lived holder of the matting models.

    ViTMatte is the matting model shared by both segmentation backends, so it is
    loaded once and reused. Each segmentation backend is lazily loaded once on
    first use and cached by name. Per-request parameters (``screen_color``,
    ``seg_model``) never trigger a model reload -- ``screen_color`` only affects
    this request's chroma/despill, and ``seg_model`` just selects an already-cached
    (or first-time-loaded) segmenter.
    """

    def __init__(self, base_settings: PipelineSettings, preload: bool = True):
        self._base = base_settings
        self._runner: ViTMatteRunner | None = None
        self._segmenters: dict[str, object] = {}
        self._load_lock = threading.Lock()
        if preload:
            if base_settings.model.matting_method == "vitmatte":
                self._get_runner()
            if base_settings.segmentation.enabled:
                self._get_segmenter(base_settings.segmentation.backend)

    @property
    def default_seg_model(self) -> str:
        return self._base.segmentation.backend

    @property
    def device(self) -> str:
        return self._base.model.device

    def _get_runner(self) -> ViTMatteRunner:
        if self._runner is None:
            with self._load_lock:
                if self._runner is None:
                    logger.info(
                        "loading ViTMatte: model=%s device=%s",
                        self._base.model.model_name,
                        self._base.model.device,
                    )
                    self._runner = ViTMatteRunner(self._base.model)
        return self._runner

    def _get_segmenter(self, backend: str):
        seg = self._segmenters.get(backend)
        if seg is None:
            with self._load_lock:
                seg = self._segmenters.get(backend)
                if seg is None:
                    seg_settings = replace(
                        self._base.segmentation,
                        backend=backend,
                        model_name=_SEG_MODELS[backend],
                    )
                    logger.info(
                        "loading segmenter: backend=%s model=%s device=%s",
                        backend,
                        seg_settings.model_name,
                        seg_settings.device,
                    )
                    seg = make_segmenter(seg_settings)
                    self._segmenters[backend] = seg
        return seg

    def process(
        self, rgb: np.ndarray, *, screen_color: str | None, seg_model: str
    ) -> MattingResult:
        if seg_model not in _SEG_MODELS:
            raise ValueError(f"invalid seg_model '{seg_model}'")
        settings = override_settings(
            self._base, screen_color=screen_color, seg_backend=seg_model
        )
        runner = (
            self._get_runner()
            if settings.model.matting_method == "vitmatte"
            else None
        )
        segmenter = (
            self._get_segmenter(seg_model) if settings.segmentation.enabled else None
        )
        pipeline = MattingPipeline(settings, runner=runner, segmenter=segmenter)
        return pipeline.run_rgb(rgb)