"""
Nowhere Post
============
Name a place that doesn't exist. The post office of Nowhere issues a stamp for it.
Collect them in your album.

Track 2 toy for the "Small Models Big Adventures" hackathon.
Model: black-forest-labs/FLUX.1-schnell (~12B, Apache-2.0, under the 32B ceiling).

Every result is composed into a perforated postage-stamp frame with PIL, so it
always looks like a stamp. In model mode FLUX paints the picture; in keeper mode a
procedural scene is drawn, so the app boots and demos with no GPU.
"""

import os
import random
import inspect

import gradio as gr
from PIL import Image, ImageDraw, ImageFont

_BLOCKS_HAS_CSS = "css" in inspect.signature(gr.Blocks.__init__).parameters
_LAUNCH_HAS_SSR = "ssr_mode" in inspect.signature(gr.Blocks.launch).parameters

MODEL_ID = os.environ.get("POST_MODEL", "black-forest-labs/FLUX.1-schnell")
DEBUG = os.environ.get("POST_DEBUG", "").strip().lower() in {"1", "true", "yes"}
ART = 384                                       # inner picture size (square)
STEPS = int(os.environ.get("POST_STEPS", "4"))  # schnell is happy at 4

EXAMPLES = ["the Isle of Lost Umbrellas", "Cloudfall", "the Last Lighthouse",
            "the Market at the Edge of Sleep"]

def _noop_gpu(*a, **k):
    def wrap(fn):
        return fn
    return wrap(a[0]) if a and callable(a[0]) else wrap

if os.environ.get("SPACES_ZERO_GPU", "").lower() in {"true", "1"}:
    try:
        import spaces
        GPU = spaces.GPU
    except Exception:                          # noqa: BLE001
        GPU = _noop_gpu
else:
    GPU = _noop_gpu

_pipe = None
MODE = "keeper"
_warmed = False


def load_model():
    global _pipe, MODE
    try:
        import torch
        from diffusers import FluxPipeline
        _pipe = FluxPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)
        MODE = "model"
        print(f"[Post] Loaded {MODEL_ID} -- model mode.")
    except Exception as exc:                    # noqa: BLE001
        MODE = "keeper"
        print(f"[Post] Could not load {MODEL_ID} ({exc}). Keeper mode active.")


load_model()


def _font(size):
    for path in ("DejaVuSans-Bold.ttf", "DejaVuSans.ttf",
                 "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf"):
        try:
            return ImageFont.truetype(path, size)
        except Exception:                       # noqa: BLE001
            continue
    return ImageFont.load_default()


@GPU(duration=60)
def _flux_art(place):
    import torch
    if torch.cuda.is_available() and getattr(_pipe, "device", None) is not None \
            and _pipe.device.type != "cuda":
        _pipe.to("cuda")
    prompt = (f"a vintage postage stamp illustration of {place}, storybook engraving "
              f"style, soft muted inks, fine linework, whimsical, centered scene")
    g = torch.Generator(device=_pipe.device).manual_seed(random.randint(0, 2**31 - 1))
    img = _pipe(prompt, num_inference_steps=STEPS, guidance_scale=0.0,
                height=ART, width=ART, max_sequence_length=256, generator=g).images[0]
    return img


def _procedural_art(place):
    """A seeded little scene for keeper mode, so a stamp always has a picture."""
    rng = random.Random(place.strip().lower())
    palettes = [((214, 198, 150), (122, 150, 120), (90, 110, 140)),
                ((232, 210, 170), (180, 140, 110), (120, 90, 120)),
                ((205, 220, 225), (120, 160, 170), (70, 100, 120)),
                ((235, 200, 160), (200, 120, 90), (110, 70, 90))]
    sky, land, accent = rng.choice(palettes)
    img = Image.new("RGB", (ART, ART), sky)
    d = ImageDraw.Draw(img)
    horizon = int(ART * rng.uniform(0.55, 0.72))
    d.rectangle([0, horizon, ART, ART], fill=land)
    # a sun or moon
    r = rng.randint(28, 46)
    cxp = rng.randint(r, ART - r)
    d.ellipse([cxp - r, horizon - 90 - r, cxp + r, horizon - 90 + r], fill=accent)
    # a few hills / shapes
    for _ in range(rng.randint(2, 4)):
        hw = rng.randint(80, 200)
        hx = rng.randint(-40, ART)
        hh = rng.randint(30, 80)
        d.ellipse([hx, horizon - hh, hx + hw, horizon + hh], fill=accent)
    return img


