"""Modal backend — LLM + image generation on cloud GPUs, single app deployment.

Deploy once:
    modal deploy modal_app.py

Download models to volumes before first run:
    modal run modal_app.py::download_model
    modal run modal_app.py::download_image_model

Then launch locally:
    uv run python app.py
"""

from __future__ import annotations

import modal

hf_secret = modal.Secret.from_name("huggingface")

# =========================================================================== #
#  Container images
# =========================================================================== #

llm_image = (
    modal.Image.from_registry("nvidia/cuda:12.4.0-devel-ubuntu22.04", add_python="3.12")
    .pip_install(
        "llama-cpp-python",
        extra_options="--extra-index-url https://abetlen.github.io/llama-cpp-python/whl/cu124",
    )
    .pip_install("huggingface_hub")
)

painter_image = (
    modal.Image.from_registry("nvidia/cuda:12.4.0-devel-ubuntu22.04", add_python="3.12")
    .pip_install(
        "torch==2.5.1",
        "torchvision==0.20.1",
        extra_options="--index-url https://download.pytorch.org/whl/cu124",
    )
    .pip_install(
        "diffusers==0.35.1",
        "transformers==4.47.1",
        "accelerate==1.8.1",
        "peft",  # required by diffusers load_lora_weights / fuse_lora
        "safetensors",
        "Pillow",
        "huggingface_hub",
        "rembg",  # background removal for sprites
        "onnxruntime",  # rembg runtime dep
    )
)

# =========================================================================== #
#  Volumes  (model weights persist across cold starts)
# =========================================================================== #

llm_volume = modal.Volume.from_name("vn-models", create_if_missing=True)
LLM_DIR = "/models"

image_volume = modal.Volume.from_name("vn-image-models", create_if_missing=True)
IMAGE_DIR = "/image-models"

# =========================================================================== #
#  Single app — both classes deploy together with `modal deploy modal_app.py`
# =========================================================================== #

app = modal.App("vn-app")

GGUF_REPO = "Qwen/Qwen3-14B-GGUF"
GGUF_FILE = "Qwen3-14B-Q8_0.gguf"
# Painter: SDXL-base-1.0 + ByteDance Lightning LoRA (matches local SdxlLightningPainter)
SDXL_BASE_REPO = "stabilityai/stable-diffusion-xl-base-1.0"
SDXL_BASE_LOCAL = "sdxl-base"
LIGHTNING_LORA_REPO = "ByteDance/SDXL-Lightning"
LIGHTNING_LORA_FILE = "sdxl_lightning_4step_lora.safetensors"


# --------------------------------------------------------------------------- #
#  Download helpers (run once each)
# --------------------------------------------------------------------------- #


@app.function(image=llm_image, volumes={LLM_DIR: llm_volume}, timeout=600, secrets=[hf_secret])
def download_model() -> None:
    from huggingface_hub import hf_hub_download

    print(f"Downloading {GGUF_REPO}/{GGUF_FILE} ...")
    path = hf_hub_download(repo_id=GGUF_REPO, filename=GGUF_FILE, local_dir=LLM_DIR)
    llm_volume.commit()
    print(f"Saved -> {path}")


@app.function(
    image=painter_image, volumes={IMAGE_DIR: image_volume}, timeout=900, secrets=[hf_secret]
)
def download_image_model() -> None:
    """Download SDXL-base-1.0 weights (~6.5 GB) to the volume.

    The Lightning LoRA (ByteDance/SDXL-Lightning) is small (~400 MB) and loaded
    at inference time directly from HF Hub — no need to pre-download it.
    """
    from huggingface_hub import snapshot_download

    print(f"Downloading {SDXL_BASE_REPO} -> {IMAGE_DIR}/{SDXL_BASE_LOCAL} ...")
    snapshot_download(SDXL_BASE_REPO, local_dir=f"{IMAGE_DIR}/{SDXL_BASE_LOCAL}")
    image_volume.commit()
    print("SDXL-base saved to volume.")


# --------------------------------------------------------------------------- #
#  LLM backend  (llama.cpp + Qwen3-14B on A10G)
# --------------------------------------------------------------------------- #


@app.cls(
    image=llm_image,
    gpu="A10G",
    volumes={LLM_DIR: llm_volume},
    timeout=300,
    scaledown_window=600,  # keep warm 10 min — reloading the 15 GB GGUF mid-session is worse
)
class ModalLLMBackend:
    @modal.enter()
    def load(self) -> None:
        from llama_cpp import Llama

        self.llm = Llama(
            model_path=f"{LLM_DIR}/{GGUF_FILE}",
            n_ctx=8192,
            n_gpu_layers=-1,
            verbose=False,
        )
        print("[modal] LLM loaded on GPU")

    @modal.method()
    def complete(self, messages: list[dict], **kw) -> str:
        out = self.llm.create_chat_completion(messages=messages, **kw)
        return out["choices"][0]["message"]["content"]

    @modal.method()
    def complete_json(self, messages: list[dict], schema: dict, **kw) -> dict:
        import json

        prompt_chars = sum(len(m.get("content", "")) for m in messages)
        print(f"[modal] LLM complete_json: {len(messages)} msgs, {prompt_chars} chars in")
        out = self.llm.create_chat_completion(
            messages=messages,
            response_format={"type": "json_object", "schema": schema},
            temperature=kw.get("temperature", 0.7),
            top_p=kw.get("top_p", 0.9),
            max_tokens=kw.get("max_tokens", 512),
            presence_penalty=kw.get("presence_penalty", 0.0),
        )

        content = out["choices"][0]["message"]["content"]
        print(f"[modal] LLM complete_json: {len(content)} chars out")
        return json.loads(content)


