"""
ABot-World — Interactive Action-Conditioned World Rollout
gradio.Server + WebSocket live-backend edition.

Given an uploaded starting image (i2v conditioning), a scene prompt, and live
WASD / IJKL controls, the model autoregressively rolls out an action-conditioned
navigable world and streams decoded frames to the browser over a WebSocket.

This mirrors the live backend/infrastructure of
https://huggingface.co/spaces/Overworld/waypoint-1-5 (gradio.Server for
ZeroGPU-friendly start/stop + a raw WebSocket for real-time binary JPEG frame
streaming and control input), with a cleaner custom UI and image-upload seeding.

Multi-user safe: every endpoint is keyed by a per-client `session_id` so
concurrent players never share seed images, frame queues, or status messages.

ZeroGPU quota: the incoming request's ZeroGPU proxy token (the `x-ip-token` /
`x-api-token` header injected by the HF iframe) is captured per-session and
propagated into the worker thread's gradio request context, so the GPU work is
billed against the *requesting user's* quota — not the Space owner's.

Upstream: https://github.com/amap-cvlab/ABot-World
Model:    https://huggingface.co/acvlab/ABot-World-0-5B-LF  (built on Wan2.2-TI2V-5B)
"""
import os
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # must precede torch / CUDA-touching imports

import sys
import io
import time
import queue
import asyncio
import struct
import tempfile
import threading
import contextvars
import uuid
from collections import deque
from dataclasses import dataclass, field
from multiprocessing import Queue as MPQueue
from pathlib import Path
from typing import Dict, Optional, Set

import numpy as np
import torch
from PIL import Image
from omegaconf import OmegaConf
from huggingface_hub import snapshot_download

from fastapi import UploadFile, File, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse, JSONResponse
from gradio import Server
from gradio.context import LocalContext

# ── Repo paths ───────────────────────────────────────────────────────────────
APP_DIR = Path(__file__).resolve().parent
if str(APP_DIR) not in sys.path:
    sys.path.insert(0, str(APP_DIR))

MODEL_ID = "acvlab/ABot-World-0-5B-LF"
CKPT_DIR = APP_DIR / "checkpoints" / "ABot-World-0-5B-LF"

# ── Stream / rollout configuration ───────────────────────────────────────────
# 704x1280 is the native training resolution used by the upstream web client.
STREAM_HEIGHT = 704
STREAM_WIDTH = 1280
JPEG_QUALITY = 82
MAX_BLOCKS_PER_SESSION = 512     # hard cap so a session can't run forever
SESSION_IDLE_TIMEOUT = 600       # seconds; janitor reaps abandoned sessions
GPU_DURATION = 120               # seconds per @spaces.GPU allocation

# Actions map to the 8-key one-hot the model was trained on (W A S D I J K L).
# The browser sends the currently-held key set; we translate to this dict.
KEY_ORDER = ["W", "A", "S", "D", "I", "J", "K", "L"]

DEFAULT_PROMPT = (
    "A realistic outdoor world scene with a navigable path, natural lighting, "
    "detailed ground texture, and stable forward motion."
)

# ── Download weights (once, at startup) ──────────────────────────────────────
print(f"[startup] downloading {MODEL_ID} weights ...", flush=True)
snapshot_download(
    repo_id=MODEL_ID,
    repo_type="model",
    local_dir=str(CKPT_DIR),
)
print("[startup] weights downloaded.", flush=True)

os.environ["TAEW2_2_CHECKPOINT"] = str(CKPT_DIR / "taew2_2.pth")

# ── Build the causal streaming pipeline (module scope, eager to CUDA) ─────────
from pipeline import CausalInferencePipeline
from utils.misc import set_seed
from utils.wan_wrapper import create_vae_from_config

_CONFIG_PATH = APP_DIR / "configs" / "long_forcing_dmd.yaml"
_DEFAULT_CFG_PATH = APP_DIR / "configs" / "default_config.yaml"


