"""avatar E2E hub — realtime streaming v2.

The browser connects to OpenAI Realtime directly over WebRTC. The assistant's
remote audio is captured in the browser in parallel with delayed live playback,
then sent here as 0.8s PCM16 chunks. A persistent body worker and the warm
causal face model return time-aligned frames before that delayed audio is heard.

  GET  /                       -> conversation.html
  POST /api/session            -> mint OpenAI Realtime ephemeral token (server key stays here)
  POST /api/warmup             -> warm face + body once, before conversation starts
  POST /api/stream/start       -> create stateful face/body stream
  POST /api/stream/chunk       -> PCM16 chunk -> newest face/body frames
  POST /api/stream/end         -> flush final short chunk and release stream
  POST /api/infer_turn {audio_b64, emotion, id?}
                               -> render_turn(wav, emotion) -> {id, dir, meta}; rebuilds turns index
  GET  /outputs/turns/<id>/*   -> the produced face.json / body.vrma / audio.wav (viewer reads these)

    OPENAI_API_KEY=sk-... python server.py [--port 8318]

Needs OPENAI_API_KEY only for /api/session; /api/infer_turn works without it (offline test).
"""
import argparse
import base64
import glob
import json
import os
import re
import threading
import time
import uuid
from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer
from urllib.parse import unquote, urlsplit

import numpy as np
import requests

HERE = os.path.dirname(os.path.abspath(__file__))
PSEUDO_NILOTOON_SRC = os.path.realpath(os.path.join(
    HERE, "..", "shader", "pseudo-nilotoon", "npm-package", "src"))
PSEUDO_NILOTOON_CORE_SRC = os.path.realpath(os.path.join(
    HERE, "..", "shader", "pseudo-nilotoon", "src"))
PSEUDO_NILOTOON_JSM = os.path.realpath(os.path.join(
    HERE, "..", "shader", "pseudo-nilotoon", "vendor", "jsm"))
PSEUDO_NILOTOON_VENDOR = os.path.realpath(os.path.join(
    HERE, "..", "shader", "pseudo-nilotoon", "vendor"))
PSEUDO_NILOTOON_ASSETS = os.path.realpath(os.path.join(
    HERE, "..", "shader", "pseudo-nilotoon", "assets"))
IDLE_VRMA_PATH = os.path.realpath(os.path.join(
    HERE, "..", "lipsync-wasm", "v2", "assets", "idle01.vrma"))
TURNS_DIR = os.path.join(HERE, "outputs", "turns")
MIC_DIR = os.path.join(HERE, "outputs", "mic")


def _load_dotenv(path):
    """Minimal .env loader (no dependency): KEY=VALUE per line, # comments, optional
    quotes. Does NOT override vars already in the environment (real env wins). The
    .env holds the OpenAI secret and is git-ignored — never log its values."""
    try:
        with open(path) as f:
            for line in f:
                line = line.strip()
                if not line or line.startswith("#") or "=" not in line:
                    continue
                k, v = line.split("=", 1)
                os.environ.setdefault(k.strip(), v.strip().strip('"').strip("'"))
    except FileNotFoundError:
        pass


_load_dotenv(os.path.join(HERE, ".env"))          # loads OPENAI_API_KEY (+ optional REALTIME_*)

# 16 emotions — must match face_v3_infer EMOTION_LABELS and the body tag set.
EMOTIONS = ["neutral", "joy", "laughter", "excitement", "agreement", "gratitude",
            "sadness", "crying", "sulk", "apology", "struggle", "anger", "refusal",
            "surprise", "fluster", "shy"]

REALTIME_MODEL = os.environ.get("REALTIME_MODEL", "gpt-realtime")
REALTIME_VOICE = os.environ.get("REALTIME_VOICE", "marin")
INSTRUCTIONS = (
    "당신은 친근한 한국어 아바타입니다. 답변은 짧고 대화체로. "
    "매 사용자 발화에 대한 첫 응답에서는 음성을 생성하지 말고 set_emotion 도구를 "
    "정확히 한 번 호출해 다음 발화의 톤에 가장 맞는 감정 하나를 넘기세요. "
    "도구 결과를 받은 뒤 이어지는 응답에서만 짧게 한 번 말하세요."
)
SET_EMOTION_TOOL = {
    "type": "function",
    "name": "set_emotion",
    "description": "Before speaking, report the single emotion that best matches the coming reply.",
    "parameters": {
        "type": "object",
        "properties": {"emotion": {"type": "string", "enum": EMOTIONS}},
        "required": ["emotion"],
    },
}

