app.py

"""Flask app: JSON API + the static frontend, one process.

Routes are deliberately thin. They validate input, hand off to a provider, and
serialise the result — no domain logic lives here. See CONTRACT.md.
"""

from __future__ import annotations

import base64
import hashlib
import json
import os
import sys
import tempfile
from pathlib import Path

from flask import Flask, Response, jsonify, request, send_file, send_from_directory

import bedrock
import config
import export
import gltf
import providers
from auth import RateLimiter, limiter, require_token
from jobs import runner
from logs import log
from rewards import TERM_LABELS
from schemas import (
    BONES,
    Clip,
    ProviderError,
    Rig,
    TrainConfig,
    schema_json,
    validate_clip,
    validate_rig,
)
from store import store
from training import TrainingRun

app = Flask(__name__, static_folder=None)
app.config["MAX_CONTENT_LENGTH"] = config.MAX_UPLOAD_BYTES

# --------------------------------------------------------------------------
# The animator/ directory (formerly phase2/) is a standalone CLI tool, not a
# package -- its files use bare `import gltf_utils` etc. assuming their own
# directory is on sys.path. Inserting it at position 0 makes `import animator`
# resolve to animator/animator.py (the CLI's dispatcher module) rather than
# the animator/ directory itself, which Python would otherwise also be
# willing to treat as an (empty) implicit namespace package from repo root.
# Deliberately does NOT import animator/main.py, which pulls in pywebview
# (a desktop GUI toolkit) at module level -- not something to load into a
# server process.
# --------------------------------------------------------------------------
ANIMATOR_DIR = Path(__file__).parent / "animator"
sys.path.insert(0, str(ANIMATOR_DIR))
import animator as llm_animator_core  # noqa: E402  (animator/animator.py)
import gltf_utils as llm_animator_gltf_utils  # noqa: E402
import llm_client as llm_animator_client  # noqa: E402

# Keep dict order as declared. The bone tree is written parents-first in
# schemas.py and reads far better that way in the API too; Flask would
# otherwise sort it alphabetically.
app.json.sort_keys = False

STATIC_DIR = Path(__file__).parent / "static"

# Separate bucket from the LLM-endpoint limiter in auth.py: this one guards a
# route with no bearer token, so it's keyed by client address alone.
render_limiter = RateLimiter("RENDER_RATE_PER_MINUTE", "RENDER_RATE_PER_DAY")
tpose_limiter = RateLimiter("TPOSE_RATE_PER_MINUTE", "TPOSE_RATE_PER_DAY")


# --------------------------------------------------------------------------
# Errors — every failure reaches the browser in the same shape, so the frontend
# has exactly one error path to render.
# --------------------------------------------------------------------------

def fail(message: str, status: int = 400):
    return jsonify({"error": message}), status


@app.errorhandler(404)
def _not_found(_):
    if request.path.startswith("/api/"):
        return fail("Not found.", 404)
    return send_from_directory(STATIC_DIR, "index.html")


@app.errorhandler(413)
def _too_large(_):
    mb = config.MAX_UPLOAD_BYTES // (1024 * 1024)
    return fail(f"That image is too big — keep it under {mb}MB.", 413)


# --------------------------------------------------------------------------
# Static frontend
# --------------------------------------------------------------------------

@app.get("/")
def index():
    return send_from_directory(STATIC_DIR, "index.html")


@app.get("/<path:filename>")
def static_files(filename: str):
    return send_from_directory(STATIC_DIR, filename)


# --------------------------------------------------------------------------
# Schema + status
# --------------------------------------------------------------------------

@app.get("/api/schema")
def get_schema():
    """The browser reads bone names and reward terms from here rather than
    hard-coding them, so the contract can only ever be defined in one place."""
    return jsonify({
        **schema_json(),
        "reward_labels": TERM_LABELS,
        "providers": providers.active(),
    })


# --------------------------------------------------------------------------
# Phase 1 — rendered image to rigged avatar
# --------------------------------------------------------------------------

@app.post("/api/avatars")
def create_avatar():
    upload = request.files.get("image")
    if upload is None:
        return fail("No image was uploaded.")

    image_bytes = upload.read()
    if not image_bytes:
        return fail("That image was empty. Try drawing something!")
    mime = upload.mimetype or "image/png"

    rigger = providers.get_rigger()

    def work(progress):
        rig = validate_rig(rigger.rig(image_bytes, mime, progress))
        avatar = store.add_avatar(rig, image_bytes=image_bytes, mime=mime)
        return avatar.to_json()

    job = runner.submit(work, message="Waking up your avatar...")
    return jsonify(job.to_json()), 202


