"""
modal_painter.py — FLUX.2 Klein img2img inference on Modal GPU.

Setup:
    pip install modal
    modal setup   # opens browser to link your account

Deploy (once):
    modal deploy modal_painter.py

This prints a URL like:
    https://<your-workspace>--story-shapes-painter-paint.modal.run

Set that as STORY_SHAPES_MODAL_URL in your local .env / Space secrets.
Then set STORY_SHAPES_PAINT_BACKEND=modal in the same place.

The endpoint accepts POST with JSON:
    { "image_b64": "<base64 PNG>", "prompt": "...", "strength": 0.45, "steps": 4 }
And returns JSON:
    { "image_b64": "<base64 PNG of the painting>" }

GPU choice: L4 is the sweet spot for FLUX.2 Klein 4B — cheaper than A100,
faster than T4, and 24GB VRAM fits the model comfortably without CPU offload.
Change to gpu="a10g" if L4 is unavailable in your region.
"""
import io, base64, random, modal

app = modal.App("story-shapes-painter")

# Container image: diffusers from source (needed for Flux2KleinPipeline) + deps
image = (
    modal.Image.debian_slim(python_version="3.11")
    .pip_install(
        "diffusers>=0.38.0",
        "torch",
        "transformers",
        "accelerate",
        "sentencepiece",
        "pillow",
        "fastapi[standard]",
    )
)

MODEL_ID = "black-forest-labs/FLUX.2-klein-4B"

@app.cls(
    image=image,
    gpu="a10g",
    scaledown_window=600,
    timeout=120,
)
class Painter:
    @modal.enter()
    def load(self):
        import torch
        from diffusers import Flux2KleinPipeline
        print(f"Loading {MODEL_ID}…")
        self.pipe = Flux2KleinPipeline.from_pretrained(
            MODEL_ID, 
            torch_dtype=torch.bfloat16
        ).to("cuda")
        print("Model loaded.")

    @modal.method()
    def paint(self, image_b64: str, prompt: str,
              strength: float = 0.45, steps: int = 4) -> str:
        import torch
        from PIL import Image
        init = Image.open(io.BytesIO(base64.b64decode(image_b64))).convert("RGB")
        init = init.resize((1024, 1024))
        result = self.pipe(
            prompt=prompt,
            image=init,
            # strength=strength,
            num_inference_steps=steps,
            # guidance_scale=3.5,
            guidance_scale=1,
            generator=torch.Generator(device="cuda").manual_seed(random.randint(0, 2**31 - 1)),
        ).images[0]
        buf = io.BytesIO()
        result.save(buf, format="PNG")
        return base64.b64encode(buf.getvalue()).decode()


# Web endpoint — called by painter.py over plain HTTP
@app.function(image=image)
# @modal.web_endpoint(method="POST")
@modal.fastapi_endpoint(method="POST")
def paint_endpoint(body: dict) -> dict:
    """
    POST body: { image_b64, prompt, strength, steps }
    Returns:   { image_b64 }
    """
    painter = Painter()
    out = painter.paint.remote(
        image_b64=body["image_b64"],
        prompt=body["prompt"],
        strength=float(body.get("strength", 0.45)),
        steps=int(body.get("steps", 4)),
    )
    return {"image_b64": out}


# Local test entrypoint — modal run modal_painter.py
@app.local_entrypoint()
def test():
    from PIL import Image
    # tiny solid-colour init image
    img = Image.new("RGB", (64, 64), color=(20, 20, 40))
    buf = io.BytesIO(); img.save(buf, "PNG")
    b64 = base64.b64encode(buf.getvalue()).decode()
    result = paint_endpoint.remote(
        {"image_b64": b64, "prompt": "abstract painting test", "strength": 0.9, "steps": 4}
    )
    print("test OK — returned", len(result["image_b64"]), "b64 chars")