api_server.py

"""OmniNPC FastAPI Server — the backend for the HTML / Three.js demo.

Routes:
  POST /infer        — JPEG-base64 frame → Nano action JSON (perception)
  POST /infer_with_thinking — same as /infer but returns the raw thinking trace
  POST /super        — request a fresh strategic directive (Super agent)
  POST /dialogue     — when two NPCs are near each other, get a dialogue line
  GET  /state        — current global state snapshot (for the HUD)
  GET  /metrics      — latency / call counts
  GET  /health       — health check
  GET  /config       — frontend-safe config (model names, intervals)

CORS is open so the static HTML page can call this from any origin (file:// included).
On Hugging Face Spaces this server runs alongside the static site.

Requires NVIDIA_NIM_API_KEY in the environment.
"""

from __future__ import annotations

import argparse
import asyncio
import base64
import logging
import os
import time
from collections import Counter, deque

import cv2
import numpy as np
from dotenv import load_dotenv

from metrics.latency_tracker import LatencyTracker
from metrics.metrics_logger import MetricsLogger
latency_tracker = LatencyTracker()
load_dotenv()

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
)
logger = logging.getLogger("api_server")

try:
    from fastapi import FastAPI, HTTPException
    from fastapi.middleware.cors import CORSMiddleware
    from fastapi.responses import JSONResponse, FileResponse
    from fastapi.staticfiles import StaticFiles
    from pydantic import BaseModel
except ImportError:
    logger.error("FastAPI not installed. Run: pip install fastapi uvicorn")
    raise

from actions.action_schema import ActionResult, InferenceEvent
from capture.frame_differ import FrameDiffer
from config import load_config
from inference.safety_verifier import SafetyVerifier
from inference.nim_client import NIMClient
from inference.prompt_templates import (
    build_nano_system_prompt,
    build_nano_user_message,
    SUPER_SYSTEM_PROMPT,
    build_super_user_message,
)
from orchestrator.pipeline_state import NPCStateMachine


# ─── Request / Response models ────────────────────────────────────────────────

class InferRequest(BaseModel):
    image_b64: str           # raw base64 (no data: prefix)
    audio_text: str = ""     # transcribed voice command (or "")
    npc_role: str = "hostile"  # "hostile" | "ally" | "wanderer"
    directive: str = ""      # current Super directive (frontend-cached)


class SuperRequest(BaseModel):
    scene_summary: str
    mission_context: str = "Compete with the player and other NPCs to collect coins."
    recent_actions: list[str] = []


class DialogueRequest(BaseModel):
    npc_a: str               # role of NPC A
    npc_b: str               # role of NPC B
    scene_summary: str = ""
    last_event: str = ""     # e.g., "ally collected a coin"


class CommandRequest(BaseModel):
    text: str
    target_npc: str = "hostile"


# ─── App state ────────────────────────────────────────────────────────────────

class ServerState:
    def __init__(self) -> None:
        self.config = load_config()
        self.api_key = os.getenv("NVIDIA_NIM_API_KEY", "")
        self.nim = NIMClient(
            api_key=self.api_key, base_url=self.config.nvidia.base_url,
            timeout=self.config.nvidia.timeout,
        )
        self.last_directive = "Compete with the player to collect the most coins."
        self.start_time = time.time()
        self.total_infer_calls = 0
        self.total_super_calls = 0
        self.total_dialogue_calls = 0
        self.recent_latencies: deque[float] = deque(maxlen=50)
        self.recent_actions: deque[dict] = deque(maxlen=50)
        self.recent_thinks: deque[str] = deque(maxlen=20)
        self.action_histogram: Counter = Counter()
        self.metrics_logger: MetricsLogger | None = None
        # frame dumping removed — keep API behavior unchanged otherwise

    # frame dumping removed


state = ServerState()

caches = {
    "HOSTILE": {"last_frame": None, "last_action": None, "repeat_count": 0},
    "ALLY": {"last_frame": None, "last_action": None, "repeat_count": 0},
    "WANDERER": {"last_frame": None, "last_action": None, "repeat_count": 0},
}
fsm = {
    "HOSTILE": NPCStateMachine("HOSTILE"),
    "ALLY": NPCStateMachine("ALLY"),
    "WANDERER": NPCStateMachine("WANDERER"),
}
verifier = SafetyVerifier()
differ = FrameDiffer()


# ─── App + CORS ───────────────────────────────────────────────────────────────

app = FastAPI(
    title="OmniNPC API",
    description="Hierarchical Nemotron Super+Nano backend for the OmniNPC HTML demo.",
    version="2.0.0",
)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=False,
    allow_methods=["*"],
    allow_headers=["*"],
)