_lock = threading.Lock()                 # one GPU, one inference command at a time
_infer = None                            # lazy module import (heavy: torch + face model)
_streams = {}
_streams_lock = threading.Lock()


def inference_module():
    global _infer
    if _infer is None:
        import infer_turn
        _infer = infer_turn
    return _infer


def render_turn(*a, **kw):
    return inference_module().render_turn(*a, **kw)


def _decode_pcm16(value):
    if not value:
        return np.zeros(0, dtype=np.float32)
    raw = base64.b64decode(value, validate=True)
    if len(raw) % 2:
        raise ValueError("PCM16 payload has an odd byte length")
    return np.frombuffer(raw, dtype="<i2").astype(np.float32) / 32768.0


def _clean_id(value, prefix="stream"):
    value = value or f"{prefix}_{uuid.uuid4().hex[:12]}"
    return re.sub(r"[^\w-]+", "_", value)[:64]


def warm_models():
    with _lock:
        return inference_module().warm_models()


def _remove_stream(stream_id, state):
    """Remove only this state instance; a reused id must never be removed."""
    with _streams_lock:
        if _streams.get(stream_id) is state:
            _streams.pop(stream_id, None)


def _cancel_body_quietly(mod, stream_id):
    """Best-effort worker cleanup after a partially completed operation."""
    try:
        with _lock:
            mod.body_stream_cancel(stream_id)
    except Exception:
        pass


def start_stream(stream_id, emotion, guidance=3.0):
    mod = inference_module()
    state = {
        "face": None,
        "emotion": emotion,
        "created": time.time(),
        "status": "starting",
        "lock": threading.Lock(),
    }
    with _streams_lock:
        if stream_id in _streams:
            raise KeyError(f"stream {stream_id!r} already exists")
        # Reserve caller-selected ids before starting either model. Concurrent
        # requests for the same id will wait until initialization completes.
        _streams[stream_id] = state

    body_started = False
    with state["lock"]:
        try:
            with _lock:
                body_meta = mod.body_stream_start(
                    stream_id, emotion=emotion, guidance=guidance, emit_frames=24)
                body_started = True
                face_stream = mod.FaceStream(mod.get_face_model(), scenario_id=stream_id)
            state["face"] = face_stream
            state["status"] = "active"
        except Exception:
            state["status"] = "failed"
            if body_started:
                _cancel_body_quietly(mod, stream_id)
            _remove_stream(stream_id, state)
            raise

    return {"stream_id": stream_id, "fps": body_meta["fps"],
            "sample_rate": body_meta["sample_rate"],
            "emit_frames": body_meta["emit_frames"],
            "bones": body_meta["bones"],
            "rotation_space": body_meta.get(
                "rotation_space", "vrm_normalized_humanoid"),
            "expressions": list(mod.ARKIT_52_NAMES)}


def stream_frames(stream_id, pcm_b64, emotion=None, end=False):
    with _streams_lock:
        state = _streams.get(stream_id)
    if state is None:
        raise KeyError(f"unknown stream {stream_id!r}")

    # Decode before mutating stream state. A malformed request stays retryable.
    pcm = _decode_pcm16(pcm_b64)
    t0 = time.perf_counter()
    mod = inference_module()

    # BodyStream and FaceStream are both mutable. Serialize their complete
    # transaction per stream, including end/cancel lifecycle transitions.
    with state["lock"]:
        if state["status"] != "active":
            raise KeyError(f"stream {stream_id!r} is {state['status']}")
        if end:
            state["status"] = "closing"
        if emotion in EMOTIONS:
            state["emotion"] = emotion
        emotion = state["emotion"]

        try:
            with _lock:
                lock_acquired_at = time.perf_counter()
                if end:
                    body = mod.body_stream_end(
                        stream_id, pcm_b64=pcm_b64, emotion=emotion)
                else:
                    body = mod.body_stream_chunk(
                        stream_id, pcm_b64=pcm_b64, emotion=emotion)
                body_finished_at = time.perf_counter()
                frame_start = int(body.get("frame_start", 0))
                n_frames = int(body.get("frames", 0))
                face = state["face"].push_pcm(
                    pcm, emotion, n_frames, frame_start=frame_start)
                face_finished_at = time.perf_counter()

            packet = {
                "stream_id": stream_id,
                "frame_start": frame_start,
                "frames": n_frames,
                "body": body.get("body", []),
                "rotation_space": body.get(
                    "rotation_space", "vrm_normalized_humanoid"),
                "face": np.round(face, 4).tolist(),
                "emotion": emotion,
                "inference_ms": int(round((time.perf_counter() - t0) * 1000)),
                "queue_ms": int(round((lock_acquired_at - t0) * 1000)),
                "body_ms": int(round((body_finished_at - lock_acquired_at) * 1000)),
                "face_ms": int(round((face_finished_at - body_finished_at) * 1000)),
            }
        except Exception:
            # Body may already have advanced (and stream_end deletes it in the
            # worker) before face inference fails. Never leave a retryable-
            # looking split-brain stream in the server.
            state["status"] = "failed"
            _cancel_body_quietly(mod, stream_id)
            _remove_stream(stream_id, state)
            raise

        if end:
            state["status"] = "closed"
            _remove_stream(stream_id, state)
        return packet


