"""sdcpp_python_load — carga un StableDiffusion (stable-diffusion-cpp-python) con cache global."""

from __future__ import annotations

import sys
import os
from typing import Any

sys.path.insert(0, os.path.dirname(__file__))

from model_ref import ModelRef

# Cache global: (model_key, wtype, n_threads) -> StableDiffusion object
_SD_CACHE: dict[tuple[str, str, int], Any] = {}


def _get_model_key(model: ModelRef) -> str:
    """Retorna la clave de cache para un ModelRef."""
    return model.path if model.path else model.name


def sdcpp_python_load(
    model: ModelRef,
    n_threads: int = -1,
    wtype: str = "default",
    rng_type: str = "cuda",
) -> Any:
    """Carga un StableDiffusion via stable-diffusion-cpp-python con cache global.

    Instancia StableDiffusion con el checkpoint indicado por model.path (o model.name
    si path es None). Los objetos se cachean en memoria por (model_key, wtype, n_threads)
    — una segunda llamada con los mismos parametros retorna la instancia cacheada sin
    recargar el modelo del disco.

    Args:
        model:     Referencia al modelo. model.path se usa si esta presente;
                   si no, model.name se pasa como model_path (ruta local o hub).
        n_threads: Numero de hilos de CPU para inferencia. -1 usa todos los disponibles.
        wtype:     Tipo de pesos / cuantizacion en memoria.
                   Valores: "default" | "f32" | "f16" | "q8_0" | "q5_1" | "q5_0" | "q4_1" | "q4_0".
                   "default" respeta el tipo original del checkpoint.
        rng_type:  Tipo de generador de numeros aleatorios.
                   Valores: "std_default" | "cuda".
                   "cuda" produce resultados compatibles con la implementacion CUDA
                   incluso en CPU.

    Returns:
        Instancia StableDiffusion lista para llamar a generate_image().

    Raises:
        ImportError: Si stable_diffusion_cpp no esta instalado.
                     Instalar con: pip install stable-diffusion-cpp-python
        OSError:     Si el path del modelo no existe o es invalido.
    """
    try:
        from stable_diffusion_cpp import StableDiffusion
    except ImportError as exc:
        raise ImportError(
            "sdcpp_python_load requiere stable-diffusion-cpp-python. "
            "Instalar con: pip install stable-diffusion-cpp-python\n"
            "Para compilar sin CUDA: "
            "CMAKE_ARGS='-DSD_CUDA=OFF' pip install stable-diffusion-cpp-python"
        ) from exc

    model_key = _get_model_key(model)
    cache_key = (model_key, wtype, n_threads)

    if cache_key in _SD_CACHE:
        return _SD_CACHE[cache_key]

    load_path = model.path if model.path else model.name

    sd = StableDiffusion(
        model_path=load_path,
        wtype=wtype,
        n_threads=n_threads,
        rng_type=rng_type,
    )

    _SD_CACHE[cache_key] = sd
    return sd


def _clear_sd_cache() -> None:
    """Limpia el cache global de instancias StableDiffusion (uso interno y tests)."""
    _SD_CACHE.clear()