"""Local Qwen Image Edit backend (needs a CUDA GPU).

Mirrors ``linoyts/Qwen-Image-Edit-Angles``: the base Qwen Image Edit 2509
pipeline with Phr00t's Rapid-AIO transformer for fast 4-step inference and
dx8152's multiple-angles LoRA fused in. Heavy imports are deferred so the
rest of the app works on CPU-only machines.
"""

from __future__ import annotations

import random
from typing import Optional

from PIL import Image

from ..config import (
    ANGLES_LORA_REPO,
    ANGLES_LORA_WEIGHT,
    QWEN_BASE_MODEL,
    QWEN_RAPID_TRANSFORMER,
)
from ..images import fit_image
from .base import ImageEditBackend


class LocalQwenBackend(ImageEditBackend):
    source = "local_qwen_image_edit"

    def __init__(self, image_size: int, lora_scale: float = 1.25) -> None:
        self.image_size = image_size
        self.lora_scale = lora_scale
        self._pipe = None
        self._device = "cpu"

    def prepare(self) -> None:
        import torch  # local import: heavy dep
        from diffusers import QwenImageEditPlusPipeline, QwenImageTransformer2DModel

        if not torch.cuda.is_available():
            raise RuntimeError("Local Qwen backend requires a CUDA GPU.")

        self._device = "cuda"
        dtype = torch.bfloat16

        transformer = QwenImageTransformer2DModel.from_pretrained(
            QWEN_RAPID_TRANSFORMER,
            subfolder="transformer",
            torch_dtype=dtype,
            device_map="cuda",
        )
        pipe = QwenImageEditPlusPipeline.from_pretrained(
            QWEN_BASE_MODEL,
            transformer=transformer,
            torch_dtype=dtype,
        ).to(self._device)

        pipe.load_lora_weights(
            ANGLES_LORA_REPO,
            weight_name=ANGLES_LORA_WEIGHT,
            adapter_name="angles",
        )
        pipe.set_adapters(["angles"], adapter_weights=[1.0])
        pipe.fuse_lora(adapter_names=["angles"], lora_scale=self.lora_scale)
        pipe.unload_lora_weights()
        self._pipe = pipe

    def edit(
        self,
        image: Image.Image,
        prompt: str,
        seed: int,
        num_inference_steps: int,
        true_guidance_scale: float,
    ) -> Image.Image:
        import torch  # local import: heavy dep

        base = fit_image(image.convert("RGB"), self.image_size)
        if not prompt.strip():
            return base
        if self._pipe is None:
            raise RuntimeError("Backend not prepared; call prepare() first.")

        generator = torch.Generator(device=self._device).manual_seed(int(seed))
        result = self._pipe(
            image=[base],
            prompt=prompt,
            width=base.width,
            height=base.height,
            num_inference_steps=num_inference_steps,
            generator=generator,
            true_cfg_scale=true_guidance_scale,
            num_images_per_prompt=1,
        ).images[0]
        return fit_image(result.convert("RGB"), self.image_size)