def compose_stamp(inner, place):
    """Frame a square picture as a perforated postage stamp with a caption."""
    pad, cap = 26, 56
    W, H = ART + pad * 2, ART + pad * 2 + cap
    stamp = Image.new("RGB", (W, H), (253, 250, 240))
    d = ImageDraw.Draw(stamp)
    # perforations
    step = 16
    for x in range(0, W, step):
        d.ellipse([x - 5, -5, x + 5, 5], fill=(243, 234, 215))
        d.ellipse([x - 5, H - 5, x + 5, H + 5], fill=(243, 234, 215))
    for y in range(0, H, step):
        d.ellipse([-5, y - 5, 5, y + 5], fill=(243, 234, 215))
        d.ellipse([W - 5, y - 5, W + 5, y + 5], fill=(243, 234, 215))
    d.rectangle([10, 10, W - 10, H - 10], outline=(170, 100, 28), width=2)
    stamp.paste(inner.resize((ART, ART)), (pad, pad))
    d.rectangle([pad, pad, pad + ART, pad + ART], outline=(170, 100, 28), width=1)
    # header + caption
    fhead, fcap = _font(15), _font(20)
    d.text((pad, 12), "NOWHERE  POST", font=fhead, fill=(122, 46, 46))
    denom = random.choice(["1", "2", "3", "5", "7"])
    d.text((W - pad - 30, 12), f"{denom} \u2606", font=fhead, fill=(122, 46, 46))
    name = place if len(place) <= 30 else place[:29] + "\u2026"
    tw = d.textlength(name, font=fcap)
    d.text(((W - tw) / 2, pad + ART + 14), name, font=fcap, fill=(58, 47, 37))
    return stamp


def make_stamp(place):
    inner = None
    if MODE == "model":
        try:
            inner = _flux_art(place)
        except Exception as exc:                # noqa: BLE001
            print(f"[Post] generation error: {exc}")
            inner = None
    if inner is None:
        inner = _procedural_art(place)
    return compose_stamp(inner, place.strip())


def on_issue(place, album):
    global _warmed
    place = (place or "").strip()
    if not place:
        yield None, album, album, ""
        return
    if MODE == "model" and not _warmed:
        # show a placeholder frame while FLUX warms up
        warming = compose_stamp(_procedural_art(place), place)
        yield warming, album, album, ""
    stamp = make_stamp(place)
    if MODE == "model":
        _warmed = True
    album = (album or []) + [(stamp, place)]
    yield stamp, [s for s, _ in album], album, ""


def on_reset():
    return None, [], [], ""


def use_example(text):
    return text


