"""Lazy Diffusers construction kept outside the MCP layer."""

from __future__ import annotations

from pathlib import Path
from typing import Any

from .config import AppConfig
from .errors import ConfigError


class PipelineFactory:
    """Build SD1.x pipelines in FP16 without network access."""

    def __init__(self, config: AppConfig) -> None:
        self.config = config

    @staticmethod
    def _runtime() -> tuple[Any, Any]:
        try:
            import diffusers
            import torch
        except ImportError as exc:
            raise ConfigError(f"inference dependency is unavailable: {exc}") from exc
        if not torch.cuda.is_available():
            raise ConfigError("CUDA GPU is required but torch.cuda.is_available() is false")
        return torch, diffusers

    def create_base(self) -> Any:
        torch, diffusers = self._runtime()
        cache_dir = str(self.config.model.root / "huggingface")
        try:
            pipeline = diffusers.StableDiffusionPipeline.from_single_file(
                str(self.config.model.checkpoint_path()),
                torch_dtype=torch.float16,
                local_files_only=True,
                cache_dir=cache_dir,
            )
            vae_path = self.config.model.vae_path()
            if vae_path is not None:
                # from_single_file needs a Diffusers config directory to avoid
                # network access.  Reuse the VAE config that the base pipeline
                # already resolved from the local huggingface cache.
                vae_config_dir = self._vae_config_dir(cache_dir)
                pipeline.vae = diffusers.AutoencoderKL.from_single_file(
                    str(vae_path),
                    torch_dtype=torch.float16,
                    local_files_only=True,
                    cache_dir=cache_dir,
                    config=vae_config_dir,
                )
            if getattr(pipeline, "vae", None) is None:
                raise ConfigError(
                    "the checkpoint does not provide a VAE and no external VAE was loaded"
                )
            adapter_names: list[str] = []
            adapter_weights: list[float] = []
            for index, (item, path) in enumerate(
                zip(self.config.model.loras, self.config.model.lora_paths(), strict=True)
            ):
                name = f"configured_lora_{index}"
                pipeline.load_lora_weights(
                    str(path.parent),
                    weight_name=path.name,
                    adapter_name=name,
                    local_files_only=True,
                )
                adapter_names.append(name)
                adapter_weights.append(item.weight)
            if adapter_names:
                pipeline.set_adapters(adapter_names, adapter_weights=adapter_weights)
            pipeline.scheduler = self._create_scheduler(pipeline.scheduler, diffusers)
            pipeline.set_progress_bar_config(disable=True)
            return pipeline.to("cuda")
        except ConfigError:
            raise
        except Exception as exc:
            raise ConfigError(f"failed to load SD1.x model configuration: {exc}") from exc

    def create_controlnet(self, path: Path) -> Any:
        torch, diffusers = self._runtime()
        try:
            return diffusers.ControlNetModel.from_single_file(
                str(path),
                torch_dtype=torch.float16,
                local_files_only=True,
                cache_dir=str(self.config.model.root / "huggingface"),
            ).to("cuda")
        except Exception as exc:
            raise ConfigError(f"failed to load ControlNet model {path}: {exc}") from exc

    @staticmethod
    def create_guided_pipeline(base: Any, controlnet: Any) -> Any:
        from diffusers.pipelines.controlnet.pipeline_controlnet import (
            StableDiffusionControlNetPipeline,
        )

        components: dict[str, Any] = dict(
            vae=base.vae,
            text_encoder=base.text_encoder,
            tokenizer=base.tokenizer,
            unet=base.unet,
            controlnet=controlnet,
            scheduler=base.scheduler,
            safety_checker=base.safety_checker,
            feature_extractor=base.feature_extractor,
            requires_safety_checker=getattr(base, "requires_safety_checker", False),
        )
        if getattr(base, "image_encoder", None) is not None:
            components["image_encoder"] = base.image_encoder
        pipeline = StableDiffusionControlNetPipeline(**components)
        pipeline.set_progress_bar_config(disable=True)
        return pipeline.to("cuda")

    @staticmethod
    def create_inpaint_pipeline(base: Any, controlnet: Any) -> Any:
        from diffusers.pipelines.controlnet.pipeline_controlnet_inpaint import (
            StableDiffusionControlNetInpaintPipeline,
        )

        components: dict[str, Any] = dict(
            vae=base.vae,
            text_encoder=base.text_encoder,
            tokenizer=base.tokenizer,
            unet=base.unet,
            controlnet=controlnet,
            scheduler=base.scheduler,
            safety_checker=base.safety_checker,
            feature_extractor=base.feature_extractor,
            requires_safety_checker=getattr(base, "requires_safety_checker", False),
        )
        if getattr(base, "image_encoder", None) is not None:
            components["image_encoder"] = base.image_encoder
        pipeline = StableDiffusionControlNetInpaintPipeline(**components)
        pipeline.set_progress_bar_config(disable=True)
        return pipeline.to("cuda")

    def _create_scheduler(self, current: Any, diffusers: Any) -> Any:
        name = self.config.generation.scheduler
        if name == "dpmpp_2m_karras":
            return diffusers.DPMSolverMultistepScheduler.from_config(
                current.config,
                algorithm_type="dpmsolver++",
                solver_order=2,
                use_karras_sigmas=True,
            )
        if name == "dpmpp_2m_sde_karras":
            return diffusers.DPMSolverMultistepScheduler.from_config(
                current.config,
                algorithm_type="sde-dpmsolver++",
                solver_order=2,
                use_karras_sigmas=True,
            )
        if name == "ddim_trailing":
            return diffusers.DDIMScheduler.from_config(
                current.config,
                timestep_spacing="trailing",
            )
        classes = {
            "euler_a": diffusers.EulerAncestralDiscreteScheduler,
            "ddim": diffusers.DDIMScheduler,
            "lcm": diffusers.LCMScheduler,
        }
        try:
            scheduler_class = classes[name]
        except KeyError as exc:
            raise ConfigError(f"unsupported scheduler: {name}") from exc
        return scheduler_class.from_config(current.config)

    @staticmethod
    def _vae_config_dir(cache_dir: str) -> str:
        """Find the VAE config.json inside the local huggingface cache."""
        snapshots = (
            Path(cache_dir) / "models--stable-diffusion-v1-5--stable-diffusion-v1-5" / "snapshots"
        )
        if not snapshots.is_dir():
            raise ConfigError(
                "cannot locate VAE config: huggingface cache snapshots directory not found"
            )
        for snapshot in snapshots.iterdir():
            vae_dir = snapshot / "vae"
            if (vae_dir / "config.json").is_file():
                return str(vae_dir)
        raise ConfigError("cannot locate VAE config.json in the local huggingface cache")