"""Modal deployment: FLUX comic-panel image generation (FLUX.2 [klein] 4B).

Renders the AI-written `image_prompt` into a single landscape comic strip
(the prompt itself describes the 1-3 horizontal panels). Separate GPU class
from the llama.cpp text model; same bearer auth.

Deploy:  modal deploy modal_app/image.py
Secrets: modal secret create bds-auth BDS_TOKEN=<hex>        (shared w/ text)
         modal secret create huggingface HF_TOKEN=<hf token> (FLUX is gated)
Then set on the HF Space / locally:
  FLUX_URL=<printed generate_image endpoint url>
  FLUX_TOKEN=<same BDS_TOKEN hex>

Model is env-swappable at deploy time. If the FLUX.2-klein deps/VRAM are
troublesome, set  FLUX_MODEL_ID=black-forest-labs/FLUX.1-schnell  (Apache-2.0,
rock-solid 4-step) — the DiffusionPipeline auto-resolves either one.
LICENSE NOTE: the FLUX.2 line is typically non-commercial — fine for a
hackathon demo; flag before any commercial use.
"""
from __future__ import annotations

import base64
import io
import os
import time

import modal

MODEL_ID = os.environ.get("FLUX_MODEL_ID", "black-forest-labs/FLUX.2-klein-4B")
STEPS = int(os.environ.get("FLUX_STEPS", "4"))        # distilled → few steps
GUIDANCE = float(os.environ.get("FLUX_GUIDANCE", "1.0"))  # klein-4B card value
WIDTH = int(os.environ.get("FLUX_WIDTH", "1024"))     # landscape comic strip
HEIGHT = int(os.environ.get("FLUX_HEIGHT", "576"))


def _download() -> None:
    from huggingface_hub import snapshot_download
    snapshot_download(MODEL_ID, token=os.environ.get("HF_TOKEN"))


image = (
    modal.Image.debian_slim(python_version="3.11")
    .apt_install("git")  # needed to pip-install diffusers from its git repo
    .pip_install(
        "torch",
        "transformers",
        "accelerate",
        "sentencepiece",
        "protobuf",
        "pillow",
        "fastapi[standard]",
        "huggingface_hub",
        # FLUX.2 needs a recent diffusers — pull from git so the newest
        # pipeline classes are present
        "git+https://github.com/huggingface/diffusers.git",
    )
    .run_function(_download, secrets=[modal.Secret.from_name("huggingface")])
)

app = modal.App("brad-comic", image=image)

with image.imports():
    import torch
    from diffusers import Flux2KleinPipeline
    from fastapi import HTTPException, Request


@app.cls(
    gpu="A10G",  # 24GB; bump to A100-40GB if FLUX.2's text encoder OOMs
    scaledown_window=300,
    timeout=300,
    secrets=[
        modal.Secret.from_name("bds-auth"),
        modal.Secret.from_name("huggingface"),
    ],
)
class Flux:
    @modal.enter()
    def load(self):
        self.pipe = Flux2KleinPipeline.from_pretrained(
            MODEL_ID,
            torch_dtype=torch.bfloat16,
            token=os.environ.get("HF_TOKEN"),
        )
        # klein-4B needs ~13GB → fits the A10G's 24GB directly (fast). Set
        # FLUX_CPU_OFFLOAD=1 to trade speed for VRAM if a bigger model OOMs.
        if os.environ.get("FLUX_CPU_OFFLOAD") == "1":
            self.pipe.enable_model_cpu_offload()
        else:
            self.pipe.to("cuda")

    @modal.fastapi_endpoint(method="POST")
    def generate_image(self, body: dict, request: Request):
        expected = os.environ.get("BDS_TOKEN", "")
        sent = request.headers.get("authorization", "")
        if expected and sent != f"Bearer {expected}":
            raise HTTPException(401, "bad token")

        if body.get("warmup"):  # cold-start ping (runs @enter, loads weights)
            return {"ok": True, "warm": True}

        prompt = (body.get("prompt") or "").strip()
        if not prompt:
            return {"ok": False, "error": "empty prompt"}

        t0 = time.time()
        try:
            result = self.pipe(
                prompt=prompt,  # keyword required — pos-0 isn't prompt on FLUX.2
                num_inference_steps=STEPS,
                guidance_scale=GUIDANCE,
                width=WIDTH,
                height=HEIGHT,
            )
            img = result.images[0]
            buf = io.BytesIO()
            img.save(buf, format="PNG")
            b64 = base64.b64encode(buf.getvalue()).decode()
            return {"ok": True, "image_b64": b64,
                    "ms": int((time.time() - t0) * 1000)}
        except Exception as exc:  # caller skips the overlay; never crash
            return {"ok": False, "error": str(exc)[:200],
                    "ms": int((time.time() - t0) * 1000)}