CSS = """
@import url('https://fonts.googleapis.com/css2?family=Fraunces:ital,opsz,wght@0,9..144,400;0,9..144,600;1,9..144,400&family=Spectral:ital,wght@0,400;0,500;1,400&display=swap');
:root{--paper:#f3ead7;--paper-2:#efe3c9;--ink:#2f2a22;--ink-soft:#6b5a47;
  --post:#7a2e2e;--amber-deep:#aa641c;--line:#c9b48f;}
.gradio-container,.gradio-container.dark,.dark{
  --body-background-fill:transparent;--background-fill-primary:#fffaf0;--background-fill-secondary:#f6edd8;
  --block-background-fill:#fffaf0;--block-border-color:var(--line);--border-color-primary:var(--line);
  --body-text-color:var(--ink);--body-text-color-subdued:var(--ink-soft);
  --block-label-text-color:var(--post);--block-title-text-color:var(--ink);
  --block-label-background-fill:#efe3c9;--block-title-background-fill:transparent;
  --input-background-fill:#fffaf0;--input-border-color:var(--line);--input-placeholder-color:var(--ink-soft);
  --button-primary-background-fill:var(--post);--button-primary-background-fill-hover:#641f1f;
  --button-primary-text-color:#fbf7ec;--button-primary-border-color:#641f1f;
  --button-secondary-background-fill:#efe3c9;--button-secondary-background-fill-hover:#e7d6b3;
  --button-secondary-text-color:var(--ink);--button-secondary-border-color:var(--line);
  --color-accent:var(--amber-deep);--color-accent-soft:#f3e3c4;}
.gradio-container{background:radial-gradient(120% 80% at 80% -10%,#f7efdd,var(--paper) 55%,var(--paper-2));
  font-family:'Spectral',Georgia,serif !important;color:var(--ink) !important;max-width:960px !important;}
.gradio-container textarea,.gradio-container input[type="text"],.gradio-container input:not([type]){
  background:#fffaf0 !important;color:var(--ink) !important;-webkit-text-fill-color:var(--ink) !important;border-color:var(--line) !important;}
.gradio-container textarea::placeholder,.gradio-container input::placeholder{color:var(--ink-soft) !important;-webkit-text-fill-color:var(--ink-soft) !important;opacity:1;}
.np-title{font-family:'Fraunces',serif;font-weight:600;font-size:2.5rem;line-height:1;margin:.2rem 0 0;}
.np-title em{font-style:italic;color:var(--post);}
.np-sub{font-style:italic;color:var(--ink-soft);margin:.35rem 0 1rem;font-size:1.05rem;}
.np-mode{display:inline-block;font-size:.72rem;letter-spacing:.12em;text-transform:uppercase;color:var(--post);border:1px solid var(--line);border-radius:999px;padding:.15rem .6rem;}
.np-foot{color:var(--ink-soft);font-size:.82rem;font-style:italic;text-align:center;margin-top:10px;}
footer{display:none !important;}
"""

_bk = {"title": "Nowhere Post"}
if _BLOCKS_HAS_CSS:
    _bk["css"] = CSS
    _bk["theme"] = gr.themes.Soft()

with gr.Blocks(**_bk) as demo:
    album = gr.State([])
    public = "FLUX.1-schnell · on-device" if MODE == "model" else "Nowhere Post"
    mode_label = f"{public}  ·  [{MODE}]" if DEBUG else public
    gr.HTML(f"""
      <div><div class="np-title">Nowhere <em>Post</em></div>
      <div class="np-sub">Name a place that doesn't exist. Receive its postage stamp.</div>
      <span class="np-mode">{mode_label}</span></div>""")

    with gr.Row():
        with gr.Column(scale=3):
            with gr.Row():
                place = gr.Textbox(placeholder="a place that doesn't exist... (the Isle of Lost Umbrellas)",
                                   show_label=False, scale=7, autofocus=True)
                go = gr.Button("Issue stamp", variant="primary", scale=2)
            with gr.Row():
                ex_btns = [gr.Button(e, size="sm") for e in EXAMPLES]
            stamp = gr.Image(label="Latest stamp", type="pil", height=460, show_label=False)
            reset = gr.Button("Start a new album", size="sm")
        with gr.Column(scale=2):
            gallery = gr.Gallery(label="Your album", columns=2, height=520, show_label=True)

    gr.HTML('<div class="np-foot">Every place that isn\'t real still deserves a stamp.</div>')

    go.click(on_issue, [place, album], [stamp, gallery, album, place])
    place.submit(on_issue, [place, album], [stamp, gallery, album, place])
    reset.click(on_reset, None, [stamp, gallery, album, place])
    for b, e in zip(ex_btns, EXAMPLES):
        b.click(use_example, gr.State(e), place)

if __name__ == "__main__":
    _lk = {}
    if not _BLOCKS_HAS_CSS:
        _lk["css"] = CSS
        _lk["theme"] = gr.themes.Soft()
    if _LAUNCH_HAS_SSR:
        _lk["ssr_mode"] = False
    demo.queue(max_size=24).launch(**_lk)