# ─── Helpers ──────────────────────────────────────────────────────────────────

import re
import json


def _strip_thinking(raw: str) -> tuple[str, str]:
    """Return (think_block, cleaned_response)."""
    think_match = re.search(r"<think>(.*?)</think>", raw, flags=re.DOTALL)
    think = think_match.group(1).strip() if think_match else ""
    cleaned = re.sub(r"<think>.*?</think>", "", raw, flags=re.DOTALL).strip()
    cleaned = re.sub(r"```[a-z]*", "", cleaned).replace("```", "").strip()
    return think, cleaned


def _extract_json(cleaned: str) -> dict:
    """Pull the first {...} JSON object out of a model response."""
    match = re.search(r"\{.*\}", cleaned, re.DOTALL)
    if not match:
        return {}
    try:
        return json.loads(match.group(0))
    except json.JSONDecodeError:
        return {}


def _role_persona(role: str) -> str:
    """Return persona text injected into Nano's system prompt for each NPC role."""
    role = role.lower()
    if role == "hostile":
        return ("You are the HOSTILE NPC. Your goal: collect coins faster than the player. "
                "If the player gets within 3m, evade or sprint away. Be competitive.")
    if role == "ally":
        return ("You are the ALLY NPC. You team up with the player. Stay near the player "
                "(within 4-6m), help collect coins, and warn about the hostile NPC.")
    if role == "wanderer":
        return ("You are the WANDERER NPC. You do your own thing — explore, collect coins, "
                "occasionally dance. Friendly but independent.")
    return ""


# ─── Routes ───────────────────────────────────────────────────────────────────

@app.get("/health")
async def health():
    return {
        "status": "ok",
        "version": "2.0.0",
        "uptime_s": round(time.time() - state.start_time, 1),
        "api_key_set": bool(state.api_key),
        "nano_model": state.config.nvidia.nano_omni_model,
        "super_model": state.config.nvidia.super_model,
    }


@app.get("/config")
async def get_config():
    return {
        "nano_model": state.config.nvidia.nano_omni_model,
        "super_model": state.config.nvidia.super_model,
        "perception_interval_ms": int(state.config.orchestrator.perception_interval * 1000),
        "strategy_interval_ms": int(state.config.orchestrator.strategy_interval * 1000),
    }


@app.get("/metrics")
async def metrics():
    lats = list(state.recent_latencies)
    avg = sum(lats) / len(lats) if lats else 0.0
    p95 = sorted(lats)[int(len(lats) * 0.95)] if len(lats) >= 5 else (max(lats) if lats else 0.0)
    uptime_s = time.time() - state.start_time
    rate_hz = round(state.total_infer_calls / uptime_s, 3) if uptime_s > 0 else 0.0
    return {
        "total_infer_calls": state.total_infer_calls,
        "total_super_calls": state.total_super_calls,
        "total_dialogue_calls": state.total_dialogue_calls,
        "avg_latency_ms": round(avg, 1),
        "p95_latency_ms": round(p95, 1),
        "recent_latencies_ms": [round(x, 1) for x in lats[-20:]],
        "current_directive": state.last_directive,
        "inference_rate_hz": rate_hz,
        "total_input_tokens": state.nim.total_input_tokens,
        "total_output_tokens": state.nim.total_output_tokens,
        "total_nim_errors": state.nim.total_errors,
        "action_histogram": dict(state.action_histogram.most_common()),
    }


@app.get("/state")
async def state_snapshot():
    return {
        "directive": state.last_directive,
        "recent_actions": list(state.recent_actions)[-10:],
        "recent_thinks": list(state.recent_thinks)[-5:],
    }


