"""Inference backend boundary shared by one-shot and server generation."""

from __future__ import annotations

from pathlib import Path
from typing import Protocol

import torch

from diffusion_cli.config import GenerationConfig, ModelProfile, ModelSources
from diffusion_cli.errors import DiffusionCliError


class InferenceBackend(Protocol):
    """Generate NCHW RGB image tensors for one local model profile."""

    @property
    def modelId(self) -> str:
        """Return the stable identifier exposed to clients."""

    def generate(self, config: GenerationConfig) -> torch.Tensor:
        """Generate RGB image tensors in the range zero through one."""


def _stageError(profile: ModelProfile | None, stage: str, error: Exception):
    """Turn a CUDA OOM into an actionable stage-specific CLI error."""

    if isinstance(error, DiffusionCliError):
        raise error
    if isinstance(error, RuntimeError) and "out of memory" in str(error).lower():
        name = profile.name if profile is not None else "local"
        raise DiffusionCliError(
            f"CUDA out of memory in {stage} for model profile {name}"
        ) from error
    raise error


class ZImageBackend:
    """Inference backend retaining the existing Z-Image generation path."""

    def __init__(
        self,
        profile: ModelProfile | None,
        sources: ModelSources,
        residency: str,
        *,
        text_encoder_factory=None,
        model_factory=None,
        vae_factory=None,
        sampler=None,
        component_dtype=None,
    ) -> None:
        from diffusion_cli.text_encoder import ZImageTextEncoder
        from diffusion_cli.vae import ZImageVae
        from diffusion_cli.zimage_model import ZImageModel
        from diffusion_cli.sampling import sampleLatents

        self.profile = profile
        self.sources = sources
        self.residency = residency
        self.text_encoder_factory = text_encoder_factory or ZImageTextEncoder
        self.model_factory = model_factory or ZImageModel
        self.vae_factory = vae_factory or ZImageVae
        self.sampler = sampler or sampleLatents
        self.component_dtype = component_dtype
        self._text_encoder = None
        self._model = None
        self._vae = None

    @property
    def modelId(self) -> str:
        """Return the configured profile name or legacy identity."""

        return self.profile.name if self.profile is not None else "z-image-local"

    def generate(self, config: GenerationConfig):
        """Generate RGB tensors using staged or CPU-cached Z-Image residency."""

        if self.residency == "cpu-cache":
            return self._generateCpuCache(config)
        return self._generateStaged(config)

    def _dtype(self, config, source):
        if self.component_dtype is not None:
            return self.component_dtype(config, source)
        if config.dtype_name != "auto":
            return config.dtype
        from diffusion_cli.checkpoint import inspectSourceTorchDtype

        return inspectSourceTorchDtype(source) or config.dtype

    def _generateStaged(self, config):
        text_dtype = self._dtype(config, self.sources.text_encoder)
        diffusion_dtype = self._dtype(config, self.sources.diffusion_model)
        vae_dtype = self._dtype(config, self.sources.vae)
        try:
            text_encoder = self.text_encoder_factory(
                self.sources.text_encoder,
                config.tokenizer_path,
                config.device,
                text_dtype,
            )
            conditioning = text_encoder.encodePrompts(
                config.prompt,
                config.negative_prompt,
            )
            del text_encoder
            _releaseMemory()
        except Exception as exc:
            _stageError(self.profile, "text encoder", exc)

        try:
            model = self.model_factory(
                self.sources.diffusion_model,
                config.device,
                diffusion_dtype,
            )
            latent = self.sampler(
                model,
                conditioning,
                batch_size=config.batch_size,
                height=config.height,
                width=config.width,
                seed=config.seed,
                steps=config.steps,
                cfg=config.cfg,
                device=config.device,
                dtype=diffusion_dtype,
            )
            del model
            _releaseMemory()
        except Exception as exc:
            _stageError(self.profile, "diffusion model", exc)

        try:
            vae = self.vae_factory(
                self.sources.vae,
                config.device,
                vae_dtype,
            )
            images = vae.decode(latent)
            del vae
            _releaseMemory()
            return images
        except Exception as exc:
            _stageError(self.profile, "VAE", exc)

    def _generateCpuCache(self, config):
        cpu = torch.device("cpu")
        text_dtype = self._dtype(config, self.sources.text_encoder)
        diffusion_dtype = self._dtype(config, self.sources.diffusion_model)
        vae_dtype = self._dtype(config, self.sources.vae)
        try:
            if self._text_encoder is None:
                self._text_encoder = self.text_encoder_factory(
                    self.sources.text_encoder,
                    config.tokenizer_path,
                    cpu,
                    text_dtype,
                )
            self._text_encoder.toDevice(config.device, text_dtype)
            try:
                conditioning = self._text_encoder.encodePrompts(
                    config.prompt,
                    config.negative_prompt,
                )
            finally:
                self._text_encoder.toCpu()
                _releaseMemory()
        except Exception as exc:
            _stageError(self.profile, "text encoder", exc)

        try:
            if self._model is None:
                self._model = self.model_factory(
                    self.sources.diffusion_model,
                    cpu,
                    diffusion_dtype,
                )
            self._model.toDevice(config.device, diffusion_dtype)
            try:
                latent = self.sampler(
                    self._model,
                    conditioning,
                    batch_size=config.batch_size,
                    height=config.height,
                    width=config.width,
                    seed=config.seed,
                    steps=config.steps,
                    cfg=config.cfg,
                    device=config.device,
                    dtype=diffusion_dtype,
                )
            finally:
                self._model.toCpu()
                _releaseMemory()
        except Exception as exc:
            _stageError(self.profile, "diffusion model", exc)

        try:
            if self._vae is None:
                self._vae = self.vae_factory(
                    self.sources.vae,
                    cpu,
                    vae_dtype,
                )
            self._vae.toDevice(config.device, vae_dtype)
            try:
                images = self._vae.decode(latent)
            finally:
                self._vae.toCpu()
                _releaseMemory()
            return images
        except Exception as exc:
            _stageError(self.profile, "VAE", exc)