def _build_config():
    cfg = OmegaConf.merge(
        OmegaConf.load(str(_DEFAULT_CFG_PATH)),
        OmegaConf.load(str(_CONFIG_PATH)),
    )
    cfg.taew2_2_checkpoint = str(CKPT_DIR / "taew2_2.pth")
    cfg.lightvae_encoder_checkpoint = str(CKPT_DIR / "Wan2.2_VAE.pth")
    cfg.model_kwargs.model_name = str(CKPT_DIR)
    cfg.text_encoder_kwargs.tokenizer_path = str(CKPT_DIR / "google" / "umt5-xxl") + "/"
    cfg.text_encoder_kwargs.encoder_pth_path = str(CKPT_DIR / "models_t5_umt5-xxl-enc-bf16.pth")
    cfg.vae_kwargs.pretrained_path = str(CKPT_DIR / "Wan2.2_VAE.pth")
    cfg.vae_type = "taew2_2"
    cfg.use_fp8_gemm = False
    return cfg


print("[startup] building pipeline ...", flush=True)
set_seed(42)
torch.set_grad_enabled(False)

CONFIG = _build_config()
_vae = create_vae_from_config(CONFIG)
pipeline = CausalInferencePipeline(CONFIG, device=torch.device("cuda"), vae=_vae)

try:
    from wan.modules.helios_kernels import (
        replace_all_norms_with_flash_norms,
        replace_rope_with_flash_rope,
    )
    replace_all_norms_with_flash_norms(pipeline.generator.model)
    replace_rope_with_flash_rope()
    print("[startup] helios Triton kernels enabled.", flush=True)
except Exception as e:  # pragma: no cover
    print(f"[startup] helios kernels disabled ({e!r}); using eager norms/rope.", flush=True)

pipeline = pipeline.to(dtype=torch.bfloat16)
pipeline.text_encoder.to(device="cuda")
pipeline.generator.to(device="cuda")
pipeline.vae.to(device="cuda")
if pipeline.encoder is not None:
    pipeline.encoder.to(device="cuda")
pipeline.torch_dtype = torch.bfloat16
print("[startup] pipeline ready.", flush=True)

NUM_FPB = int(getattr(CONFIG, "num_frame_per_block", 3))
_vae_for_shape = pipeline.encoder if pipeline.encoder is not None else pipeline.vae
LATENT_CHANNELS = _vae_for_shape.z_dim
UPSAMPLE = getattr(_vae_for_shape, "upsampling_factor", 16)
LATENT_H = STREAM_HEIGHT // UPSAMPLE
LATENT_W = STREAM_WIDTH // UPSAMPLE

# Only one rollout may touch the shared pipeline at a time.
_infer_lock = threading.Lock()


def _keys_from_buttons(buttons) -> Dict[str, bool]:
    """Translate a set of held key names (e.g. {'W','A'}) into the model dict."""
    held = {k.upper() for k in (buttons or [])}
    return {k: True for k in KEY_ORDER if k in held}


def _decode_block_to_frames(lat_block):
    frames = []

    class _Cap:
        def append_data(self, f):
            frames.append(np.asarray(f))

    pipeline.decode_block_and_write(lat_block, _Cap())
    return frames


# ── Command types (browser -> worker) ────────────────────────────────────────
@dataclass
class ControlCommand:
    buttons: Set[str]
    prompt: str


@dataclass
class StopCommand:
    pass


# ── Per-session state ────────────────────────────────────────────────────────
# NOTE on queues: the @spaces.GPU rollout runs in a forked subprocess, so any
# object it reads must cross the fork boundary. `command_queue` is therefore a
# multiprocessing Queue (browser controls / stop reach the GPU loop through it).
# `frame_queue` / `status_queue` are plain queue.Queue used only in the parent
# process (frames arrive back via the ZeroGPU generator IPC and are forwarded
# to the WebSocket by the worker thread).
@dataclass
class GameSession:
    session_id: str
    command_queue: "MPQueue"
    frame_queue: "queue.Queue"
    status_queue: "queue.Queue"
    stop_event: threading.Event
    seed_path: str
    prompt: str
    seed: int
    worker_thread: Optional[threading.Thread] = None
    frame_times: deque = field(default_factory=lambda: deque(maxlen=30))
    last_active: float = field(default_factory=time.time)

    def touch(self):
        self.last_active = time.time()

    def stop(self):
        self.stop_event.set()
        try:
            self.command_queue.put_nowait(StopCommand())
        except Exception:
            pass
        if self.worker_thread and self.worker_thread.is_alive():
            self.worker_thread.join(timeout=4.0)


