"""S10 browser sliders with one CPU environment and bounded JPEG storage."""
import argparse
import gc
from io import BytesIO
import os
from pathlib import Path
import threading
import time

from play_keyboard import DEFAULT_CHECKPOINT, keyboard_command_cfg
from s10_supermimic.core.gui import ControlState, make_server


def main():
  parser = argparse.ArgumentParser(description=__doc__)
  parser.add_argument("--checkpoint", type=Path, default=DEFAULT_CHECKPOINT)
  parser.add_argument("--port", type=int, default=8765)
  parser.add_argument("--device", default="cpu")
  parser.add_argument("--terrain-kind", choices=("flat", "up", "down", "platform", "slope", "rubble"), default="rubble")
  parser.add_argument("--terrain-level", type=int, choices=range(6), default=0)
  parser.add_argument("--fps", type=int, default=50)
  parser.add_argument("--shadows", action="store_true", help="Enable costly shadows/reflections")
  parser.add_argument("--ghost-horizon", type=int, default=32,
                      help="World-model future steps drawn as ghost poses")
  parser.add_argument("--ghost-stride", type=int, default=4,
                      help="Draw every Nth ghost (1 = all, 4 = 8 ghosts)")
  parser.add_argument("--no-ghost", action="store_true",
                      help="Disable the world-model ghost/error overlay")
  parser.add_argument("--steps", type=int, default=0, help="Exit after N simulation steps; 0 means until Ctrl+C")
  args = parser.parse_args()
  if not args.checkpoint.is_file():
    parser.error(f"Checkpoint not found: {args.checkpoint}")
  if not 1 <= args.fps <= 50 or not 1 <= args.port <= 65535 or args.steps < 0:
    parser.error("Require fps 1..50, port 1..65535 and nonnegative steps")
  if args.ghost_horizon < 1 or args.ghost_stride < 1:
    parser.error("ghost-horizon and ghost-stride must be positive")
  os.environ.setdefault("MUJOCO_GL", "egl")
  os.environ.setdefault("OMP_NUM_THREADS", "1")
  import torch
  from PIL import Image
  from mjlab.envs import ManagerBasedRlEnv
  from s10_supermimic.core.bridge import S10Bridge
  from s10_supermimic.core.policy import load_actor, evaluation_env_cfg
  from s10_supermimic.core.world_preview import WorldModelPreview

  torch.set_num_threads(1)
  state = ControlState()
  server = make_server(state, args.port)
  env = None
  worker = None
  try:
    checkpoint = torch.load(args.checkpoint, map_location=args.device, weights_only=True)
    actor = load_actor(checkpoint, args.device)
    # Keep the world model so the play view can draw the predicted future.
    checkpoint = {k: checkpoint[k] for k in
                  ("format", "history", "config", "env_action_order", "dt", "step",
                   "world_model")}
    checkpoint["actor_wheel_vel"] = actor.actor_wheel_vel
    def build_scene(kind, level):
      cfg = evaluation_env_cfg(checkpoint, level, kind)
      cfg.scene.num_envs = 1
      # Playback only needs the actor and the world-model preview; training
      # rewards/scans are not consumed here.
      cfg.rewards = {}
      cfg.metrics = {}
      cfg.observations.pop("critic", None)
      cfg.episode_length_s = 1e9
      cfg.commands["twist"] = keyboard_command_cfg(state)
      cfg.viewer.width, cfg.viewer.height = 640, 480
      cfg.viewer.enable_shadows = args.shadows
      cfg.viewer.enable_reflections = args.shadows
      cfg.viewer.distance, cfg.viewer.elevation, cfg.viewer.azimuth = 2.5, -20., 140.
      new_env = ManagerBasedRlEnv(cfg=cfg, device=args.device, render_mode="rgb_array")
      if checkpoint["format"] == "s10-terrain-policy-v1":
        from s10_supermimic.core.terrain import TerrainBridge, TerrainField
        field = TerrainField.from_env(new_env)
        bridge = TerrainBridge(new_env, field)
      else:
        field = None
        bridge = S10Bridge(new_env)
      if checkpoint["env_action_order"] != bridge.to_env or abs(checkpoint["dt"] - new_env.step_dt) > 1e-8:
        new_env.close()
        raise ValueError("Checkpoint action order/control timestep mismatch")
      obs, _ = new_env.reset()
      preview = None
      if not args.no_ghost:
        try:
          preview = WorldModelPreview(
            checkpoint, new_env, bridge, field, actor,
            horizon=args.ghost_horizon, device=args.device,
            env_action_order=checkpoint["env_action_order"])
        except Exception as exc:
          print(f"World-model preview disabled: {exc}", flush=True)
      return new_env, bridge, obs, preview

    kind, level = args.terrain_kind, args.terrain_level
    env, bridge, obs, preview = build_scene(kind, level)
    metadata = dict(checkpoint=args.checkpoint.name, iteration=checkpoint["step"],
                    terrain=checkpoint["format"] == "s10-terrain-policy-v1", device=args.device,
                    terrain_revision=checkpoint["config"].get("terrain_revision", 1))
    metadata.update(scene=f"{kind} · 等级 {level}", terrain_kind=kind, terrain_level=level)
    state.publish([0., 0., 0.], 0, **metadata)
    worker = threading.Thread(target=server.serve_forever, daemon=True)
    worker.start()
    print(f"S10 GUI ready: http://127.0.0.1:{server.server_port}/", flush=True)
    print(f"One environment on {args.device}. Ctrl+C exits. Commands expire after 2 seconds without a browser heartbeat.", flush=True)
    step, next_render = 0, 0.
    stats_since = time.monotonic()
    stats_step, frames = 0, 0
    render_fps = realtime_factor = 0.
    previous = torch.zeros(1, 16, device=args.device)
    ghosts = []
    with torch.inference_mode():
      while not args.steps or step < args.steps:
        began = time.monotonic()
        selection = state.take_scene()
        if selection is not None:
          # Release the old simulator before constructing the new one.
          state.clear()
          env.close()
          del bridge, obs
          d = None
          env = None
          gc.collect()
          if args.device.startswith("cuda"):
            torch.cuda.empty_cache()
          kind, level = selection
          env, bridge, obs, preview = build_scene(kind, level)
          previous.zero_()
          ghosts = []
          metadata.update(scene=f"{kind} · 等级 {level}", terrain_kind=kind, terrain_level=level)
          state.finish_scene()
          stats_since = next_render = time.monotonic()
          stats_step, frames = step, 0
        if state.take_reset():
          obs, _ = env.reset()
          previous.zero_()
          ghosts = []
          if preview is not None:
            preview.clear()
        with state.lock:
          paused = state.paused
        if not paused:
          state_before, pose_before = bridge.state(), bridge.pose()
          if preview is not None:
            preview.update_errors(step, state_before, pose_before)
            ghosts = preview.predict(
              obs["actor"], previous, env.command_manager.get_command("twist"),
              bridge.context(), bridge.robot.data.encoder_bias[:, bridge.legs],
              state_before, pose_before)
            preview.push_errors(step, ghosts)
          action = actor(obs["actor"])
          if not torch.isfinite(action).all():
            raise RuntimeError("Non-finite policy action")
          obs, _, terminated, truncated, _ = env.step(bridge.env_action(action))
          previous = action
          if not torch.isfinite(obs["actor"]).all():
            raise RuntimeError("Non-finite observation")
          if (terminated | truncated).any():
            state.clear()
            previous.zero_()
            ghosts = []
            if preview is not None:
              preview.clear()
          step += 1
        d = bridge.robot.data
        actual = [d.root_link_lin_vel_b[0, 0].item(), d.root_link_lin_vel_b[0, 1].item(),
                  d.root_link_ang_vel_b[0, 2].item()]
        state.publish(actual, step, render_fps=render_fps, realtime_factor=realtime_factor,
                      wm_error=(preview.summary() if preview is not None else ""),
                      wm_horizon=(args.ghost_horizon if preview is not None else 0),
                      **metadata)
        if began >= next_render:
          # Only the simulation thread touches the native renderer camera.
          with state.lock:
            azimuth, elevation, distance = state.camera
          camera = env._offline_renderer._cam
          camera.azimuth, camera.elevation, camera.distance = azimuth, elevation, distance
          rgb = (preview.render(ghosts, stride=args.ghost_stride)
                 if preview is not None and ghosts else env.render())
          buffer = BytesIO()
          Image.fromarray(rgb).save(buffer, format="JPEG", quality=80)
          with state.lock:
            state.jpeg = buffer.getvalue()
          frames += 1
          # Schedule from frame start, not completion: rendering time is part
          # of the frame budget rather than an extra delay on every frame.
          next_render += 1 / args.fps
          if next_render < began:
            next_render = began + 1 / args.fps
        now = time.monotonic()
        if now - stats_since >= 1.:
          render_fps = frames / (now - stats_since)
          realtime_factor = (step - stats_step) * env.step_dt / (now - stats_since)
          frames, stats_step, stats_since = 0, step, now
        time.sleep(max(0., env.step_dt - (time.monotonic() - began)))
    print(f"Render: {render_fps:.1f} FPS | realtime: {realtime_factor:.2f}x", flush=True)
  except KeyboardInterrupt:
    print("Closing S10 GUI.", flush=True)
  finally:
    state.clear()
    if worker is not None:
      server.shutdown()
      worker.join(timeout=2.)
    server.server_close()
    if env is not None:
      env.close()


if __name__ == "__main__":
  main()