@app.post("/api/avatars/glb")
def sideload_avatar():
    """Create an avatar directly from a compatible rigged GLB."""
    # Rigged models are routinely much larger than sketch images. Override the
    # app-wide sketch limit before Werkzeug parses the multipart body.
    request.max_content_length = request.content_length or (1 << 63) - 1
    upload = request.files.get("glb")
    if upload is None:
        return fail("No GLB was uploaded.")

    glb_bytes = upload.read()
    if not glb_bytes:
        return fail("That GLB was empty.")

    try:
        gltf.validate_avatar_glb(glb_bytes)
    except gltf.GlbError as exc:
        log(f"[GLB sideload] rejected {upload.filename!r}: {exc}")
        return fail(f"That GLB can't be used: {exc}.")

    rig = validate_rig(Rig(
        format="glb",
        skeleton=list(BONES),
        glb_bytes=glb_bytes,
        notes="Sideloaded rigged GLB.",
    ))
    avatar = store.add_avatar(rig)
    return jsonify(avatar.to_json()), 201


@app.get("/api/avatars/<avatar_id>")
def get_avatar(avatar_id: str):
    avatar = store.get_avatar(avatar_id)
    if not avatar:
        return fail("That avatar doesn't exist.", 404)
    return jsonify(avatar.to_json())


@app.get("/api/avatars/<avatar_id>/image")
def get_avatar_image(avatar_id: str):
    avatar = store.get_avatar(avatar_id)
    if not avatar or not avatar.image_path:
        return fail("No source image saved for that avatar.", 404)
    return send_file(avatar.image_path)


@app.get("/api/avatars/<avatar_id>/glb")
def get_avatar_glb(avatar_id: str):
    avatar = store.get_avatar(avatar_id)
    if not avatar or not avatar.rig.glb_bytes:
        return fail("That avatar has no GLB — it's drawn procedurally.", 404)
    # Without validators a browser heuristically caches a 3MB binary and never
    # asks again — so a swapped fixture keeps rendering the old body with no
    # sign anything is stale. An ETag over the bytes plus no-cache means the
    # browser always revalidates, but pays for the transfer only when the
    # content actually differs (304 otherwise).
    etag = hashlib.sha256(avatar.rig.glb_bytes).hexdigest()[:32]
    if request.if_none_match.contains(etag):
        return Response(status=304, headers={"ETag": f'"{etag}"',
                                             "Cache-Control": "no-cache"})

    return Response(avatar.rig.glb_bytes, mimetype="model/gltf-binary",
                    headers={"ETag": f'"{etag}"', "Cache-Control": "no-cache"})


@app.post("/api/renders")
def render_sketch():
    """Render an uploaded drawing without creating or rigging an avatar."""
    upload = request.files.get("image")
    if upload is None:
        return fail("No image was uploaded.")

    image_bytes = upload.read()
    if not image_bytes:
        return fail("That image was empty. Try drawing something!")

    prompt = request.form.get("prompt", "").strip()
    if not prompt:
        return fail("Describe how you'd like this rendered.")
    if len(prompt) > config.BEDROCK_MAX_PROMPT_CHARS:
        return fail(f"Prompt is too long — the limit is "
                    f"{config.BEDROCK_MAX_PROMPT_CHARS} characters.", 413)

    refusal = render_limiter.check(request.remote_addr or "unknown")
    if refusal:
        return fail(refusal, 429)

    def work(progress):
        try:
            # Rigging and the later T-pose transform both consume this render,
            # so a close-up here cannot be recovered downstream.
            result = bedrock.render_sketch(
                image_bytes, f"{prompt}, {bedrock.FULL_BODY_HINT}",
                negative_prompt=bedrock.FULL_BODY_NEGATIVE_HINT)
        except bedrock.BedrockError as exc:
            raise ProviderError(exc.message, detail=exc.detail) from exc
        return {
            "image_base64": base64.b64encode(result["image_bytes"]).decode("ascii"),
            "output_format": result["output_format"],
            "seed": result["seed"],
        }

    job = runner.submit(work, message="Rendering your character...")
    return jsonify(job.to_json()), 202


