"""
model/painter.py — painting backend abstraction.

BACKENDS (set via STORY_SHAPES_PAINT_BACKEND):
  "modal"      — HTTP POST to a deployed Modal endpoint (recommended).
                 Fast, no local VRAM needed.
                 Requires STORY_SHAPES_PAINT_MODAL_URL to be set.
                 Deploy once with: modal deploy modal_painter.py
  "flux_local" — local diffusers inference (slow on 8GB, left for reference).

STORY_SHAPES_PAINT_MODAL_URL: the URL printed after `modal deploy modal_painter.py`,
  e.g. https://<workspace>--story-shapes-painter-paint-endpoint.modal.run
"""
import os, io, base64, json, logging, urllib.request, random

import spaces

log = logging.getLogger("story_shapes.painter")

PAINT_BACKEND  = os.environ.get("STORY_SHAPES_PAINT_BACKEND", "modal")
MODAL_URL      = os.environ.get("STORY_SHAPES_PAINT_MODAL_URL", "")
FLUX_MODEL     = os.environ.get("STORY_SHAPES_FLUX_MODEL", "black-forest-labs/FLUX.2-klein-4B")

STRENGTH = {"strong": 0.45, "loose": 0.80}


# Rich per-style descriptors so the chosen style genuinely steers the look,
# instead of a single word tacked onto a fixed "art-print" base (which made
# every style render nearly identical). Keys are lowercased; the frontend's
# "Surprise me!" list maps onto these.
STYLE_PROMPTS = {
    "art print":    "fine-art giclée print, layered paper-and-paint forms, rich grain, gallery quality",
    "cut paper":    "cut-paper collage, crisp torn edges, layered construction-paper relief, bold flat color, Matisse-like",
    "ink & wash":   "sumi-e ink and wash, flowing brushwork, bleeding washes on rice paper, generous negative space, muted tones",
    "neon glass":   "luminous neon glass, glowing translucent forms, dark backdrop, electric rim light, vivid saturated color",
    "oil impasto":  "thick oil impasto, heavy palette-knife strokes, visible ridges of paint, dramatic light",
    "risograph":    "risograph print, limited spot-color inks, halftone grain, slight misregistration, retro zine look",
    "dream poster": "surreal dream poster, soft gradients, hazy atmospheric glow, vintage offset print",
}

def _build_prompt(theme, labels):
    base = ("expressive non-literal abstract art that preserves the original shapes, "
            "colors, positions, and composition of the reference image")
    t = (theme or "").strip()
    style = STYLE_PROMPTS.get(t.lower(), t) if t else "mixed-media abstract"
    # Lead with the style so it carries weight, then the preservation clause.
    return f"{style}; {base}"


# ---------------------------------------------------------------------------
# Modal backend — POST to the deployed web endpoint
# ---------------------------------------------------------------------------
def _paint_modal(image_b64, prompt, strength, steps):
    if not MODAL_URL:
        raise RuntimeError(
            "STORY_SHAPES_PAINT_MODAL_URL is not set. "
            "Run `modal deploy modal_painter.py` and set the printed URL."
        )
    body = json.dumps({
        "image_b64": image_b64,
        "prompt":    prompt,
        "strength":  strength,
        "steps":     steps,
    }).encode()
    req = urllib.request.Request(
        MODAL_URL,
        data=body,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    log.info("paint → Modal  prompt=%r  strength=%.2f", prompt, strength)
    with urllib.request.urlopen(req, timeout=300) as r:
        resp = json.loads(r.read())
    return resp["image_b64"]


# ---------------------------------------------------------------------------
# Local backend — diffusers FLUX.2 Klein (kept for reference / HF Space)
# ---------------------------------------------------------------------------
_pipe = None

def _load_local_pipe():
    global _pipe
    if _pipe is not None:
        return _pipe
    import torch
    from diffusers import Flux2KleinPipeline
    log.info("loading FLUX.2 Klein locally (%s)…", FLUX_MODEL)
    _pipe = Flux2KleinPipeline.from_pretrained(
        FLUX_MODEL,
        torch_dtype=torch.bfloat16
    )
    _pipe.enable_model_cpu_offload()
    return _pipe

def free_painter():
    global _pipe
    if _pipe is not None:
        del _pipe; _pipe = None
        try:
            import torch, gc; gc.collect(); torch.cuda.empty_cache()
        except Exception:
            pass
        log.info("freed local FLUX pipeline")

# @_GPU(duration=90)
@spaces.GPU(duration=60)
def _paint_local(image_b64, prompt, strength, steps):
    import torch
    from PIL import Image
    pipe = _load_local_pipe()
    init = Image.open(io.BytesIO(base64.b64decode(image_b64))).convert("RGB")
    init = init.resize((1024, 1024))
    device = "cuda" if torch.cuda.is_available() else "cpu"
    result = pipe(
        prompt=prompt,
        image=init,
        num_inference_steps=steps,
        generator=torch.Generator(device=device).manual_seed(random.randint(0, 2**31 - 1)),
    ).images[0]
    buf = io.BytesIO(); result.save(buf, "PNG")
    return base64.b64encode(buf.getvalue()).decode()


# ---------------------------------------------------------------------------
# Public entry point
# ---------------------------------------------------------------------------
def paint(composition_png_b64, theme="", labels=None, mode="strong", steps=4):
    """
    composition_png_b64: data-URL or bare base64 PNG of the current scene.
    Returns bare base64 PNG of the generated painting.
    """
    labels   = labels or []
    prompt   = _build_prompt(theme, labels)
    strength = STRENGTH.get(mode, STRENGTH["strong"])
    log.info("paint  backend=%s  mode=%s  strength=%.2f  prompt=%r",
             PAINT_BACKEND, mode, strength, prompt)

    # strip data-URL prefix if present
    if "," in composition_png_b64[:32]:
        composition_png_b64 = composition_png_b64.split(",", 1)[1]

    if PAINT_BACKEND == "modal":
        return _paint_modal(composition_png_b64, prompt, strength, steps)
    elif PAINT_BACKEND == "flux_local":
        return _paint_local(composition_png_b64, prompt, strength, steps)
    else:
        raise ValueError(f"Unknown STORY_SHAPES_PAINT_BACKEND: {PAINT_BACKEND!r}")