def cancel_stream(stream_id):
    with _streams_lock:
        state = _streams.get(stream_id)
    if state is None:
        return {"stream_id": stream_id, "cancelled": False}

    mod = inference_module()
    with state["lock"]:
        if state["status"] in ("closed", "failed"):
            _remove_stream(stream_id, state)
            return {"stream_id": stream_id, "cancelled": False}
        state["status"] = "closing"
        try:
            with _lock:
                mod.body_stream_cancel(stream_id)
        finally:
            # If the worker died/restarted, keeping server state would only
            # turn the next chunk into an unknown-body failure.
            state["status"] = "closed"
            _remove_stream(stream_id, state)
    return {"stream_id": stream_id, "cancelled": True}


def rebuild_index():
    idx = []
    for meta in sorted(glob.glob(os.path.join(TURNS_DIR, "*", "meta.json")),
                       key=os.path.getmtime):
        m = json.load(open(meta)); m["dir"] = f"outputs/turns/{m['id']}"; idx.append(m)
    json.dump(idx, open(os.path.join(TURNS_DIR, "index.json"), "w"), ensure_ascii=False, indent=1)
    return idx


def mint_ephemeral_token():
    key = os.environ.get("OPENAI_API_KEY")
    if not key:
        raise RuntimeError("OPENAI_API_KEY not set on the server")
    body = {"session": {
        "type": "realtime", "model": REALTIME_MODEL,
        "audio": {"output": {"voice": REALTIME_VOICE}},
        "instructions": INSTRUCTIONS,
        "output_modalities": ["audio"],
        # The browser explicitly creates the spoken continuation with
        # tool_choice=none after acknowledging this required animation tag.
        "tools": [SET_EMOTION_TOOL], "tool_choice": "required",
    }}
    r = requests.post("https://api.openai.com/v1/realtime/client_secrets",
                      headers={"Authorization": f"Bearer {key}",
                               "Content-Type": "application/json"},
                      json=body, timeout=20)
    r.raise_for_status()
    d = r.json()
    # ephemeral secret has appeared as top-level {value,expires_at} or nested; be defensive.
    val = d.get("value") or (d.get("client_secret") or {}).get("value")
    if not val:
        raise RuntimeError(f"no ephemeral token in response: {json.dumps(d)[:300]}")
    return {"value": val, "expires_at": d.get("expires_at"),
            "model": REALTIME_MODEL, "emotions": EMOTIONS}