class Krea2Backend:
    """Local Krea 2 Turbo/Raw backend with model-specific sampling."""

    def __init__(self, profile: ModelProfile, sources: ModelSources, residency: str):
        self.profile = profile
        self.sources = sources
        self.residency = residency
        self._text_encoder = None
        self._model = None
        self._vae = None

    @property
    def modelId(self) -> str:
        """Return the selected Krea profile name."""

        return self.profile.name

    def _constructText(self, source, tokenizer, device, dtype):
        from diffusion_cli.krea2_text_encoder import Krea2TextEncoder

        return Krea2TextEncoder(source, tokenizer, device, dtype)

    def _constructModel(self, source, device, dtype):
        from diffusion_cli.krea2_model import buildKrea2Model

        return buildKrea2Model(source, device=device, dtype=dtype)

    def _constructVae(self, source, device, dtype):
        from diffusion_cli.qwen_image_vae import QwenImageVae

        return QwenImageVae(source, device, dtype)

    def _sample(self, model, positive, negative, config, dtype):
        from diffusion_cli.krea2_sampling import sampleKrea2

        return sampleKrea2(
            model,
            positive,
            negative_conditioning=negative,
            batch_size=config.batch_size,
            height=config.height,
            width=config.width,
            seed=config.seed,
            steps=config.steps,
            cfg=config.cfg,
            device=config.device,
            dtype=dtype,
            mu=config.mu,
            shift_y1=config.shift_y1,
            shift_y2=config.shift_y2,
        )

    def _generateStages(self, config, *, cache: bool):
        import gc

        cpu = torch.device("cpu")
        device = config.device
        dtype = config.dtype
        try:
            if cache and self._text_encoder is not None:
                text_encoder = self._text_encoder
            else:
                text_encoder = self._constructText(
                    self.sources.text_encoder,
                    config.tokenizer_path,
                    cpu if cache else device,
                    dtype,
                )
                if cache:
                    self._text_encoder = text_encoder
            if cache:
                text_encoder.toDevice(device, dtype)
            positive = text_encoder.encodePrompt(config.prompt)
            negative = (
                text_encoder.encodePrompt(config.negative_prompt)
                if config.cfg > 0
                else None
            )
            if cache:
                text_encoder.toCpu()
            else:
                del text_encoder
            gc.collect()
        except Exception as exc:
            _stageError(self.profile, "text encoder", exc)

        try:
            if cache and self._model is not None:
                model = self._model
            else:
                model = self._constructModel(
                    self.sources.diffusion_model,
                    cpu if cache else device,
                    dtype,
                )
                if cache:
                    self._model = model
            if cache:
                model.toDevice(device, dtype)
            latent = self._sample(model, positive, negative, config, dtype)
            if cache:
                model.toCpu()
            else:
                del model
            gc.collect()
        except Exception as exc:
            _stageError(self.profile, "diffusion model", exc)

        try:
            if cache and self._vae is not None:
                vae = self._vae
            else:
                vae = self._constructVae(
                    self.sources.vae,
                    cpu if cache else device,
                    dtype,
                )
                if cache:
                    self._vae = vae
            if cache:
                vae.toDevice(device, dtype)
            images = vae.decode(latent)
            if cache:
                vae.toCpu()
            else:
                del vae
            gc.collect()
            return images
        except Exception as exc:
            _stageError(self.profile, "VAE", exc)

    def generate(self, config: GenerationConfig):
        """Generate RGB images with Krea's flow-matching sampler."""

        return self._generateStages(config, cache=self.residency == "cpu-cache")


def _releaseMemory() -> None:
    """Release Python and CUDA caches between staged components."""

    import gc

    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


BACKEND_FACTORIES = {
    "z-image": ZImageBackend,
    "krea2": Krea2Backend,
}


def createBackend(
    profile: ModelProfile,
    sources: ModelSources,
    residency: str,
    **kwargs,
) -> InferenceBackend:
    """Construct the registered backend for one configured profile."""

    factory = BACKEND_FACTORIES.get(profile.architecture)
    if factory is None:
        raise DiffusionCliError(
            f"Unknown model architecture: {profile.architecture}"
        )
    if factory is ZImageBackend:
        return factory(profile, sources, residency, **kwargs)
    if kwargs:
        raise TypeError("Krea2Backend does not accept factory overrides")
    return factory(profile, sources, residency)