@app.post("/infer")
async def infer(req: InferRequest, directive: str = ""):
    """Run Nano perception on a single frame. Returns action JSON."""
    if not state.api_key:
        raise HTTPException(500, "NVIDIA_NIM_API_KEY not configured on server")
    if not req.image_b64:
        raise HTTPException(400, "image_b64 missing")

    try:
        frame_bytes = base64.b64decode(req.image_b64)
    except Exception as exc:
        raise HTTPException(400, "image_b64 is not valid base64") from exc

    frame_np = cv2.imdecode(
        np.frombuffer(frame_bytes, np.uint8),
        cv2.IMREAD_COLOR,
    )
    if frame_np is None:
        raise HTTPException(400, "image_b64 did not decode to a valid JPEG frame")
    role = req.npc_role
    role_key = role.upper()
    cache = caches[role_key]
    event = InferenceEvent(
        voice_command=bool(getattr(req, "voice_override", False)),
        new_directive=False,
        health_changed=False,
    )
    should_infer, reason = differ.should_infer(
        cache["last_frame"], frame_np, cache["repeat_count"]
    )
    latency_tracker.log_inference_decision(
        did_infer=should_infer,
        reason=reason,
        role=req.npc_role,
    )
    if not should_infer and cache["last_action"] is not None:
        cache["repeat_count"] = int(cache.get("repeat_count", 0)) + 1
        cached_response = dict(cache["last_action"])
        cached_response["cached"] = True
        cached_action = ActionResult(
            action=cached_response.get("action", "idle"),
            target=cached_response.get("target", ""),
            direction=cached_response.get("direction", "none"),
            magnitude=float(cached_response.get("magnitude", 0.5)),
            confidence=float(cached_response.get("confidence", 0.5)),
            say=cached_response.get("say", ""),
            scene_description=cached_response.get("scene_description", ""),
            capture_ts=time.perf_counter(),
            inference_ts=time.perf_counter(),
            inference_latency_ms=0.0,
        )
        state.metrics_logger.log_action(
            cached_action,
            latency_tracker.get_stats_dict(),
            did_infer=should_infer,
            skip_reason=reason,
            npc_role=role,
        )
        return cached_response

    persona = _role_persona(req.npc_role)
    directive = directive or req.directive or state.last_directive
    full_directive = f"{directive} {persona}".strip()
    system_prompt = build_nano_system_prompt(full_directive)

    audio_ctx = (
        f"Voice command heard: '{req.audio_text}'. Respond accordingly."
        if req.audio_text else "No voice command. React to visual input only."
    )
    user_text = build_nano_user_message(audio_ctx)

    content = [
        {"type": "text", "text": user_text},
        {"type": "image_url",
         "image_url": {"url": f"data:image/jpeg;base64,{req.image_b64}"}},
    ]
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": content},
    ]

    temp = 0.3 if any(
        w in directive.upper()
        for w in ["EVADE", "THREAT", "MOVE RIGHT", "MOVE LEFT", "MOVE BACK"]
    ) else 0.1
    t0 = time.perf_counter()
    raw = await state.nim.complete(
        model=state.config.nvidia.nano_omni_model,
        messages=messages,
        max_tokens=state.config.nvidia.max_tokens,
        temperature=temp,
    )
    latency_ms = (time.perf_counter() - t0) * 1000
    state.recent_latencies.append(latency_ms)
    state.total_infer_calls += 1

    if not raw:
        err = state.nim.last_error or "unknown error"
        return {"action": "idle", "magnitude": 0, "confidence": 0,
                "say": "",
                "scene_description": f"NIM error: {err}",
                "thinking": "",
                "error": err,
                "latency_ms": round(latency_ms, 1),
                "npc_role": req.npc_role}

    # frame dumping removed

    think, cleaned = _strip_thinking(raw)
    data = _extract_json(cleaned)

    response = {
        "action": data.get("action", "idle"),
        "target": data.get("target", ""),
        "direction": data.get("direction", "none"),
        "magnitude": float(data.get("magnitude", 0.5)),
        "confidence": float(data.get("confidence", 0.5)),
        "say": data.get("say", ""),
        "scene_description": data.get("scene_description", ""),
        "thinking": think[:500],
        "latency_ms": round(latency_ms, 1),
        "npc_role": req.npc_role,
        "ts": time.time(),
        "cached": False,
    }
    safe_action, passed = verifier.verify(response["action"], fsm[role_key].get_state())
    response["action"] = safe_action
    action_result = ActionResult(
        action=response["action"],
        target=response["target"],
        direction=response["direction"],
        magnitude=response["magnitude"],
        confidence=response["confidence"],
        say=response["say"],
        scene_description=response["scene_description"],
        capture_ts=t0,
        inference_ts=t0 + (latency_ms / 1000.0),
        inference_latency_ms=latency_ms,
    )
    state.metrics_logger.log_action(
        action_result,
        latency_tracker.get_stats_dict(),
        did_infer=should_infer,
        skip_reason=reason,
        npc_role=role,
    )
    cache["last_frame"] = frame_np
    cache["last_action"] = response
    cache["repeat_count"] = 0
    state.action_histogram[response["action"]] += 1
    state.recent_actions.append({k: response[k] for k in ("action", "say", "npc_role", "ts")})
    if think:
        state.recent_thinks.append(think[:200])
    fsm[role_key].tick()
    return response


