"""Persistent warm body worker — one process, model loaded ONCE, jobs over stdin.

Why a separate process (not just an in-process BodyEngine): the V3 face code and the
GVRM body code both use top-level package names `models` / `scripts` / `dataloaders`,
which collide in a single interpreter (whichever imports first wins → the other breaks
with e.g. `No module named 'models.vrm_lsm'`). Isolating the body engine in its own
process gives it its own module namespace (GVRM's), while the server keeps the V3 face
model warm in-process. Both stay warm.

Protocol (line-based JSON jobs, stdout sentinels):
  emits `@@READY@@` once the model is loaded,
  then for each stdin job emits
  `@@RESULT@@ {"ok":true,"frames":N,"vrma":...}` (or ok:false,error). Non-sentinel
  stdout/stderr lines are logs and are ignored by the parent.

Supported operations:
  generate      batch wav -> vrma (the original warm turn path)
  stream_start  allocate a stateful BodyStream
  stream_chunk  feed raw little-endian PCM16 and return new quaternion frames
  stream_end    flush the last short PCM remainder and release the stream

    python body_worker.py --ckpt ... --vqvae_dir ... --meanstd_dir ... --npz ... --proto_seed ...
"""
import argparse
import base64
import json
import os
import sys

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from body_engine import BodyEngine


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", required=True)
    ap.add_argument("--vqvae_dir", required=True)
    ap.add_argument("--meanstd_dir", required=True)
    ap.add_argument("--npz", required=True)
    ap.add_argument("--proto_seed", default=None)
    ap.add_argument("--steps", type=int, default=8)
    ap.add_argument("--model_kind", choices=("tag3", "base"), default="tag3")
    ap.add_argument("--codec_selection", default=None)
    ap.add_argument("--continuous_decode", action="store_true")
    ap.add_argument("--template_vrma", default=None,
                    help="베이크 템플릿 vrma — 주면 그 구조/rest로 발행")
    a = ap.parse_args()

    eng = BodyEngine(ckpt=a.ckpt, vqvae_dir=a.vqvae_dir, meanstd_dir=a.meanstd_dir,
                     npz=a.npz, proto_seed=a.proto_seed, fingers=True, steps=a.steps,
                     model_kind=a.model_kind, codec_selection=a.codec_selection,
                     continuous_decode=a.continuous_decode,
                     template_vrma=a.template_vrma)
    streams = {}
    print("@@READY@@", flush=True)

    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        try:
            job = json.loads(line)
            op = job.get("op", "generate")
            if op == "ping":
                result = {"ok": True, "op": op}
            elif op == "generate":
                q, T, root = eng.generate_take(
                    job["wav"], job["emotion"],
                    guidance=float(job.get("guidance", 3.0)))
                eng.write_vrma(q, job["out_vrma"], root_translation=root)
                result = {"ok": True, "op": op, "frames": int(T),
                          "vrma": job["out_vrma"]}
            elif op == "stream_start":
                stream_id = job["stream_id"]
                streams[stream_id] = eng.new_stream(
                    job.get("emotion", "neutral"),
                    style=job.get("style", "general"),
                    guidance=float(job.get("guidance", 3.0)),
                    emit_frames=int(job.get("emit_frames", eng.PRE * 4)),
                )
                result = {"ok": True, "op": op, "stream_id": stream_id,
                          "fps": 30, "sample_rate": 16000,
                          "emit_frames": streams[stream_id].emit,
                          # Contract: streamed body values are loader-equivalent
                          # normalized-humanoid quaternions, not source-local basis.
                          # The source VRMA has no mapped upperChest, so it is omitted
                          # here exactly as it is in batch VRMA playback.
                          "bones": list(eng.stream_bone_names),
                          "rotation_space": "vrm_normalized_humanoid"}
            elif op in ("stream_chunk", "stream_end"):
                stream_id = job["stream_id"]
                if stream_id not in streams:
                    raise KeyError(f"unknown body stream {stream_id!r}")
                stream = streams[stream_id]
                if job.get("emotion"):
                    stream.set_emotion(job["emotion"])
                pcm = _decode_pcm16(job.get("pcm_b64", ""))
                frame_start = stream.frames_emitted
                # GestureVRM emits source-rig local basis rotations. Batch VRMA
                # playback reconstructs source channels and the animation loader
                # conjugates them by each bone's source world-rest rotation. Do the
                # same conversion here because streaming bypasses that loader.
                q_basis = stream.push_pcm(
                    pcm, flush=(op == "stream_end" or bool(job.get("flush"))))
                q = eng.basis_to_humanoid(q_basis, mapped_only=True)
                result = {"ok": True, "op": op, "stream_id": stream_id,
                          "frame_start": int(frame_start), "frames": int(len(q)),
                          "body": np.round(q, 5).tolist(),
                          "rotation_space": "vrm_normalized_humanoid"}
                if op == "stream_end":
                    del streams[stream_id]
            elif op == "stream_cancel":
                streams.pop(job["stream_id"], None)
                result = {"ok": True, "op": op, "stream_id": job["stream_id"]}
            else:
                raise ValueError(f"unknown body worker op {op!r}")
            print("@@RESULT@@ " + json.dumps(result, separators=(",", ":")), flush=True)
        except Exception as e:                                   # noqa: BLE001
            print("@@RESULT@@ " + json.dumps({"ok": False, "error": str(e)}), flush=True)


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


if __name__ == "__main__":
    main()