@app.post("/api/avatars/<avatar_id>/tpose")
def tpose_avatar(avatar_id: str):
    """Turn the avatar into a forward-facing, T-pose, transparent-background
    PNG, via bedrock.tpose_transform. Fixed shape, no request body — always
    the same transform. The avatar image is the rendered PNG that was supplied
    to rigging, so both downstream operations use the same character design."""
    avatar = store.get_avatar(avatar_id)
    if not avatar or not avatar.image_path:
        return fail("That avatar doesn't exist.", 404)

    refusal = tpose_limiter.check(request.remote_addr or "unknown")
    if refusal:
        return fail(refusal, 429)

    image_bytes = Path(avatar.image_path).read_bytes()

    def work(progress):
        try:
            result = bedrock.tpose_transform(image_bytes)
        except bedrock.BedrockError as exc:
            raise ProviderError(exc.message, detail=exc.detail) from exc
        return {
            "image_base64": base64.b64encode(result["image_bytes"]).decode("ascii"),
            "output_format": result["output_format"],
        }

    job = runner.submit(work, message="Posing your character...")
    return jsonify(job.to_json()), 202


# --------------------------------------------------------------------------
# Phase 2 — prompt to pose
# --------------------------------------------------------------------------

@app.post("/api/avatars/<avatar_id>/poses")
def create_pose(avatar_id: str):
    avatar = store.get_avatar(avatar_id)
    if not avatar:
        return fail("That avatar doesn't exist.", 404)

    prompt = (request.json or {}).get("prompt", "").strip()
    if not prompt:
        return fail("Tell me what you'd like your avatar to do!")

    poser = providers.get_poser()

    def work(progress):
        clip = validate_clip(poser.pose(prompt, avatar.rig, progress))
        clip.prompt = clip.prompt or prompt
        return store.add_clip(clip).to_json()

    job = runner.submit(work, message="Thinking about that move...")
    return jsonify(job.to_json()), 202


@app.get("/api/clips/<clip_id>")
def get_clip(clip_id: str):
    clip = store.get_clip(clip_id)
    if not clip:
        return fail("That move doesn't exist.", 404)
    return jsonify(clip.to_json())


# --------------------------------------------------------------------------
# Jobs
# --------------------------------------------------------------------------

@app.get("/api/jobs/<job_id>")
def get_job(job_id: str):
    job = runner.get(job_id)
    if not job:
        return fail("That job has expired.", 404)
    return jsonify(job.to_json())


# --------------------------------------------------------------------------
# Phases 3 & 4 — training
# --------------------------------------------------------------------------

@app.post("/api/training/runs")
def create_run():
    body = request.json or {}
    avatar = store.get_avatar(body.get("avatar_id", ""))
    if not avatar:
        return fail("That avatar doesn't exist.", 404)

    target = store.get_clip(body.get("target_clip_id", ""))
    if not target:
        return fail("Pick a move to train towards first.", 404)

    cfg = TrainConfig.from_json(body.get("config"))
    run = TrainingRun(avatar.id, target, cfg)
    store.add_run(run)
    run.set_speed(body.get("speed", 1.0))
    run.start(providers.get_trainer(), avatar.rig)
    return jsonify(run.to_json()), 201


@app.get("/api/training/runs/<run_id>")
def get_run(run_id: str):
    run = store.get_run(run_id)
    if not run:
        return fail("That training run doesn't exist.", 404)
    return jsonify(run.to_json())


# --------------------------------------------------------------------------
# Training output
#
# What a finished run produces, for whatever comes next. Two artifacts because
# they have different audiences: the GLB is for rendering and 3D tools, the
# JSON is for the animation step. See export.py — in particular, do not hand
# the GLB to a language model.
# --------------------------------------------------------------------------

def _export_inputs(run_id: str):
    """Resolve a run to (run, avatar, clip), or raise ExportError."""
    run = store.get_run(run_id)
    if not run:
        raise export.ExportError("That training run doesn't exist.", 404)
    avatar = store.get_avatar(run.avatar_id)
    if not avatar:
        raise export.ExportError("That avatar is no longer around.", 404,
                                 "avatar evicted from the in-memory store")
    return run, avatar, run.target_clip


@app.get("/api/training/runs/<run_id>/export.glb")
def export_run_glb(run_id: str):
    """The learned pose, baked into the avatar's own model."""
    try:
        run, avatar, clip = _export_inputs(run_id)
        data, _ = export.build_glb(run, avatar, clip)
    except export.ExportError as exc:
        if exc.detail:
            log(f"[export] {run_id}: {exc.detail}")
        return fail(exc.user_message, exc.status)

    name = f"{clip.name.replace(' ', '-').lower()}-{run_id}.glb"
    return Response(data, mimetype="model/gltf-binary", headers={
        "Content-Disposition": f'attachment; filename="{name}"',
        "Cache-Control": "no-cache",
    })


