"""
Lulluna — Illustration Engine (SD Turbo, ~860 M params)

Two usage modes:
  1. Direct import  — called from app.py on HF Spaces (unified env)
  2. Subprocess     — called by app.py on local dev (separate .venv_tts)
     python illustrate.py  <reads JSON prompt from stdin, writes base64 PNG JSON to stdout>

The diffusers pipeline is cached after first load so repeated illustration
requests within the same process (import mode) don't reload the weights each time.
"""
import sys
import os
import json
import base64
import io
import torch
from diffusers import AutoPipelineForText2Image
from PIL import Image

# Device: MPS on Apple Silicon, CUDA on NVIDIA (HF Spaces GPU), CPU fallback.
DEVICE = (
    "mps"  if torch.backends.mps.is_available() else
    "cuda" if torch.cuda.is_available()          else
    "cpu"
)
DTYPE       = torch.float16 if DEVICE in ("mps", "cuda") else torch.float32
MODEL_ID    = "stabilityai/sd-turbo"
STEPS       = int(os.getenv("ILLUS_STEPS",  "4"))
WIDTH       = int(os.getenv("ILLUS_WIDTH",  "512"))
HEIGHT      = int(os.getenv("ILLUS_HEIGHT", "512"))

STYLE_SUFFIX = (
    ", children's book illustration, watercolor painting, "
    "soft pastel colors, cute, whimsical, storybook art style, "
    "warm lighting, no text"
)

# Module-level cache — populated on first call, reused for all subsequent calls.
_pipe: "AutoPipelineForText2Image | None" = None


def _get_pipe() -> "AutoPipelineForText2Image":
    global _pipe
    if _pipe is None:
        _pipe = AutoPipelineForText2Image.from_pretrained(
            MODEL_ID,
            torch_dtype=DTYPE,
            variant="fp16" if DTYPE == torch.float16 else None,
        )
        _pipe.to(DEVICE)
        _pipe.set_progress_bar_config(disable=True)
    return _pipe


def generate(prompt: str) -> dict:
    """Generate a storybook illustration for *prompt*. Returns {image: base64, mime: image/png}."""
    pipe = _get_pipe()
    full_prompt = prompt.strip() + STYLE_SUFFIX
    image: Image.Image = pipe(
        prompt=full_prompt,
        num_inference_steps=STEPS,
        guidance_scale=0.0,   # SD Turbo is distilled — no CFG needed
        width=WIDTH,
        height=HEIGHT,
    ).images[0]
    buf = io.BytesIO()
    image.save(buf, format="PNG")
    b64 = base64.b64encode(buf.getvalue()).decode()
    return {"image": b64, "mime": "image/png"}


if __name__ == "__main__":
    # Subprocess mode: read JSON prompt from stdin, write JSON to stdout.
    data = json.loads(sys.stdin.read().strip())
    prompt = data.get("prompt", "")
    if not prompt:
        print("ERROR: empty prompt", file=sys.stderr)
        sys.exit(1)
    try:
        result = generate(prompt)
        print(json.dumps(result), end="")
    except Exception as e:
        print(f"ERROR: {e}", file=sys.stderr)
        sys.exit(1)