_sessions: Dict[str, GameSession] = {}
_sessions_lock = threading.Lock()

# Contextvar carrying the active session's status queue (inherited by the worker
# thread via contextvars.copy_context()).
_current_status_queue: "contextvars.ContextVar[Optional[queue.Queue]]" = contextvars.ContextVar(
    "abot_status_queue", default=None
)


def broadcast_status(msg: str):
    q = _current_status_queue.get()
    if q is None:
        return
    try:
        q.put_nowait(msg)
    except queue.Full:
        pass


def _get_session(session_id: str) -> Optional[GameSession]:
    with _sessions_lock:
        return _sessions.get(session_id)


def _drop_session(session_id: str) -> Optional[GameSession]:
    with _sessions_lock:
        return _sessions.pop(session_id, None)


def _reap_idle_sessions():
    while True:
        time.sleep(60)
        now = time.time()
        to_drop = []
        with _sessions_lock:
            for sid, sess in list(_sessions.items()):
                worker_dead = sess.worker_thread is None or not sess.worker_thread.is_alive()
                idle = (now - sess.last_active) > SESSION_IDLE_TIMEOUT
                if worker_dead and idle:
                    to_drop.append(sid)
            for sid in to_drop:
                _sessions.pop(sid, None)
        if to_drop:
            print(f"Janitor reaped {len(to_drop)} idle session(s)", flush=True)


threading.Thread(target=_reap_idle_sessions, daemon=True).start()


# ── GPU worker ───────────────────────────────────────────────────────────────
def gpu_worker_thread(session: "GameSession"):
    """Parent-thread driver: consumes frames yielded by the ZeroGPU generator,
    computes FPS, and forwards frames to the WebSocket via `frame_queue`.

    Status/stop live in the parent process; the GPU loop is steered purely
    through the (picklable, cross-fork) `command_queue`.
    """
    try:
        broadcast_status("GPU allocated — starting world…")
        gen = create_gpu_rollout_loop(
            session.command_queue, session.seed_path, session.prompt, session.seed,
        )
        first = True
        while not session.stop_event.is_set():
            try:
                frame, frame_count = next(gen)
            except StopIteration:
                print("Rollout generator exhausted", flush=True)
                break
            except Exception as e:
                if "aborted" in str(e).lower() or "duration" in str(e).lower():
                    print(f"GPU time expired: {e}", flush=True)
                else:
                    print(f"Worker error: {e}", flush=True)
                    broadcast_status(f"error:{e}")
                break

            if first:
                broadcast_status("Rolling out — use WASD / IJKL to steer.")
                first = False

            now = time.time()
            session.frame_times.append(now)
            fps = 0.0
            if len(session.frame_times) >= 2:
                elapsed = session.frame_times[-1] - session.frame_times[0]
                fps = (len(session.frame_times) - 1) / elapsed if elapsed > 0 else 0.0
            # Keep only the freshest frame if the consumer fell behind.
            while session.frame_queue.qsize() > 2:
                try:
                    session.frame_queue.get_nowait()
                except queue.Empty:
                    break
            try:
                session.frame_queue.put_nowait((frame, frame_count, round(fps, 1)))
            except queue.Full:
                pass
    finally:
        session.stop_event.set()
        print("Worker thread finished", flush=True)