@app.get("/api/training/runs/<run_id>/export.json")
def export_run_json(run_id: str):
    """Bones, start pose, end pose, provenance — the animation step's input."""
    try:
        run, avatar, clip = _export_inputs(run_id)
        document = export.build_document(run, avatar, clip)
    except export.ExportError as exc:
        if exc.detail:
            log(f"[export] {run_id}: {exc.detail}")
        return fail(exc.user_message, exc.status)
    return jsonify(document)


@app.get("/api/training/runs/<run_id>/events")
def stream_run(run_id: str):
    run = store.get_run(run_id)
    if not run:
        return fail("That training run doesn't exist.", 404)

    def generate():
        # Flush a comment straight away so the proxy commits to the response
        # and the browser's EventSource opens rather than sitting on a buffer.
        yield ": open\n\n"
        for episode in run.events():
            if episode is None:
                yield ": keepalive\n\n"
                continue
            # The id lets a reconnecting browser tell us where it got to.
            yield (f"id: {episode.episode}\n"
                   f"event: episode\ndata: {json.dumps(episode.to_json())}\n\n")
        yield f"event: end\ndata: {json.dumps(run.to_json())}\n\n"

    return Response(generate(), mimetype="text/event-stream", headers={
        "Cache-Control": "no-cache",
        "X-Accel-Buffering": "no",   # don't let a proxy swallow the stream
        "Connection": "keep-alive",
    })


@app.post("/api/training/runs/<run_id>/<action>")
def control_run(run_id: str, action: str):
    run = store.get_run(run_id)
    if not run:
        return fail("That training run doesn't exist.", 404)

    if action == "stop":
        run.stop()
    elif action == "pause":
        run.pause()
    elif action == "resume":
        run.resume()
    elif action == "speed":
        run.set_speed((request.json or {}).get("speed", 1.0))
    else:
        return fail(f"Unknown action '{action}'.", 404)
    return jsonify(run.to_json())


# --------------------------------------------------------------------------
# Behaviours — the bridge from "a pose" to "a thing my avatar does"
# --------------------------------------------------------------------------

@app.get("/api/behaviours")
def list_behaviours():
    avatar_id = request.args.get("avatar_id")
    return jsonify({"behaviours": [b.to_json()
                                   for b in store.list_behaviours(avatar_id)]})


@app.post("/api/behaviours")
def create_behaviour():
    body = request.json or {}
    avatar = store.get_avatar(body.get("avatar_id", ""))
    if not avatar:
        return fail("That avatar doesn't exist.", 404)

    name = (body.get("name") or "").strip()
    if not name:
        return fail("Give this behaviour a name.")

    clip_id = body.get("clip_id")
    if clip_id:
        clip = store.get_clip(clip_id)
        if not clip:
            return fail("That move doesn't exist.", 404)
    elif body.get("clip"):
        clip = store.add_clip(validate_clip(body["clip"]))
    else:
        return fail("A behaviour needs a move to play.")

    behaviour = store.add_behaviour(
        name=name, clip=clip, avatar_id=avatar.id,
        trained=bool(body.get("trained")),
        best_reward=float(body.get("best_reward", 0.0)))
    return jsonify(behaviour.to_json()), 201


# --------------------------------------------------------------------------
# LLM animator — thin wrapper around the standalone animator/ CLI tool
#
# NOT part of the Rigger/Poser/Trainer contract (see CONTRACT.md): that
# contract's "next step: animation" is a future Animator provider working on
# the app's 16-bone skeleton and returning a Clip. This is a different, older
# tool that bakes a glTF Animation directly into a GLB on whatever skeleton
# the uploaded file has (e.g. a ~86-joint Mixamo rig) -- kept deliberately
# separate, not retrofitted into the contract. One request field per CLI arg;
# runs synchronously (no job/poll) since a call takes ~5-15s, and this
# prototype's OpenRouter account is capped, so no rate limiting either.
# --------------------------------------------------------------------------

_LLM_ANIMATOR_OPENROUTER_MODEL = "anthropic/claude-sonnet-5"
_LLM_ANIMATOR_OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1"


def _truthy(value: str | None) -> bool:
    return (value or "").strip().lower() in ("1", "true", "yes", "on")


