"""HTTP-independent generation service and residency management."""

from __future__ import annotations

from dataclasses import dataclass
import gc
from pathlib import Path
from threading import Lock

from diffusion_cli.checkpoint import inspectSourceTorchDtype
from diffusion_cli.config import (
    GenerationConfig,
    ImageGenerationRequest,
    UserConfig,
    buildGenerationConfigFromRequest,
)
from diffusion_cli.image_io import encodeImages, saveImages
from diffusion_cli.paths import resolveModelSourcesFromConfig
from diffusion_cli.sampling import sampleLatents
from diffusion_cli.text_encoder import ZImageTextEncoder
from diffusion_cli.vae import ZImageVae
from diffusion_cli.zimage_model import ZImageModel

MODEL_RESIDENCY_VALUES = ("staged", "cpu-cache")


@dataclass(frozen=True)
class GeneratedImage:
    """One generated image encoded for transport or file output."""

    data: bytes
    format: str
    seed: int


def releaseMemory() -> None:
    """Release Python and CUDA caches between large model stages."""

    gc.collect()
    try:
        import torch

        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    except ImportError:
        pass


def componentDtype(config: GenerationConfig, source):
    """Return the runtime dtype for one component source."""

    if config.dtype_name != "auto":
        return config.dtype
    return inspectSourceTorchDtype(source) or config.dtype


class GenerationService:
    """Generate images from internal requests with a residency policy."""

    def __init__(
        self,
        user_config: UserConfig,
        *,
        model_residency: str = "staged",
    ) -> None:
        if model_residency not in MODEL_RESIDENCY_VALUES:
            raise ValueError(f"Unknown model residency: {model_residency}")
        self.user_config = user_config
        self.model_residency = model_residency
        self._lock = Lock()
        self._text_encoder: ZImageTextEncoder | None = None
        self._model: ZImageModel | None = None
        self._vae: ZImageVae | None = None

    def generateImages(
        self,
        request: ImageGenerationRequest,
    ) -> list[GeneratedImage]:
        """Generate final-format image bytes for one request."""

        with self._lock:
            config = buildGenerationConfigFromRequest(
                request,
                self.user_config,
            )
            images = self._generateTensor(config)
            encoded_images = encodeImages(
                images,
                config.output_extension,
                config.output_quality,
            )
            return [
                GeneratedImage(
                    data=data,
                    format=config.output_extension,
                    seed=config.seed,
                )
                for data in encoded_images
            ]

    def generateToFiles(self, request: ImageGenerationRequest) -> list[Path]:
        """Generate images and write them to the configured output path."""

        with self._lock:
            config = buildGenerationConfigFromRequest(
                request,
                self.user_config,
            )
            images = self._generateTensor(config)
            return saveImages(
                images,
                config.output,
                config.output_extension,
                config.output_quality,
            )

    def _generateTensor(self, config: GenerationConfig):
        if self.model_residency == "cpu-cache":
            return self._generateCpuCache(config)
        return self._generateStaged(config)

    def _generateStaged(self, config: GenerationConfig):
        model_sources = resolveModelSourcesFromConfig(self.user_config.models)
        text_dtype = componentDtype(config, model_sources.text_encoder)
        diffusion_dtype = componentDtype(config, model_sources.diffusion_model)
        vae_dtype = componentDtype(config, model_sources.vae)

        text_encoder = ZImageTextEncoder(
            model_sources.text_encoder,
            config.tokenizer_path,
            config.device,
            text_dtype,
        )
        conditioning = text_encoder.encodePrompts(
            config.prompt,
            config.negative_prompt,
        )
        del text_encoder
        releaseMemory()

        model = ZImageModel(
            model_sources.diffusion_model,
            config.device,
            diffusion_dtype,
        )
        latent = sampleLatents(
            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()

        vae = ZImageVae(model_sources.vae, config.device, vae_dtype)
        images = vae.decode(latent)
        del vae
        releaseMemory()
        return images

    def _generateCpuCache(self, config: GenerationConfig):
        import torch

        model_sources = resolveModelSourcesFromConfig(self.user_config.models)
        text_dtype = componentDtype(config, model_sources.text_encoder)
        diffusion_dtype = componentDtype(config, model_sources.diffusion_model)
        vae_dtype = componentDtype(config, model_sources.vae)
        cpu = torch.device("cpu")

        if self._text_encoder is None:
            self._text_encoder = ZImageTextEncoder(
                model_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()

        if self._model is None:
            self._model = ZImageModel(
                model_sources.diffusion_model,
                cpu,
                diffusion_dtype,
            )
        self._model.toDevice(config.device, diffusion_dtype)
        try:
            latent = sampleLatents(
                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()

        if self._vae is None:
            self._vae = ZImageVae(model_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