class Handler(SimpleHTTPRequestHandler):
    def __init__(self, *a, **kw):
        super().__init__(*a, directory=HERE, **kw)

    def translate_path(self, path):
        """Expose the canonical Pseudo NiloToon ESM sources without copying them.

        SimpleHTTPRequestHandler normally serves only HERE. The conversation
        viewer imports the shader from /pseudo-nilotoon/*, which maps to the
        sibling package's browser-ready npm source directory. Its one relaxed
        idle clip is exposed as an exact file route rather than mounting the
        surrounding lipsync package.
        """
        request_path = unquote(urlsplit(path).path)
        if request_path == "/idle01.vrma":
            return IDLE_VRMA_PATH
        mounts = (
            ("/pseudo-nilotoon-core/", PSEUDO_NILOTOON_CORE_SRC),
            ("/pseudo-nilotoon-jsm/", PSEUDO_NILOTOON_JSM),
            ("/pseudo-nilotoon-vendor/", PSEUDO_NILOTOON_VENDOR),
            ("/pseudo-nilotoon-assets/", PSEUDO_NILOTOON_ASSETS),
            ("/pseudo-nilotoon/", PSEUDO_NILOTOON_SRC),
        )
        for prefix, root in mounts:
            if not request_path.startswith(prefix):
                continue
            rel = request_path[len(prefix):].lstrip("/")
            candidate = os.path.realpath(os.path.join(root, rel))
            try:
                inside_root = os.path.commonpath(
                    (root, candidate)) == root
            except ValueError:
                inside_root = False
            if inside_root:
                return candidate
            return os.path.join(root, "__not_found__")
        return super().translate_path(path)

    def log_message(self, fmt, *args):
        if "/api/" in (self.path or ""):
            super().log_message(fmt, *args)

    def _json(self, code, obj):
        b = json.dumps(obj, ensure_ascii=False).encode()
        self.send_response(code)
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Content-Length", str(len(b)))
        self.end_headers()
        self.wfile.write(b)

    def do_POST(self):
        path = self.path.split("?")[0]
        try:
            n = int(self.headers.get("Content-Length", 0))
            req = json.loads(self.rfile.read(n)) if n else {}
        except Exception as e:
            return self._json(400, {"error": f"bad request body: {e}"})

        if path == "/api/session":
            try:
                return self._json(200, mint_ephemeral_token())
            except requests.HTTPError as e:
                return self._json(502, {"error": f"OpenAI token mint failed: {e.response.text[:300]}"})
            except Exception as e:
                return self._json(503, {"error": str(e)})

        if path == "/api/warmup":
            try:
                return self._json(200, {"ok": True, "timing": warm_models()})
            except Exception as e:
                return self._json(500, {"error": str(e)})

        if path == "/api/stream/start":
            emotion = req.get("emotion", "neutral")
            if emotion not in EMOTIONS:
                emotion = "neutral"
            stream_id = _clean_id(req.get("id"))
            try:
                return self._json(200, start_stream(
                    stream_id, emotion, guidance=float(req.get("guidance", 3.0))))
            except Exception as e:
                return self._json(500, {"error": str(e)})

        if path in ("/api/stream/chunk", "/api/stream/end"):
            stream_id = req.get("stream_id")
            if not stream_id:
                return self._json(400, {"error": "missing stream_id"})
            pcm_b64 = req.get("pcm_b64", "")
            if path.endswith("/chunk") and not pcm_b64:
                return self._json(400, {"error": "missing pcm_b64"})
            try:
                packet = stream_frames(stream_id, pcm_b64, emotion=req.get("emotion"),
                                       end=path.endswith("/end"))
                return self._json(200, packet)
            except KeyError as e:
                return self._json(404, {"error": str(e)})
            except Exception as e:
                return self._json(500, {"error": str(e)})

        if path == "/api/stream/cancel":
            stream_id = req.get("stream_id")
            if not stream_id:
                return self._json(400, {"error": "missing stream_id"})
            try:
                return self._json(200, cancel_stream(stream_id))
            except Exception as e:
                return self._json(500, {"error": str(e)})

        if path == "/api/infer_turn":
            emotion = req.get("emotion", "neutral")
            if emotion not in EMOTIONS:
                emotion = "neutral"                       # tolerate a stray tag rather than fail the turn
            b64 = req.get("audio_b64")
            if not b64:
                return self._json(400, {"error": "missing audio_b64"})
            os.makedirs(MIC_DIR, exist_ok=True)
            out_id = re.sub(r"\W+", "_", req.get("id") or f"turn_{int(time.time()*1000)}")[:48]
            wav_path = os.path.join(MIC_DIR, f"{out_id}.wav")
            with open(wav_path, "wb") as f:
                f.write(base64.b64decode(b64))
            try:
                with _lock:                               # serialise: one GPU
                    meta = render_turn(wav_path, emotion, out_id)
                rebuild_index()
                return self._json(200, {"id": out_id, "dir": f"outputs/turns/{out_id}", "meta": meta})
            except Exception as e:
                return self._json(500, {"error": str(e)})

        return self._json(404, {"error": "unknown endpoint"})


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--port", type=int, default=8318)
    a = ap.parse_args()
    os.makedirs(TURNS_DIR, exist_ok=True)
    rebuild_index()
    key = "set" if os.environ.get("OPENAI_API_KEY") else "NOT set (/api/session will 503)"
    print(f"serving {HERE} on http://127.0.0.1:{a.port}")
    print(f"  open   http://127.0.0.1:{a.port}/conversation.html")
    print(f"  model  {REALTIME_MODEL}  voice {REALTIME_VOICE}  OPENAI_API_KEY {key}")
    print(f"  body   {os.environ.get('AVATAR_BODY_PROFILE', 'v4')}")
    print("  mode   realtime v2 (24 frames / 0.8s per body chunk, 30fps output)")
    ThreadingHTTPServer(("127.0.0.1", a.port), Handler).serve_forever()


if __name__ == "__main__":
    main()