@app.post("/api/llm-animator/generate")
def llm_animator_generate():
    prompt = (request.form.get("prompt") or "").strip() or None
    input_file = request.files.get("input")
    pose_b_file = request.files.get("pose_b")

    if not prompt and not pose_b_file:
        return fail("Provide a prompt, a pose_b GLB, or both.")

    use_llm = not _truthy(request.form.get("no_llm"))
    loop = _truthy(request.form.get("loop"))
    model = (request.form.get("model") or "").strip() or None
    base_url = (request.form.get("base_url") or "").strip() or None

    if _truthy(request.form.get("openrouter")):
        if not os.environ.get("OPENROUTER_API_KEY"):
            return fail("Server isn't configured with an OpenRouter API key.", 500)
        llm_animator_client.DEFAULT_MODEL = model or _LLM_ANIMATOR_OPENROUTER_MODEL
        llm_animator_client.DEFAULT_BASE_URL = base_url or _LLM_ANIMATOR_OPENROUTER_BASE_URL
    elif model or base_url:
        llm_animator_client.DEFAULT_MODEL = model or llm_animator_client.DEFAULT_MODEL
        llm_animator_client.DEFAULT_BASE_URL = base_url or llm_animator_client.DEFAULT_BASE_URL

    with tempfile.TemporaryDirectory() as tmp_name:
        tmp = Path(tmp_name)

        if input_file:
            input_path = tmp / "input.glb"
            input_file.save(input_path)
        else:
            input_path = ANIMATOR_DIR / "rigged_human.glb"

        gltf, joints = llm_animator_gltf_utils.load_skeleton(str(input_path))

        try:
            if pose_b_file:
                pose_b_path = tmp / "pose_b.glb"
                pose_b_file.save(pose_b_path)
                _, joints_b = llm_animator_gltf_utils.load_skeleton(str(pose_b_path))
                name, tracks = llm_animator_core.generate_transition(
                    joints, joints_b, style_prompt=prompt, use_llm=use_llm, loop=loop)
            else:
                name, tracks = llm_animator_core.generate(prompt, gltf, joints, use_llm=use_llm)
        except Exception as exc:
            return fail(f"Animation generation failed: {exc}", 502)

        llm_animator_gltf_utils.add_rotation_animation(gltf, tracks, name=name)

        output_path = tmp / "result.glb"
        llm_animator_gltf_utils.save(gltf, str(output_path))
        data = output_path.read_bytes()

    return Response(data, mimetype="model/gltf-binary", headers={
        "X-Animation-Name": name,
        "Content-Disposition": 'attachment; filename="result.glb"',
    })


# --------------------------------------------------------------------------
# Bedrock prompt endpoint
#
# A utility for the teams building the real providers: run a prompt against a
# Bedrock model and see what comes back. Guarded by a bearer token, a model
# allowlist and rate limits — see auth.py and bedrock.py.
#
# NOT called by the frontend, and the token must never be shipped to the
# browser: anything the browser holds is public, and this endpoint spends money.
# --------------------------------------------------------------------------

@app.post("/api/llm/generate")
@require_token
def llm_generate():
    body = request.json or {}

    prompt = (body.get("prompt") or "").strip()
    if not prompt:
        return fail("A 'prompt' is required.")
    if len(prompt) > config.BEDROCK_MAX_PROMPT_CHARS:
        return fail(f"Prompt is too long — the limit is "
                    f"{config.BEDROCK_MAX_PROMPT_CHARS} characters.", 413)

    model_id = (body.get("model_id") or "").strip()
    if not model_id:
        return fail("A 'model_id' is required. GET /api/llm/models lists the "
                    "ones this server allows.")

    system = (body.get("system") or "").strip() or None
    if system and len(system) > config.BEDROCK_MAX_PROMPT_CHARS:
        return fail("System prompt is too long.", 413)

    try:
        temperature = body.get("temperature")
        result = bedrock.converse(
            model_id, prompt,
            system=system,
            max_tokens=int(body.get("max_tokens", 1024)),
            temperature=None if temperature is None else float(temperature),
        )
    except bedrock.BedrockError as exc:
        return fail(exc.message, exc.status)
    except (TypeError, ValueError) as exc:
        return fail(f"Invalid request: {exc}", 400)

    return jsonify(result)


@app.get("/api/llm/models")
@require_token
def llm_models():
    """What this deployment will actually run, plus current rate-limit usage."""
    return jsonify({
        "models": bedrock.allowed_models(),
        "region": config.BEDROCK_REGION or None,
        "limits": {
            **limiter.snapshot(),
            "max_tokens": config.BEDROCK_MAX_TOKENS,
            "max_prompt_chars": config.BEDROCK_MAX_PROMPT_CHARS,
        },
    })


@app.errorhandler(ProviderError)
def _provider_error(exc: ProviderError):
    return fail(exc.user_message, 502)


if __name__ == "__main__":
    config.ensure_dirs()
    app.run(debug=True, threaded=True, port=5000)