# --------------------------------------------------------------------------- #
#  Painter backend  (SDXL-base-1.0 + Lightning LoRA on A10G)
# --------------------------------------------------------------------------- #


@app.cls(
    image=painter_image,
    gpu="A10G",
    volumes={IMAGE_DIR: image_volume},
    timeout=120,
    scaledown_window=600,  # keep warm 10 min — avoids cold-start during a play session
)
class ModalPainterBackend:
    """SDXL-base-1.0 + ByteDance SDXL-Lightning 4-step LoRA.

    Mirrors the local SdxlLightningPainter exactly:
      - EulerDiscreteScheduler with trailing timestep spacing
      - LoRA fused into weights (fuse_lora) for faster inference
      - guidance_scale=0.0 (required by Lightning distillation)
      - 4 inference steps
      - Optional rembg background removal for sprites
    """

    @modal.enter()
    def load(self) -> None:
        import torch
        from diffusers import AutoencoderKL, EulerDiscreteScheduler, StableDiffusionXLPipeline

        print(f"[modal] Loading SDXL-base from {IMAGE_DIR}/{SDXL_BASE_LOCAL} ...")
        # fp16-safe VAE: decodes natively in fp16 (no "weird pixels", no per-render
        # fp32 upcast like the deprecated pipe.upcast_vae() workaround)
        vae = AutoencoderKL.from_pretrained(
            "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16
        )
        self.pipe = StableDiffusionXLPipeline.from_pretrained(
            f"{IMAGE_DIR}/{SDXL_BASE_LOCAL}",
            torch_dtype=torch.float16,
            variant="fp16",
            vae=vae,
        ).to("cuda")

        # Lightning requires EulerDiscrete with trailing timestep spacing
        self.pipe.scheduler = EulerDiscreteScheduler.from_config(
            self.pipe.scheduler.config, timestep_spacing="trailing"
        )

        # Load Lightning LoRA from HF Hub (~400 MB, fast)
        print(f"[modal] Loading Lightning LoRA from {LIGHTNING_LORA_REPO} ...")
        self.pipe.load_lora_weights(LIGHTNING_LORA_REPO, weight_name=LIGHTNING_LORA_FILE)
        self.pipe.fuse_lora()  # bake into weights for faster inference

        self.torch = torch
        print("[modal] SDXL-Lightning ready on GPU")

    @modal.method()
    def render(
        self,
        prompt: str,
        negative_prompt: str,
        seed: int,
        size: int,
        steps: int,
        guidance_scale: float = 0.0,
        remove_bg: bool = False,
    ) -> bytes:
        """Generate an image and return PNG bytes.

        Args:
            prompt: Positive prompt.
            negative_prompt: Negative prompt.
            seed: RNG seed for reproducibility.
            size: Square image side in pixels.
            steps: Inference steps (4 for Lightning).
            guidance_scale: CFG scale — 0.0 for Lightning sprites, >1.0 for backdrops.
            remove_bg: Run rembg on the output (sprites only).
        """
        import io

        gen = self.torch.Generator(device="cuda").manual_seed(seed)
        result = self.pipe(
            prompt=prompt,
            negative_prompt=negative_prompt or None,
            num_inference_steps=steps,
            guidance_scale=guidance_scale,
            height=size,
            width=size,
            generator=gen,
        )
        img = result.images[0]

        if remove_bg:
            from rembg import new_session, remove  # noqa: PLC0415

            # Reuse one ONNX session per container — remove() without a session
            # reloads the ~170MB u2net model on every sprite. Lazy (not in
            # @modal.enter()) so backdrop-only requests never pay for it.
            if getattr(self, "_rembg_session", None) is None:
                self._rembg_session = new_session()
            img = remove(img, session=self._rembg_session)  # returns RGBA PIL image

        buf = io.BytesIO()
        img.save(buf, format="PNG")  # PNG supports RGBA transparency
        return buf.getvalue()


# =========================================================================== #
#  Smoke tests
# =========================================================================== #


@app.local_entrypoint()
def smoke() -> None:
    backend = ModalLLMBackend()
    reply = backend.complete.remote(
        [{"role": "user", "content": "Say hello in one word."}], max_tokens=10
    )
    print("LLM smoke:", reply)


@app.local_entrypoint()
def smoke_painter() -> None:
    import io

    from PIL import Image

    backend = ModalPainterBackend()

    # Backdrop test
    png = backend.render.remote(
        prompt="Japanese anime forest, glowing mushrooms, painterly background",
        negative_prompt="text, watermark, characters, person",
        seed=42,
        size=512,
        steps=4,
        guidance_scale=0.0,
        remove_bg=False,
    )
    img = Image.open(io.BytesIO(png))
    img.save("smoke_backdrop.png")
    print(f"Backdrop smoke OK -> smoke_backdrop.png  {img.size}")

    # Sprite test (with rembg background removal)
    png_sprite = backend.render.remote(
        prompt="anime girl, school uniform, happy expression, white background",
        negative_prompt="text, watermark, scenery, complex background",
        seed=7,
        size=512,
        steps=4,
        guidance_scale=0.0,
        remove_bg=True,
    )
    img_sprite = Image.open(io.BytesIO(png_sprite))
    img_sprite.save("smoke_sprite.png")
    print(f"Sprite smoke OK  -> smoke_sprite.png   {img_sprite.size}  mode={img_sprite.mode}")