@app.post("/super")
async def super_directive(req: SuperRequest):
    """Request a fresh strategic directive from the Super model."""
    if not state.api_key:
        raise HTTPException(500, "NVIDIA_NIM_API_KEY not configured on server")

    user_msg = build_super_user_message(
        req.scene_summary, req.mission_context, req.recent_actions
    )
    messages = [
        {"role": "system", "content": SUPER_SYSTEM_PROMPT},
        {"role": "user", "content": user_msg},
    ]

    t0 = time.perf_counter()
    raw = await state.nim.complete(
        model=state.config.nvidia.super_model,
        messages=messages,
        max_tokens=2048,
        temperature=0.3,
    )
    latency_ms = (time.perf_counter() - t0) * 1000
    state.total_super_calls += 1

    think, cleaned = _strip_thinking(raw)
    data = _extract_json(cleaned)
    directive = data.get("directive", "Compete to collect coins.")
    state_name = str(data.get("state", "EXPLORE")).upper()
    for role in ("HOSTILE", "ALLY", "WANDERER"):
        fsm[role].transition(state_name)
    state.last_directive = directive
    if think:
        state.recent_thinks.append(f"[SUPER] {think[:400]}")

    return {
        "directive": directive,
        "priority": data.get("priority", "explore"),
        "reasoning": data.get("reasoning", ""),
        "thinking": think[:800],
        "latency_ms": round(latency_ms, 1),
    }


@app.post("/dialogue")
async def dialogue(req: DialogueRequest):
    """Generate a single line of dialogue between two NPCs."""
    if not state.api_key:
        raise HTTPException(500, "NVIDIA_NIM_API_KEY not configured on server")

    sys_prompt = (
        "You write one short line of dialogue spoken between two NPC characters "
        "in a coin-collecting game. Output ONLY a JSON object: "
        '{"speaker":"<a|b>","text":"<one line, max 12 words>"}.'
    )
    user_prompt = (
        f"NPC A is the {req.npc_a}. NPC B is the {req.npc_b}. "
        f"Scene: {req.scene_summary or 'they meet near the coins'}. "
        f"Recent event: {req.last_event or 'none'}. Output the JSON now."
    )
    messages = [
        {"role": "system", "content": sys_prompt},
        {"role": "user", "content": user_prompt},
    ]

    t0 = time.perf_counter()
    raw = await state.nim.complete(
        model=state.config.nvidia.super_model,
        messages=messages,
        max_tokens=120,
        temperature=0.7,
    )
    latency_ms = (time.perf_counter() - t0) * 1000
    state.total_dialogue_calls += 1

    _, cleaned = _strip_thinking(raw)
    data = _extract_json(cleaned)
    return {
        "speaker": data.get("speaker", "a"),
        "text": data.get("text", "..."),
        "latency_ms": round(latency_ms, 1),
    }


# ─── Startup: pre-warm NIM connection ────────────────────────────────────────

@app.on_event("startup")
async def _prewarm_nim():
    """Establish TLS + TCP to NIM before first user request (eliminates cold-start latency)."""
    if state.metrics_logger is None:
        state.metrics_logger = MetricsLogger(state.config)
    if not state.api_key:
        logger.warning("Skipping NIM pre-warm — no API key.")
        return
    logger.info("Pre-warming NIM connection...")
    try:
        await state.nim.complete(
            model=state.config.nvidia.nano_omni_model,
            messages=[{"role": "user", "content": "ping"}],
            max_tokens=5,
            temperature=0.1,
        )
        logger.info("NIM pre-warm complete — first real request will be fast.")
    except Exception as exc:
        logger.warning("NIM pre-warm failed (non-fatal): %s", exc)


@app.on_event("shutdown")
async def _shutdown_metrics():
    if state.metrics_logger is not None:
        state.metrics_logger.close()


# ─── Static HTML hosting (so HF Spaces serves both API and game) ────────────
# Mount at "/" so index.html's `./game.js` reference resolves correctly.
# Explicit API routes (above) take precedence over the static catch-all.

_HTML_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "html_game")
if os.path.isdir(_HTML_DIR):
    app.mount("/", StaticFiles(directory=_HTML_DIR, html=True), name="game")


# ─── Entry point ──────────────────────────────────────────────────────────────

def main() -> None:
    parser = argparse.ArgumentParser(description="OmniNPC API Server")
    parser.add_argument("--port", type=int, default=int(os.getenv("PORT", "8000")))
    parser.add_argument("--host", type=str, default=os.getenv("HOST", "0.0.0.0"))
    args = parser.parse_args()

    if not state.api_key:
        logger.warning("NVIDIA_NIM_API_KEY not set — /infer routes will return 500.")
    else:
        logger.info("API key OK — Nano=%s Super=%s",
                    state.config.nvidia.nano_omni_model,
                    state.config.nvidia.super_model)

    import uvicorn
    uvicorn.run(app, host=args.host, port=args.port, log_level="info")


if __name__ == "__main__":
    main()