def create_gpu_rollout_loop(command_queue, seed_path, prompt_text, seed):
    """Return a ZeroGPU generator that rolls the world out block-by-block.

    Only picklable primitives + the multiprocessing `command_queue` cross the
    fork boundary. Live controls (held key set) and stop arrive via that queue.
    """
    @spaces.GPU(duration=GPU_DURATION)
    def gpu_rollout():
        device = torch.device("cuda")
        prompt = (prompt_text or DEFAULT_PROMPT).strip() or DEFAULT_PROMPT
        set_seed(int(seed))

        with _infer_lock:
            noise = torch.randn(
                [1, NUM_FPB, LATENT_CHANNELS, LATENT_H, LATENT_W],
                device=device, dtype=torch.bfloat16,
            )

            pipeline.set_prompts([prompt], device=device)
            pipeline.set_ref_latent_mask_from_exists_paths(
                ref_dir=str(APP_DIR / "__no_such_ref_dir__"), device=device,
            )
            pipeline.reset_stream(batch_size=1, dtype=torch.bfloat16,
                                  device=device, initial_latent=None)
            pipeline.set_first_frame_latent(
                seed_path, height=STREAM_HEIGHT, width=STREAM_WIDTH, device=device,
            )

            current_keys: Dict[str, bool] = {"W": True}  # default: forward
            stop_requested = False
            try:
                for b in range(MAX_BLOCKS_PER_SESSION):
                    # Drain control commands: newest held-key set wins.
                    while True:
                        try:
                            cmd = command_queue.get_nowait()
                        except Exception:
                            break
                        if isinstance(cmd, StopCommand):
                            stop_requested = True
                            break
                        if isinstance(cmd, ControlCommand):
                            current_keys = _keys_from_buttons(cmd.buttons)
                    if stop_requested:
                        break

                    pipeline.set_act(current_keys, height=STREAM_HEIGHT, width=STREAM_WIDTH,
                                     num_frames=NUM_FPB, device=device)
                    lat_block = pipeline.generate_next_block(noise)
                    noise = torch.randn_like(noise)
                    frames = _decode_block_to_frames(lat_block)
                    for f in frames:
                        yield (f, b)
            finally:
                try:
                    pipeline.reset_stream(batch_size=1, dtype=torch.bfloat16,
                                          device=device, initial_latent=None)
                except Exception:
                    pass
                try:
                    if hasattr(pipeline.vae, "model") and hasattr(pipeline.vae.model, "clear_cache"):
                        pipeline.vae.model.clear_cache()
                except Exception:
                    pass

    return gpu_rollout()


# ── App (gradio.Server) ──────────────────────────────────────────────────────
app = Server()


@app.api(name="start_game")
def start_game(session_id: str = "", seed_path: str = "",
               prompt: str = "", seed: int = 42) -> str:
    """Start a new interactive world rollout for `session_id`.

    Args:
        session_id: per-client id (UUID) isolating this player's stream.
        seed_path: filepath (uploaded via /upload) of the starting frame image
            that seeds the i2v world rollout.
        prompt: scene description.
        seed: RNG seed for reproducibility.

    Returns:
        The session_id actually used.
    """
    if not session_id:
        session_id = str(uuid.uuid4())

    prior = _drop_session(session_id)
    if prior is not None:
        prior.stop()

    if not seed_path:
        raise ValueError("A starting image is required — please upload one first.")

    command_queue = MPQueue()  # crosses the ZeroGPU fork boundary
    frame_queue: "queue.Queue" = queue.Queue(maxsize=4)
    status_queue: "queue.Queue" = queue.Queue(maxsize=32)
    stop_event = threading.Event()

    session = GameSession(
        session_id=session_id,
        command_queue=command_queue,
        frame_queue=frame_queue,
        status_queue=status_queue,
        stop_event=stop_event,
        seed_path=seed_path,
        prompt=prompt or DEFAULT_PROMPT,
        seed=int(seed),
    )
    with _sessions_lock:
        _sessions[session_id] = session

    # Capture the *incoming request* — HF has already injected this user's
    # ZeroGPU proxy token (x-ip-token / x-api-token) into its headers. We
    # re-set it into the worker thread's gradio LocalContext so that
    # @spaces.GPU bills GPU time against THIS user's quota, not the owner's.
    gradio_request = LocalContext.request.get(None)
    status_token = _current_status_queue.set(status_queue)
    try:
        broadcast_status("Requesting GPU from ZeroGPU…")

        def _thread_entry():
            # Re-establish the request context inside the worker thread so the
            # ZeroGPU scheduler reads the requesting user's token.
            if gradio_request is not None:
                try:
                    LocalContext.request.set(gradio_request)
                except Exception:
                    pass
            gpu_worker_thread(session)

        ctx = contextvars.copy_context()
        worker = threading.Thread(target=ctx.run, args=(_thread_entry,), daemon=True)
        session.worker_thread = worker
        worker.start()
    finally:
        _current_status_queue.reset(status_token)

    return session_id


