"""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 csv
import logging
import os
import threading
import time
from collections import Counter, deque
from datetime import datetime
from pathlib import Path
from typing import Optional

from dotenv import load_dotenv

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 config import load_config
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,
)


# ─── 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"


# ─── Session logger ───────────────────────────────────────────────────────────

class _SessionLogger:
    """Lightweight CSV logger - one row per /infer call."""

    _COLUMNS = ["timestamp", "npc_role", "action", "direction",
                "confidence", "total_latency_ms", "say"]

    def __init__(self, log_dir: str) -> None:
        self._lock = threading.Lock()
        Path(log_dir).mkdir(parents=True, exist_ok=True)
        ts = datetime.now().strftime("%Y%m%d_%H%M%S")
        self._path = Path(log_dir) / f"session_{ts}.csv"
        with open(self._path, "w", newline="", encoding="utf-8") as f:
            csv.DictWriter(f, self._COLUMNS).writeheader()
        logger.info("Session log: %s", self._path)

    def log(self, npc_role: str, action: str, direction: str,
            confidence: float, latency_ms: float, say: str) -> None:
        row = {
            "timestamp": datetime.now().isoformat(),
            "npc_role": npc_role,
            "action": action,
            "direction": direction,
            "confidence": f"{confidence:.3f}",
            "total_latency_ms": f"{latency_ms:.1f}",
            "say": (say or "")[:50],
        }
        with self._lock:
            with open(self._path, "a", newline="", encoding="utf-8") as f:
                csv.DictWriter(f, self._COLUMNS).writerow(row)


# ─── 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)
        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.session_log = _SessionLogger(self.config.metrics.log_dir)


state = ServerState()


# ─── 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):
    """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")

    persona = _role_persona(req.npc_role)
    directive = 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},
    ]

    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=0.1,
    )
    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}

    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(),
    }
    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])
    state.session_log.log(
        npc_role=req.npc_role,
        action=response["action"],
        direction=response["direction"],
        confidence=response["confidence"],
        latency_ms=latency_ms,
        say=response["say"],
    )
    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.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 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)


# ─── 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()