from __future__ import annotations

from typing import Any

from PIL import Image

from config_loader import load_app_config, get_depth_backend_from_app
from model.Animation.animation_loader import (
    UnifiedAnimationConfig,
    AnimationBackend,
    build_animation_predictor,
)
from model.Depth.depth_loader import UnifiedDepthConfig, DepthBackend, build_depth_predictor
from model.Inpaint.inpaint_loader import UnifiedInpaintConfig, InpaintBackend, build_inpaint_predictor
from model.Seg.seg_loader import UnifiedSegConfig, SegBackend, build_seg_predictor

_depth_predictor = None
_depth_backend: DepthBackend | None = None

_seg_cache: dict[str, Any] = {}

_inpaint_cache: dict[str, Any] = {}

_animation_cache: dict[str, Any] = {}


def ensure_depth_predictor() -> None:
    global _depth_predictor, _depth_backend
    if _depth_predictor is not None and _depth_backend is not None:
        return

    app_cfg = load_app_config()
    backend_str = get_depth_backend_from_app(app_cfg)
    try:
        backend = DepthBackend(backend_str)
    except Exception as e:
        raise ValueError(f"config.py 中 depth.backend 不合法: {backend_str}") from e

    _depth_predictor, _depth_backend = build_depth_predictor(UnifiedDepthConfig(backend=backend))


def get_depth_predictor():
    ensure_depth_predictor()
    return _depth_predictor


def get_seg_predictor(model_name: str):
    if model_name == "mask2former_debug":
        model_name = "sam"

    if model_name in _seg_cache:
        return _seg_cache[model_name]

    if model_name == "sam":
        pred, _ = build_seg_predictor(UnifiedSegConfig(backend=SegBackend.SAM))
        _seg_cache[model_name] = pred
        return pred

    if model_name == "mask2former":
        pred, _ = build_seg_predictor(UnifiedSegConfig(backend=SegBackend.MASK2FORMER))
        _seg_cache[model_name] = pred
        return pred

    raise ValueError(f"未知 segment model_name: {model_name}")


def get_inpaint_predictor(model_name: str):
    if model_name in _inpaint_cache:
        return _inpaint_cache[model_name]

    if model_name == "copy":

        def _copy(image: Image.Image, *_args, **_kwargs) -> Image.Image:
            return image.convert("RGB")

        _inpaint_cache[model_name] = _copy
        return _copy

    if model_name == "sdxl_inpaint":
        pred, _ = build_inpaint_predictor(UnifiedInpaintConfig(backend=InpaintBackend.SDXL_INPAINT))
        _inpaint_cache[model_name] = pred
        return pred

    if model_name == "controlnet":
        pred, _ = build_inpaint_predictor(UnifiedInpaintConfig(backend=InpaintBackend.CONTROLNET))
        _inpaint_cache[model_name] = pred
        return pred

    if model_name == "lama":
        pred, _ = build_inpaint_predictor(UnifiedInpaintConfig(backend=InpaintBackend.LAMA))
        _inpaint_cache[model_name] = pred
        return pred

    raise ValueError(f"未知 inpaint model_name: {model_name}")


def get_animation_predictor(model_name: str):
    if model_name in _animation_cache:
        return _animation_cache[model_name]

    if model_name == "animatediff":
        pred, _ = build_animation_predictor(
            UnifiedAnimationConfig(backend=AnimationBackend.ANIMATEDIFF)
        )
        _animation_cache[model_name] = pred
        return pred

    raise ValueError(f"未知 animation model_name: {model_name}")