"""Diffusers-backed generation and on-demand ControlNet management."""

from __future__ import annotations

import gc
from typing import Any

import numpy as np
from PIL import Image

from .config import AppConfig
from .errors import GenerationError
from .pipeline_factory import PipelineFactory
from .preprocessors.depth import DepthPreprocessor
from .preprocessors.openpose import OpenPosePreprocessor
from .schema import GuidedImageRequest, InpaintImageRequest, TextToImageRequest


class ModelManager:
    """Own all model state; construction validates model compatibility at startup."""

    def __init__(self, config: AppConfig) -> None:
        self.config = config
        self.factory = PipelineFactory(config)
        self.base = self.factory.create_base()
        self._control_mode: str | None = None
        self._controlnet: Any | None = None
        self._control_pipeline: Any | None = None
        self._depth = DepthPreprocessor()
        self._openpose = OpenPosePreprocessor()

    def text_to_image(self, request: TextToImageRequest, seed: int) -> Image.Image:
        width, height = request.size.dimensions
        result = self.base(
            prompt=request.positive_prompt,
            negative_prompt=request.negative_prompt,
            width=width,
            height=height,
            num_inference_steps=self.config.generation.steps,
            guidance_scale=self.config.generation.cfg_scale,
            num_images_per_prompt=1,
            generator=self._generator(seed),
        )
        return self._one_image(result)

    def guided_image(
        self, request: GuidedImageRequest, image: Image.Image, seed: int
    ) -> Image.Image:
        width, height = request.size.dimensions
        image = image.resize((width, height), Image.Resampling.LANCZOS)
        if request.control_mode == "depth":
            control_image = self._depth(image)
        else:
            control_image = self._openpose(image, full=request.control_mode == "openpose_full")
        pipeline = self._pipeline_for(str(request.control_mode), inpaint=False)
        result = pipeline(
            prompt=request.positive_prompt,
            negative_prompt=request.negative_prompt,
            image=control_image,
            width=width,
            height=height,
            num_inference_steps=self.config.generation.steps,
            guidance_scale=self.config.generation.cfg_scale,
            num_images_per_prompt=1,
            generator=self._generator(seed),
        )
        return self._one_image(result)

    def inpaint_image(
        self,
        request: InpaintImageRequest,
        image: Image.Image,
        mask: Image.Image,
        seed: int,
    ) -> Image.Image:
        width, height = request.size.dimensions
        image = image.resize((width, height), Image.Resampling.LANCZOS)
        mask = mask.resize((width, height), Image.Resampling.NEAREST)
        control_image = self._inpaint_condition(image, mask)
        pipeline = self._pipeline_for("inpaint", inpaint=True)
        result = pipeline(
            prompt=request.positive_prompt,
            negative_prompt=request.negative_prompt,
            image=image,
            mask_image=mask,
            control_image=control_image,
            width=width,
            height=height,
            num_inference_steps=self.config.generation.steps,
            guidance_scale=self.config.generation.cfg_scale,
            num_images_per_prompt=1,
            generator=self._generator(seed),
        )
        return self._one_image(result)

    def _pipeline_for(self, mode: str, *, inpaint: bool) -> Any:
        model_mode = "openpose" if mode == "openpose_full" else mode
        if self._control_mode == model_mode and self._control_pipeline is not None:
            return self._control_pipeline
        self._release_controlnet()
        self._controlnet = self.factory.create_controlnet(self.config.controlnet_path(model_mode))
        self._control_pipeline = (
            self.factory.create_inpaint_pipeline(self.base, self._controlnet)
            if inpaint
            else self.factory.create_guided_pipeline(self.base, self._controlnet)
        )
        self._control_mode = model_mode
        return self._control_pipeline

    def _release_controlnet(self) -> None:
        self._control_pipeline = None
        self._controlnet = None
        self._control_mode = None
        gc.collect()
        try:
            import torch

            torch.cuda.empty_cache()
        except ImportError:
            pass

    @staticmethod
    def _generator(seed: int) -> Any:
        import torch

        return torch.Generator(device="cuda").manual_seed(seed)

    @staticmethod
    def _inpaint_condition(image: Image.Image, mask: Image.Image) -> Any:
        import torch

        image_array = np.asarray(image.convert("RGB"), dtype=np.float32) / 255.0
        mask_array = np.asarray(mask.convert("L"), dtype=np.float32) / 255.0
        image_array[mask_array > 0.5] = -1.0
        return torch.from_numpy(image_array[None].transpose(0, 3, 1, 2))

    @staticmethod
    def _one_image(result: Any) -> Image.Image:
        images = getattr(result, "images", None)
        if not images or len(images) != 1 or not isinstance(images[0], Image.Image):
            raise GenerationError("Diffusers did not return exactly one image")
        return images[0]