@app.api(name="stop_game")
def stop_game(session_id: str = "") -> str:
    """Stop the active rollout for the given client."""
    if not session_id:
        return "no_session"
    session = _drop_session(session_id)
    if session is not None:
        session.stop()
    return "stopped"


@app.websocket("/ws")
async def game_ws(websocket: WebSocket, session_id: str = ""):
    """Real-time rollout WebSocket. Requires `?session_id=...` matching /start_game."""
    await websocket.accept()
    if not session_id:
        await websocket.send_json({"type": "error", "message": "missing session_id"})
        await websocket.close(code=1008)
        return

    loop = asyncio.get_event_loop()

    async def send_frames():
        session_ended_sent = False
        while True:
            session = _get_session(session_id)

            if session is not None:
                try:
                    status_msg = session.status_queue.get_nowait()
                    if status_msg.startswith("error:"):
                        await websocket.send_json({"type": "error", "message": status_msg[6:]})
                        break
                    await websocket.send_json({"type": "status", "message": status_msg})
                except queue.Empty:
                    pass
                except (WebSocketDisconnect, RuntimeError):
                    break

            if session is None:
                await asyncio.sleep(0.05)
                continue
            if session.stop_event.is_set() and session.frame_queue.empty():
                if not session_ended_sent:
                    try:
                        await websocket.send_json({"type": "session_ended"})
                    except (WebSocketDisconnect, RuntimeError):
                        break
                    session_ended_sent = True
                await asyncio.sleep(0.4)
                continue
            try:
                result = await loop.run_in_executor(
                    None, lambda s=session: s.frame_queue.get(timeout=0.1)
                )
                frame, count, fps = result
                img = Image.fromarray(frame)
                buf = io.BytesIO()
                img.save(buf, format="JPEG", quality=JPEG_QUALITY)
                jpeg_bytes = buf.getvalue()
                header = struct.pack(">II", int(count), int(fps * 10))
                await websocket.send_bytes(header + jpeg_bytes)
                session.touch()
            except queue.Empty:
                pass
            except (WebSocketDisconnect, RuntimeError):
                break

    async def receive_controls():
        while True:
            try:
                data = await websocket.receive_json()
                session = _get_session(session_id)
                if session is None:
                    continue
                session.touch()
                msg_type = data.get("type", "control")
                if msg_type == "control":
                    buttons = set(data.get("buttons", []))
                    prompt = data.get("prompt", session.prompt)
                    try:
                        session.command_queue.put_nowait(
                            ControlCommand(buttons=buttons, prompt=prompt)
                        )
                    except queue.Full:
                        pass
                elif msg_type == "stop":
                    session.stop()
            except WebSocketDisconnect:
                break
            except Exception:
                break

    try:
        await asyncio.gather(send_frames(), receive_controls())
    except WebSocketDisconnect:
        pass


@app.post("/upload_seed")
async def upload_seed(file: UploadFile = File(...)):
    """Accept a user-uploaded starting image and stash it server-side.

    Returns the temp filepath, which the browser then passes to /start_game as
    `seed_path` to seed the image-to-video (i2v) world rollout. Only an image is
    needed — there is no video upload.
    """
    try:
        raw = await file.read()
        img = Image.open(io.BytesIO(raw)).convert("RGB")
    except Exception:
        return JSONResponse({"error": "Could not read image file."}, status_code=400)

    tmp = tempfile.NamedTemporaryFile(prefix="abot_seed_", suffix=".png", delete=False)
    img.save(tmp.name, format="PNG")
    return {"seed_path": tmp.name}


@app.get("/", response_class=HTMLResponse)
async def homepage():
    html_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "index.html")
    with open(html_path, "r", encoding="utf-8") as f:
        return f.read()


# Avoid ZeroGPU "no GPU function" error at boot.
spaces.GPU(lambda: None)

app.launch(server_name="0.0.0.0", server_port=7860)