app.py

"""
app.py — Story → Shapes backend, built on gradio.Server (FastAPI + Gradio engine).

Serves the custom HTML/JS frontend at "/" and exposes JSON API endpoints the
frontend calls. The model judgment (affect) comes from the model backend
(llama.cpp in-process, or a Modal GPU endpoint); the deterministic scoring lives
here; geometry + color rendering happen client-side (the frontend has the
ported renderer).

Run locally (llama.cpp, the default backend):
    pip install -r requirements.txt
    python app.py          # serves http://localhost:7860 (pulls the GGUF on first run)
    # CPU-only? set STORY_SHAPES_LLAMACPP_GPU_LAYERS=0

On a HF Space: keep STORY_SHAPES_BACKEND=llamacpp (in-process GGUF) or set
modal_llm + STORY_SHAPES_LLM_MODAL_URL to offload the LLM to Modal.
"""
import os
from gradio import Server
from fastapi.responses import HTMLResponse, FileResponse

from model import backend
from engine.scorer import score, PUZZLE_THRESHOLD

import logging
log = logging.getLogger("story_shapes.api")

# PyTorch Fix
import os
import platform
if platform.system() == "Windows":
    import ctypes
    from importlib.util import find_spec
    try:
        if (spec := find_spec("torch")) and spec.origin and os.path.exists(
            dll_path := os.path.join(os.path.dirname(spec.origin), "lib", "c10.dll")
        ):
            ctypes.CDLL(os.path.normpath(dll_path))
    except Exception:
        log.warning("c10.dll not found")
        pass

app = Server()
STATIC = os.path.join(os.path.dirname(__file__), "static")

# ---- API endpoints (queued/streamed by Gradio's engine) ----

# ZeroGPU has a limited GPU pool — serialise all model calls so Gradio's queue
# absorbs backpressure rather than ZeroGPU returning 429.
@app.api(name="judge_beat", concurrency_limit=1)
def judge_beat(text: str, story: str = "", mode: str = "exploration", pacing: dict = None) -> dict:
    """Read a story beat -> affect judgment (+ exploration flags) + comment."""
    return backend.judge_beat(text, story, mode, pacing)

@app.api(name="judge_beat_segmented", concurrency_limit=1)
def judge_beat_segmented(text: str, story: str = "", pacing: dict = None) -> dict:
    """Per-segment essence: chunks the beat into phrases, each with own affect."""
    return backend.judge_beat_segmented(text, story, pacing)

@app.api(name="judge_attempt", concurrency_limit=1)
def judge_attempt(target_valence: float, target_arousal: float, target_dominance: float,
                  shape_sentence: str) -> dict:
    """Puzzle mode: judge a shape-sentence attempt and score it vs the target."""
    target = {"valence": target_valence, "arousal": target_arousal, "dominance": target_dominance}
    judg = backend.judge_attempt(target, shape_sentence)
    attempt = {"valence": judg["valence"], "arousal": judg["arousal"], "dominance": judg["dominance"]}
    sc = score(target, attempt)
    result = {**judg, **sc, "cleared": sc["score"] >= PUZZLE_THRESHOLD, "threshold": PUZZLE_THRESHOLD}
    log.info("judge_attempt scored: %.2f  %s  (cleared=%s)", sc["score"], sc["band"], result["cleared"])
    return result

@app.api(name="continue_story", concurrency_limit=1)
def continue_story(story: str) -> dict:
    """Pass-the-pen: model writes one continuation sentence."""
    return {"sentence": backend.continue_story(story)}

@app.api(name="reveal", concurrency_limit=1)
def reveal(story: str, labels: list) -> dict:
    """End of story: companions + layout roles + final coherence."""
    return backend.reveal(story, labels)

@app.api(name="title_story", concurrency_limit=1)
def title_story(story: str) -> dict:
    """Name the finished story for the share card."""
    return {"title": backend.title_story(story)}

@app.api(name="paint", concurrency_limit=1)
def paint(composition_png: str, theme: str = "", labels: list = None,
          mode: str = "strong") -> dict:
    """Turn the shape composition into a painting via FLUX.2 Klein img2img.
    composition_png: base64 PNG of the current scene (incl. user drags).
    mode: 'strong' (preserve composition) or 'loose' (inspiration only)."""
    from model import painter
    png_b64 = painter.paint(composition_png, theme=theme, labels=labels or [], mode=mode)
    return {"image": png_b64}

# ---- serve the frontend ----

@app.get("/")
async def index():
    return FileResponse(os.path.join(STATIC, "index.html"))

@app.get("/health")
async def health():
    return {"status": "ok", "backend": backend.BACKEND, "model": backend.MODEL}

if __name__ == "__main__":
    app.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))