from __future__ import annotations

import base64
import importlib
import importlib.util
import socket
import logging
import webbrowser
import os
import shutil
import sys
import threading
import time
import traceback
from collections import defaultdict
from io import BytesIO
from pathlib import Path
from typing import Any

import cv2
import numpy as np
import asyncio
import json
import torch
from fastapi import FastAPI, HTTPException, Query, Body, Request
from fastapi.responses import StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles
from PIL import Image
from pydantic import BaseModel, Field

# Support both module paths:
# - world_models.ui.server
# - torchwm.world_models.ui.server
if "world_models" not in sys.modules:
    package_name = __package__ or ""
    if package_name.endswith("world_models.ui"):
        maybe_world_models_pkg = package_name.rsplit(".ui", 1)[0]
        if maybe_world_models_pkg != "world_models":
            try:
                sys.modules["world_models"] = importlib.import_module(
                    maybe_world_models_pkg
                )
            except Exception:
                pass

from ..configs.dreamer_config import DreamerConfig
from ..configs.iris_config import IRISConfig
from ..envs import list_available_atari_envs
from ..models.dreamer import DreamerAgent
from ..models.planet import Planet
from ..training.train_planet import train as planet_train
from ..training.train_iris import IRISTrainer
from ..utils.utils import flatten_dict, visualize_latent_tsne, visualize_latent_umap

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Global storage for last latents
last_latents = None


SUPPORTED_MODELS: dict[str, dict[str, str]] = {
    "dreamerv1": {
        "label": "DreamerV1",
        "description": "DreamerV1 - Model-based RL agent with latent world model.",
    },
    "dreamerv2": {
        "label": "DreamerV2",
        "description": "DreamerV2 - Improved Dreamer with discrete latent variables.",
    },
    "planet": {
        "label": "Planet",
        "description": "PlaNet-style recurrent state-space world model with MPC planning.",
    },
    "iris": {
        "label": "IRIS",
        "description": "IRIS - Transformers are Sample-Efficient World Models",
    },
}

DREAMER_ENVS = [
    "cartpole-balance",
    "cartpole-swingup",
    "cheetah-run",
    "finger-spin",
    "reacher-easy",
    "walker-walk",
    "walker-run",
    "quadruped-walk",
]

PLANET_BASE_ENVS = [
    "CartPole-v1",
    "Pendulum-v1",
    "MountainCarContinuous-v0",
    "Acrobot-v1",
    "HalfCheetah-v4",
    "Humanoid-v4",
]

GYM_ENVS = [
    "CartPole-v1",
    "Pendulum-v1",
    "MountainCarContinuous-v0",
    "Acrobot-v1",
    "HalfCheetah-v4",
    "Humanoid-v4",
    "Hopper-v4",
    "Swimmer-v4",
    "Walker2d-v4",
    "Ant-v4",
    "Reacher-v4",
    "Pusher-v4",
    "Manipulator-v4",
    "LunarLander-v3",
    "LunarLanderContinuous-v3",
    "BipedalWalker-v3",
    "BipedalWalkerHardcore-v3",
    "CarRacing-v3",
    "Blackjack-v1",
    "FrozenLake-v1",
    "FrozenLake8x8-v1",
    "Taxi-v3",
    "InvertedPendulum-v4",
    "InvertedDoublePendulum-v4",
    "HalfCheetah-v2",
    "Hopper-v2",
    "Swimmer-v2",
    "Walker2d-v2",
    "Reacher-v2",
    "Pusher-v2",
]

UNITY_ENVS = []

ATARI_ENVS: list[str] = []
try:
    ATARI_ENVS = list_available_atari_envs()
except Exception:
    ATARI_ENVS = []

ENV_BACKENDS: dict[str, dict[str, Any]] = {
    "dm_control": {
        "label": "DM Control",
        "description": "DeepMind Control Suite",
        "environments": DREAMER_ENVS,
    },
    "mujoco": {
        "label": "MuJoCo",
        "description": "MuJoCo physics environments",
        "environments": GYM_ENVS,
    },
    "gym": {
        "label": "Gym",
        "description": "OpenAI Gym environments",
        "environments": PLANET_BASE_ENVS,
    },
    "unity": {
        "label": "Unity ML Agents",
        "description": "Unity ML Agents environments",
        "environments": UNITY_ENVS,
    },
    "atari": {
        "label": "Atari",
        "description": "Atari 2600 environments via ALE",
        "environments": ATARI_ENVS,
    },
}

DEFAULT_MODEL_CONFIGS: dict[str, dict[str, Any]] = {
    "dreamerv1": {
        "seed": 1,
        "action_repeat": 2,
        "batch_size": 50,
        "train_seq_len": 50,
        "algo": "Dreamerv1",
        "env_backend": "dmc",
    },
    "dreamerv2": {
        "seed": 1,
        "action_repeat": 2,
        "batch_size": 50,
        "train_seq_len": 50,
        "algo": "Dreamerv2",
        "env_backend": "dmc",
    },
    "planet": {
        "bit_depth": 5,
        "memory_size": 100,
        "action_repeats": 1,
    },
    "iris": {
        "env_backend": "atari",
    },
}

DEFAULT_TRAINING_CONFIGS: dict[str, dict[str, Any]] = {
    "dreamerv1": {
        "total_steps": 20000,
        "seed_steps": 1000,
        "update_steps": 50,
        "collect_steps": 500,
        "test_interval": 2000,
        "test_episodes": 1,
        "preview_eval_episodes": 1,
    },
    "dreamerv2": {
        "total_steps": 20000,
        "seed_steps": 1000,
        "update_steps": 50,
        "collect_steps": 500,
        "test_interval": 2000,
        "test_episodes": 1,
        "preview_eval_episodes": 1,
    },
    "planet": {
        "epochs": 25,
        "warmup_episodes": 2,
        "steps_per_epoch": 50,
        "batch_size": 32,
        "horizon": 50,
        "beta": 1.0,
        "save_every": 10,
        "record_grads": False,
    },
    "iris": {
        "total_epochs": 100,
        "seed_steps": 1000,
        "test_interval": 10,
    },
}


def _build_env_catalog() -> dict[str, list[str]]:
    atari_envs: list[str] = []
    try:
        atari_envs = list_available_atari_envs()
    except Exception:
        atari_envs = []
    return {
        "dreamerv1": DREAMER_ENVS + GYM_ENVS,
        "dreamerv2": DREAMER_ENVS + GYM_ENVS,
        "planet": PLANET_BASE_ENVS + atari_envs[:80],
        "iris": atari_envs[:80],
    }


ENVIRONMENTS_BY_MODEL = _build_env_catalog()


class LoadModelRequest(BaseModel):
    model: str
    config: dict[str, Any] = Field(default_factory=dict)


class LoadEnvironmentRequest(BaseModel):
    environment: str
    backend: str = "dm_control"
    config: dict[str, Any] = Field(default_factory=dict)


class StartTrainingRequest(BaseModel):
    config: dict[str, Any] = Field(default_factory=dict)


def _coerce_bool(value: Any) -> bool:
    if isinstance(value, bool):
        return value
    if isinstance(value, str):
        return value.strip().lower() in {"1", "true", "yes", "on"}
    return bool(value)


def _coerce_like_template(template: Any, value: Any) -> Any:
    if isinstance(template, bool):
        return _coerce_bool(value)
    if isinstance(template, int) and not isinstance(template, bool):
        return int(value)
    if isinstance(template, float):
        return float(value)
    return value


def _to_float(value: Any) -> float | None:
    try:
        return float(value)
    except (TypeError, ValueError):
        return None


def _encode_frame_as_data_url(frame: Any) -> str | None:
    if frame is None:
        return None

    arr = np.asarray(frame)
    if arr.ndim != 3:
        return None

    # CHW -> HWC
    if arr.shape[0] in (1, 3, 4) and arr.shape[-1] not in (1, 3, 4):
        arr = np.transpose(arr, (1, 2, 0))

    if arr.shape[-1] == 1:
        arr = np.repeat(arr, 3, axis=-1)
    elif arr.shape[-1] == 4:
        arr = arr[..., :3]
    elif arr.shape[-1] != 3:
        return None

    arr = arr.astype(np.float32)
    mn, mx = float(arr.min()), float(arr.max())
    if mn >= -0.6 and mx <= 0.6:
        arr = (arr + 0.5) * 255.0
    elif mn >= 0.0 and mx <= 1.0:
        arr = arr * 255.0
    arr = np.clip(arr, 0.0, 255.0).astype(np.uint8)

    encoded_ok, png = cv2.imencode(".png", cv2.cvtColor(arr, cv2.COLOR_RGB2BGR))
    if not encoded_ok:
        return None
    raw = base64.b64encode(png.tobytes()).decode("ascii")
    return f"data:image/png;base64,{raw}"


def _encode_frames_as_gif(frames: list[np.ndarray], duration: int = 100) -> str | None:
    if not frames or len(frames) == 0:
        return None

    try:
        pil_frames = []
        for arr in frames:
            arr = np.asarray(arr)
            if arr.ndim == 4:
                arr = arr[-1]
            if arr.shape[0] in (1, 3, 4) and arr.shape[-1] not in (1, 3, 4):
                arr = np.transpose(arr, (1, 2, 0))
            if arr.shape[-1] == 1:
                arr = np.repeat(arr, 3, axis=-1)
            elif arr.shape[-1] == 4:
                arr = arr[..., :3]
            elif arr.shape[-1] != 3:
                continue

            arr = arr.astype(np.float32)
            mn, mx = float(arr.min()), float(arr.max())
            if mn >= -0.6 and mx <= 0.6:
                arr = (arr + 0.5) * 255.0
            elif mn >= 0.0 and mx <= 1.0:
                arr = arr * 255.0
            arr = np.clip(arr, 0.0, 255.0).astype(np.uint8)

            pil_frames.append(Image.fromarray(arr))

        if not pil_frames:
            return None

        sample_interval = max(1, len(pil_frames) // 30)
        sampled_frames = pil_frames[::sample_interval][:30]

        buffer = BytesIO()
        sampled_frames[0].save(
            buffer,
            format="GIF",
            save_all=True,
            append_images=sampled_frames[1:],
            duration=duration,
            loop=0,
        )
        raw = base64.b64encode(buffer.getvalue()).decode("ascii")
        return f"data:image/gif;base64,{raw}"
    except Exception:
        return None


def _reward_stats(values: Any, prefix: str) -> dict[str, float]:
    arr = np.asarray(values, dtype=np.float32).reshape(-1)
    if arr.size == 0:
        return {}
    return {
        f"{prefix}/avg_reward": float(arr.mean()),
        f"{prefix}/max_reward": float(arr.max()),
        f"{prefix}/min_reward": float(arr.min()),
        f"{prefix}/std_reward": float(arr.std()),
    }


class TrainingController:
    def __init__(self) -> None:
        self._lock = threading.RLock()
        self._thread: threading.Thread | None = None
        self._stop_event = threading.Event()

        self._model_name: str | None = None
        self._model_config: dict[str, Any] = {}
        self._environment_name: str | None = None
        self._environment_config: dict[str, Any] = {}
        self._training_config: dict[str, Any] = {}

        self._status = "idle"
        self._message = "Ready."
        self._traceback: str | None = None
        self._results_dir: str | None = None

        self._progress_current = 0
        self._progress_total = 1
        self._progress_unit = "steps"

        self._metrics: dict[str, list[dict[str, float]]] = defaultdict(list)
        self._latest_frame_data_url: str | None = None
        self._latest_gif_data_url: str | None = None

        self._started_at: float | None = None
        self._finished_at: float | None = None
        self._last_update_at: float | None = None

        self._latest_metrics: dict[str, float] = {}

        self.last_latents_ref = [last_latents]

    def available_environments(self, model_name: str | None) -> list[str]:
        if model_name is None:
            merged = sorted(
                set(ENVIRONMENTS_BY_MODEL["dreamerv1"])
                | set(ENVIRONMENTS_BY_MODEL["dreamerv2"])
                | set(ENVIRONMENTS_BY_MODEL["planet"])
            )
            return merged
        key = model_name.strip().lower()
        if key not in ENVIRONMENTS_BY_MODEL:
            raise ValueError(f"Unknown model '{model_name}'.")
        return ENVIRONMENTS_BY_MODEL[key]

    def load_model(self, model_name: str, config: dict[str, Any] | None = None) -> None:
        key = model_name.strip().lower()
        if key not in SUPPORTED_MODELS:
            raise ValueError(f"Unsupported model '{model_name}'.")

        with self._lock:
            if self._is_running_locked():
                raise RuntimeError("Cannot change model while training is running.")
            self._model_name = key
            self._model_config = dict(config or {})
            valid_envs = set(ENVIRONMENTS_BY_MODEL[key])
            if self._environment_name not in valid_envs:
                self._environment_name = None
                self._environment_config = {}
            self._status = "idle"
            self._message = f"Model '{key}' loaded."
            self._traceback = None
            self._last_update_at = time.time()

    def load_environment(
        self,
        environment_name: str,
        backend: str = "dm_control",
        config: dict[str, Any] | None = None,
    ) -> None:
        env = environment_name.strip()
        backend_key = backend.strip().lower() if backend else "dm_control"

        backend_to_env_type = {
            "dm_control": "dmc",
            "mujoco": "gym",
            "gym": "gym",
            "unity": "unity_mlagents",
        }

        with self._lock:
            if self._is_running_locked():
                raise RuntimeError(
                    "Cannot change environment while training is running."
                )
            if self._model_name is not None:
                valid = set(ENVIRONMENTS_BY_MODEL[self._model_name])
                if env not in valid:
                    raise ValueError(
                        f"Environment '{env}' is not valid for model '{self._model_name}'."
                    )
            self._environment_name = env
            merged_config = dict(config or {})
            merged_config["env_backend"] = backend_to_env_type.get(backend_key, "gym")
            self._environment_config = merged_config
            self._status = "idle"
            self._message = f"Environment '{env}' loaded (backend: {backend_key})."
            self._traceback = None
            self._last_update_at = time.time()

    def start_training(self, config: dict[str, Any] | None = None) -> None:
        with self._lock:
            if self._is_running_locked():
                raise RuntimeError("Training is already running.")
            if self._model_name is None:
                raise ValueError("Load a model before starting training.")
            if self._environment_name is None:
                raise ValueError("Load an environment before starting training.")

            defaults = DEFAULT_TRAINING_CONFIGS.get(self._model_name, {})
            merged = dict(defaults)
            merged.update(config or {})
            self._training_config = merged

            self._metrics = defaultdict(list)
            self._latest_frame_data_url = None
            self._latest_gif_data_url = None
            self._progress_current = 0
            self._progress_total = 1
            self._progress_unit = "steps"
            self._results_dir = None
            self._traceback = None

            self._status = "running"
            self._message = "Training started."
            self._started_at = time.time()
            self._finished_at = None
            self._last_update_at = self._started_at
            self._stop_event = threading.Event()

            self._thread = threading.Thread(
                target=self._run_training, name="torchwm-ui-trainer", daemon=True
            )
            self._thread.start()
            logger.info(
                "Training started: model=%s environment=%s results_dir=%s",
                self._model_name,
                self._environment_name,
                self._results_dir,
            )

    def stop_training(self) -> bool:
        with self._lock:
            if not self._is_running_locked():
                return False
            self._stop_event.set()
            self._message = "Stop requested. Waiting for current step to finish."
            self._last_update_at = time.time()
            return True

    def snapshot_state(self) -> dict[str, Any]:
        with self._lock:
            ratio = (
                float(self._progress_current) / float(self._progress_total)
                if self._progress_total > 0
                else 0.0
            )
            return {
                "model": self._model_name,
                "environment": self._environment_name,
                "status": self._status,
                "message": self._message,
                "traceback": self._traceback,
                "started_at": self._started_at,
                "finished_at": self._finished_at,
                "last_update_at": self._last_update_at,
                "results_dir": self._results_dir,
                "progress": {
                    "current": self._progress_current,
                    "total": self._progress_total,
                    "unit": self._progress_unit,
                    "ratio": max(0.0, min(1.0, ratio)),
                },
            }

    def snapshot_metrics(self, limit: int) -> dict[str, Any]:
        limit = max(1, min(limit, 5000))
        with self._lock:
            payload: dict[str, list[dict[str, float]]] = {}
            for key, values in self._metrics.items():
                payload[key] = values[-limit:]
            return {"series": payload}

    def snapshot_frame(self) -> dict[str, str | None]:
        with self._lock:
            return {
                "image": self._latest_frame_data_url,
                "gif": self._latest_gif_data_url,
            }

    def _is_running_locked(self) -> bool:
        return (
            self._status == "running"
            and self._thread is not None
            and self._thread.is_alive()
        )

    def _set_progress(self, current: int, total: int, unit: str) -> None:
        with self._lock:
            self._progress_current = int(max(0, current))
            self._progress_total = int(max(1, total))
            self._progress_unit = unit
            self._last_update_at = time.time()
            try:
                # Log a concise progress message for observability
                pct = (
                    100.0 * float(self._progress_current) / float(self._progress_total)
                )
            except Exception:
                pct = 0.0
            logger.info(
                "Progress: model=%s env=%s results=%s %d/%d %s (%.1f%%)",
                self._model_name,
                self._environment_name,
                self._results_dir,
                self._progress_current,
                self._progress_total,
                self._progress_unit,
                pct,
            )

    def _set_message(self, message: str) -> None:
        with self._lock:
            self._message = message
            self._last_update_at = time.time()

    def _append_metrics(self, step: int, metrics: dict[str, Any]) -> None:
        now = time.time()
        with self._lock:
            for key, value in metrics.items():
                numeric = _to_float(value)
                if numeric is None:
                    continue
                self._metrics[key].append(
                    {"step": float(step), "value": numeric, "timestamp": now}
                )
                self._latest_metrics[key] = numeric
                if len(self._metrics[key]) > 3000:
                    self._metrics[key] = self._metrics[key][-3000:]
            self._last_update_at = now

    def _set_frame(self, frame: Any) -> None:
        encoded = _encode_frame_as_data_url(frame)
        if encoded is None:
            return
        with self._lock:
            self._latest_frame_data_url = encoded
            self._last_update_at = time.time()

    def _set_gif(self, frames: list[Any]) -> None:
        if not frames:
            return
        with self._lock:
            self._latest_gif_data_url = _encode_frames_as_gif(frames)
            self._last_update_at = time.time()

    def _run_training(self) -> None:
        try:
            model_name = self._model_name
            if model_name in ("dreamerv1", "dreamerv2"):
                self._run_dreamer()
            elif model_name == "planet":
                self._run_planet()
            elif model_name == "iris":
                self._run_iris()
            else:
                raise ValueError(
                    f"No training implementation for model '{model_name}'."
                )

            with self._lock:
                self._finished_at = time.time()
            if self._stop_event.is_set():
                self._status = "stopped"
                self._message = "Training stopped by user."
            else:
                self._status = "completed"
                self._message = "Training completed."
            self._last_update_at = self._finished_at
            logger.info(
                "Training finished: model=%s environment=%s results_dir=%s status=%s",
                model_name,
                getattr(self, "_environment_name", None),
                getattr(self, "_results_dir", None),
                self._status,
            )
        except Exception as exc:
            # Capture full traceback and persist it so UI clients and logs
            # have a clear explanation for failures instead of silent
            # swallow. Also write a probe file into the results directory
            # when available to aid debugging across processes.
            full_tb = traceback.format_exc()
            logger.error("Training thread crashed", exc_info=True)
            with self._lock:
                self._status = "failed"
                # Use repr(exc) to ensure useful string even when str(exc) is empty
                self._message = f"{type(exc).__name__}: {exc!r}"
                self._traceback = full_tb
                self._finished_at = time.time()
                self._last_update_at = self._finished_at
            try:
                # Best-effort: write the traceback into the controller-managed
                # results dir so external tools (TensorBoard, UI) can inspect it.
                results_dir = getattr(self, "_results_dir", None)
                if results_dir:
                    err_path = Path(results_dir) / "training_exception.txt"
                    err_path.write_text(
                        f"{time.strftime('%Y-%m-%d %H:%M:%S')}\n{full_tb}"
                    )
            except Exception:
                logger.exception("Failed to write training exception to file")

    def _run_dreamer(self) -> None:
        with self._lock:
            env_name = self._environment_name
            model_cfg = dict(DEFAULT_MODEL_CONFIGS[self._model_name])
            model_cfg.update(self._model_config)
            model_cfg.update(self._environment_config)
            train_cfg = dict(self._training_config)

        if env_name is None:
            raise ValueError("Dreamer requires an environment.")

        default_root = Path(__file__).resolve().parent.parent.parent / "results"
        results_dir = str(
            train_cfg.get(
                "results_dir", default_root / f"ui_dreamer_{int(time.time())}"
            )
        )
        # Ensure directory exists and resolve to an absolute path so downstream
        # loggers/TensorBoard see the exact same location.
        Path(results_dir).mkdir(parents=True, exist_ok=True)
        results_dir = str(Path(results_dir).resolve())
        with self._lock:
            self._results_dir = results_dir
        logger.info(f"Training results_dir resolved: {results_dir}")

        cfg = DreamerConfig()
        cfg.env = env_name
        for key, value in {**model_cfg, **train_cfg}.items():
            if hasattr(cfg, key):
                current = getattr(cfg, key)
                setattr(cfg, key, _coerce_like_template(current, value))

        cfg.log_video_freq = -1
        cfg.restore = False

        # Ensure DreamerAgent writes logs into the controller-managed results_dir
        # (by default DreamerAgent used a package-local data path, which meant
        # TensorBoard started against the controller results_dir saw no events).
        agent = DreamerAgent(
            cfg, last_latents_ref=self.last_latents_ref, logdir=results_dir
        )
        total_steps = int(cfg.total_steps)
        self._set_progress(0, total_steps, "steps")

        seed_steps = max(1, int(cfg.seed_steps // max(1, cfg.action_repeat)))
        seed_rewards = agent.dreamer.collect_random_episodes(
            agent.train_env, seed_steps
        )
        seed_logs = _reward_stats(seed_rewards, "seed")
        self._append_metrics(0, seed_logs)
        agent.logger.log_scalars(seed_logs, 0)
        agent.logger.flush()
        try:
            broadcast_update(["state", "metrics"])
        except Exception:
            pass

        global_step = int(agent.dreamer.data_buffer.steps * cfg.action_repeat)
        self._set_progress(global_step, total_steps, "steps")

        print(
            f"[Dreamer] Starting training loop at global_step={global_step}, total_steps={total_steps}"
        )
        while global_step <= total_steps:
            if self._stop_event.is_set():
                print("[Dreamer] Training stopped by user")
                return

            model_loss = 0.0
            actor_loss = 0.0
            value_loss = 0.0
            print(f"[Dreamer] Running {int(cfg.update_steps)} training steps...")
            update_steps = max(1, int(cfg.update_steps))
            heartbeat_every = max(1, min(5, update_steps))
            for update_idx in range(update_steps):
                if self._stop_event.is_set():
                    return
                if update_idx == 0 or (update_idx + 1) % heartbeat_every == 0:
                    self._set_message(
                        f"Optimizing model {update_idx + 1}/{update_steps} at env step {global_step}/{total_steps}"
                    )
                    try:
                        broadcast_update(["state"])
                    except Exception:
                        pass
                model_loss, actor_loss, value_loss = agent.dreamer.train_one_batch()
            print(f"[Dreamer] Training steps done, model_loss={model_loss:.4f}")

            collect_steps = max(1, int(cfg.collect_steps // max(1, cfg.action_repeat)))
            self._set_message(
                f"Collecting {collect_steps} environment steps from env step {global_step}/{total_steps}"
            )
            try:
                broadcast_update(["state"])
            except Exception:
                pass
            print(f"[Dreamer] Collecting {collect_steps} env steps...")
            train_rewards = agent.dreamer.act_and_collect_data(
                agent.train_env, collect_steps
            )
            print(
                f"[Dreamer] Data collected, buffer has {agent.dreamer.data_buffer.steps} steps"
            )

            if global_step % 100 == 0:
                try:
                    obs = agent.test_env.reset()
                    _, action = agent.dreamer.act_with_world_model(
                        obs,
                        agent.dreamer.rssm.init_state(1, agent.dreamer.device),
                        torch.zeros(1, agent.dreamer.action_size).to(
                            agent.dreamer.device
                        ),
                        explore=False,
                    )
                    action_np = action[0].cpu().numpy()
                    next_obs, _, _, _ = agent.test_env.step(action_np)
                    if isinstance(next_obs, dict) and "image" in next_obs:
                        frame_img = next_obs["image"]
                        if hasattr(frame_img, "shape"):
                            if frame_img.shape[-1] == 3:
                                self._set_frame(frame_img)
                            elif frame_img.shape[0] == 3:
                                self._set_frame(frame_img.transpose(1, 2, 0))
                    elif hasattr(next_obs, "shape") and next_obs.shape[-1] == 3:
                        self._set_frame(next_obs)
                except Exception:
                    pass

            logs: dict[str, Any] = {
                "train/model_loss": model_loss,
                "train/actor_loss": actor_loss,
                "train/value_loss": value_loss,
            }
            logs.update(_reward_stats(train_rewards, "train"))

            self._append_metrics(global_step, logs)
            agent.logger.log_scalars(logs, global_step)
            agent.logger.flush()

            print(f"[Dreamer] Metrics logged at step {global_step}")

            # Allow more frequent realtime evaluations for visualizations.
            # Use `realtime_eval_interval` from training config if provided; default to 50 steps for Dreamer.
            eval_interval = max(1, int(cfg.test_interval))
            realtime_interval = max(0, int(train_cfg.get("realtime_eval_interval", 50)))
            should_eval = (
                (realtime_interval > 0 and global_step % realtime_interval == 0)
                or (global_step % eval_interval == 0)
                or (global_step >= total_steps)
            )
            if should_eval:
                eval_episodes = max(1, int(train_cfg.get("preview_eval_episodes", 1)))
                self._set_message(
                    f"Running realtime evaluation at env step {global_step}/{total_steps}"
                )
                try:
                    broadcast_update(["state"])
                except Exception:
                    pass
                eval_rewards, eval_videos, eval_latents = agent.dreamer.evaluate(
                    agent.test_env, eval_episodes, render=True
                )
                try:
                    if getattr(self, "last_latents_ref", None) is not None:
                        self.last_latents_ref[0] = eval_latents
                except Exception:
                    pass

                eval_logs = _reward_stats(eval_rewards, "eval")
                logs.update(eval_logs)
                agent.logger.log_scalars(eval_logs, global_step)
                if (
                    eval_videos is not None
                    and hasattr(eval_videos, "__len__")
                    and len(eval_videos) > 0
                ):
                    agent.logger.log_videos(
                        eval_videos,
                        global_step,
                        max_videos_to_save=min(
                            int(getattr(agent.args, "max_videos_to_save", 1)),
                            len(eval_videos),
                        ),
                    )
                frame = self._extract_dreamer_preview_frame(eval_videos)
                if frame is not None:
                    self._set_frame(frame)
                    self._set_gif(self._extract_dreamer_preview_frames(eval_videos))
                    self._save_preview_image(frame, global_step, results_dir)

            global_step = int(agent.dreamer.data_buffer.steps * cfg.action_repeat)
            self._set_progress(min(global_step, total_steps), total_steps, "steps")
            # Notify clients that progress/frame/latents may have updated
            try:
                broadcast_update(["state", "metrics", "frame", "latents"])
            except Exception:
                pass

    def _run_planet(self) -> None:
        with self._lock:
            env_name = self._environment_name
            model_cfg = dict(DEFAULT_MODEL_CONFIGS["planet"])
            model_cfg.update(self._model_config)
            train_cfg = dict(self._training_config)

        if env_name is None:
            raise ValueError("Planet requires an environment.")

        default_root = Path(__file__).resolve().parent.parent.parent / "results"
        results_dir = str(
            train_cfg.get("results_dir", default_root / f"ui_planet_{int(time.time())}")
        )
        Path(results_dir).mkdir(parents=True, exist_ok=True)
        results_dir = str(Path(results_dir).resolve())
        with self._lock:
            self._results_dir = results_dir
        logger.info(f"Training results_dir resolved: {results_dir}")

        planet = Planet(
            env=env_name,
            bit_depth=int(model_cfg.get("bit_depth", 5)),
            memory_size=int(model_cfg.get("memory_size", 100)),
            action_repeats=int(model_cfg.get("action_repeats", 1)),
            headless=True,
            results_dir=results_dir,
        )

        warmup_episodes = int(train_cfg.get("warmup_episodes", 1))
        if warmup_episodes > 0:
            planet.warmup(n_episodes=warmup_episodes, random_policy=True)

        epochs = max(1, int(train_cfg.get("epochs", 20)))
        steps_per_epoch = max(1, int(train_cfg.get("steps_per_epoch", 50)))
        batch_size = max(1, int(train_cfg.get("batch_size", 32)))
        horizon = max(2, int(train_cfg.get("horizon", 50)))
        beta = float(train_cfg.get("beta", 1.0))
        save_every = max(1, int(train_cfg.get("save_every", 10)))
        record_grads = _coerce_bool(train_cfg.get("record_grads", False))

        self._set_progress(0, epochs, "epochs")

        for epoch in range(1, epochs + 1):
            if self._stop_event.is_set():
                return

            epoch_metrics: dict[str, list[float]] = defaultdict(list)
            for _ in range(steps_per_epoch):
                if self._stop_event.is_set():
                    return
                train_metrics = planet_train(
                    planet.memory,
                    planet.rssm.train(),
                    planet.optimizer,
                    planet.device,
                    N=batch_size,
                    H=horizon,
                    beta=beta,
                    grads=record_grads,
                )
                flat = flatten_dict(train_metrics)
                for key, value in flat.items():
                    numeric = _to_float(value)
                    if numeric is not None:
                        epoch_metrics[f"train/{key.replace('.', '/')}"].append(numeric)

            epoch_summary: dict[str, float] = {
                key: float(np.mean(values))
                for key, values in epoch_metrics.items()
                if values
            }

            planet.memory.append(planet.rollout_gen.rollout_once(explore=True))
            eval_episode, eval_frames, eval_metrics = planet.rollout_gen.rollout_eval()
            planet.memory.append(eval_episode)

            flat_eval = flatten_dict(eval_metrics)
            for key, value in flat_eval.items():
                metric_name = key.replace(".", "/")
                if isinstance(value, (list, tuple, np.ndarray)):
                    arr = np.asarray(value, dtype=np.float32).reshape(-1)
                    if arr.size > 0:
                        epoch_summary[f"{metric_name}_mean"] = float(arr.mean())
                else:
                    numeric = _to_float(value)
                    if numeric is not None:
                        epoch_summary[metric_name] = numeric

            if eval_frames is not None and len(eval_frames) > 0:
                self._set_frame(eval_frames[-1])

            self._append_metrics(epoch, epoch_summary)
            self._set_progress(epoch, epochs, "epochs")
            self._set_message(f"Planet epoch {epoch}/{epochs}")

            if epoch % save_every == 0:
                ckpt = Path(results_dir) / f"ckpt_{epoch}.pth"
                torch.save(planet.rssm.state_dict(), ckpt)

    def _run_iris(self) -> None:
        with self._lock:
            env_name = self._environment_name
            train_cfg = dict(self._training_config)

        if env_name is None:
            raise ValueError("IRIS requires an environment.")

        default_root = Path(__file__).resolve().parent.parent.parent / "results"
        results_dir = str(
            train_cfg.get("results_dir", default_root / f"ui_iris_{int(time.time())}")
        )
        Path(results_dir).mkdir(parents=True, exist_ok=True)
        results_dir = str(Path(results_dir).resolve())
        with self._lock:
            self._results_dir = results_dir
        logger.info(f"Training results_dir resolved: {results_dir}")

        config = IRISConfig()
        total_epochs = int(train_cfg.get("total_epochs", 100))
        self._set_progress(0, total_epochs, "epochs")

        trainer = IRISTrainer(
            game=env_name,
            device="cuda" if torch.cuda.is_available() else "cpu",
            config=config,
        )

        seed_steps = int(train_cfg.get("seed_steps", 1000))
        self._set_message(f"Collecting {seed_steps} seed steps...")
        try:
            broadcast_update(["state"])
        except Exception:
            pass
        seed_return = trainer.collect_experience(seed_steps, epsilon=1.0)
        self._append_metrics(0, {"seed_return": seed_return})

        test_interval = int(train_cfg.get("test_interval", 10))

        for epoch in range(total_epochs):
            if self._stop_event.is_set():
                print("[IRIS] Training stopped by user")
                return

            self._set_message(f"IRIS epoch {epoch}/{total_epochs}")
            try:
                broadcast_update(["state"])
            except Exception:
                pass

            train_metrics = trainer.train_epoch(epoch)
            self._append_metrics(epoch, train_metrics)

            self._set_progress(epoch, total_epochs, "epochs")

            if epoch % test_interval == 0 and epoch > 0:
                logger.info("IRIS evaluation start: epoch=%d", epoch)
                try:
                    # Request rendered evaluation so we get videos and latents
                    eval_rewards, eval_videos, eval_latents = trainer.evaluate(
                        num_episodes=1, render=True
                    )
                    # Record latents for frontend visualization
                    try:
                        if getattr(self, "last_latents_ref", None) is not None:
                            # eval_latents may be empty array
                            self.last_latents_ref[0] = eval_latents
                    except Exception:
                        logger.exception("Failed to set last_latents_ref for IRIS")

                    # Append scalar metrics
                    eval_logs = _reward_stats(eval_rewards, "eval")
                    self._append_metrics(epoch, eval_logs)

                    # If we have video frames, set preview frame/gif and save preview
                    if (
                        eval_videos is not None
                        and hasattr(eval_videos, "__len__")
                        and len(eval_videos) > 0
                    ):
                        frame = self._extract_dreamer_preview_frame(eval_videos)
                        if frame is not None:
                            self._set_frame(frame)
                            self._set_gif(
                                self._extract_dreamer_preview_frames(eval_videos)
                            )
                            self._save_preview_image(frame, epoch, results_dir)

                    # Notify SSE clients that state/metrics/frame/latents updated
                    try:
                        broadcast_update(["state", "metrics", "frame", "latents"])
                    except Exception:
                        pass
                except Exception:
                    logger.exception("IRIS evaluation failed at epoch=%d", epoch)

        ckpt = Path(results_dir) / "iris_final.pth"
        torch.save(trainer.agent.state_dict(), ckpt)

    @staticmethod
    def _extract_dreamer_preview_frame(videos: Any) -> np.ndarray | None:
        if videos is None:
            return None

        if isinstance(videos, np.ndarray) and videos.dtype != object:
            if videos.ndim >= 5:
                return np.asarray(videos[0, -1])
            if videos.ndim == 4:
                return np.asarray(videos[-1])
            return None

        try:
            first_video = videos[0]
            if len(first_video) == 0:
                return None
            return np.asarray(first_video[-1])
        except Exception:
            return None

    @staticmethod
    def _extract_dreamer_preview_frames(videos: Any) -> list[np.ndarray]:
        if videos is None:
            return []

        frames = []
        try:
            if isinstance(videos, np.ndarray) and videos.dtype != object:
                if videos.ndim >= 5:
                    for ep in videos[:1]:
                        frames.extend([np.asarray(f) for f in ep])
                elif videos.ndim == 4:
                    frames = [np.asarray(f) for f in videos]
            else:
                first_video = videos[0]
                if len(first_video) > 0:
                    frames = [np.asarray(f) for f in first_video]
        except Exception:
            pass
        return frames

    def _save_preview_image(
        self, frame: np.ndarray, step: int, results_dir: str | None
    ) -> None:
        if results_dir is None:
            return
        try:
            preview_dir = Path(results_dir) / "previews"
            preview_dir.mkdir(parents=True, exist_ok=True)
            frame_path = preview_dir / f"step_{step:08d}.png"
            if frame is not None:
                arr = np.asarray(frame)
                if arr.ndim == 3 and arr.shape[-1] == 3:
                    cv2.imwrite(str(frame_path), cv2.cvtColor(arr, cv2.COLOR_RGB2BGR))
        except Exception:
            pass


controller = TrainingController()
app = FastAPI(title="TorchWM UI Backend", version="0.1.0")

# Simple subscriber list for Server-Sent Events (SSE) notifications.
# Each entry is a tuple (asyncio.Queue, loop) so background threads can push
# updates into the event loop with loop.call_soon_threadsafe(q.put_nowait, msg).
SUBSCRIBERS: list[tuple[asyncio.Queue, asyncio.AbstractEventLoop]] = []


def broadcast_update(keys: list[str], extra: dict | None = None) -> None:
    """Notify all connected SSE subscribers that resources updated.

    keys: list of strings like 'latents', 'frame', 'gif', 'state'.
    extra: optional small metadata.
    """
    payload = {"type": "update", "keys": keys, "extra": extra or {}}
    text = json.dumps(payload)
    for q, loop in list(SUBSCRIBERS):
        try:
            loop.call_soon_threadsafe(q.put_nowait, text)
        except Exception:
            # ignore per-subscriber failures
            pass


@app.get("/api/events")
async def sse_events():
    """Server-Sent Events endpoint that streams simple JSON update notifications.

    Clients should connect and listen for `data: {...}` lines. Each message is a JSON
    object with keys: type, keys, extra.
    """
    q: asyncio.Queue = asyncio.Queue()
    loop = asyncio.get_event_loop()
    SUBSCRIBERS.append((q, loop))

    async def event_generator():
        try:
            while True:
                msg = await q.get()
                yield f"data: {msg}\n\n"
        except asyncio.CancelledError:
            return
        finally:
            # remove subscriber
            try:
                SUBSCRIBERS.remove((q, loop))
            except Exception:
                pass

    return StreamingResponse(event_generator(), media_type="text/event-stream")


cors_origins = [
    "http://127.0.0.1:5173",
    "http://localhost:5173",
    "http://127.0.0.1:*",
    "http://localhost:*",
    "https://torchwm.vercel.app",
    "https://torchwm.onrender.com",
    "file://",
]

if os.environ.get("ELECTRON_RUN") == "true":
    cors_origins.extend(
        [
            "http://0.0.0.0:*",
            "http://*",
        ]
    )

app.add_middleware(
    CORSMiddleware,
    allow_origins=cors_origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

BASE_DIR = Path(__file__).resolve().parent.parent.parent
DIST_DIR = BASE_DIR / "torchwm_ui" / "dist"
# Central results root for UI-launched trainings and TensorBoard default.
BASE_RESULTS_DIR = BASE_DIR / "results"
BASE_RESULTS_DIR.mkdir(parents=True, exist_ok=True)

if DIST_DIR.exists():
    app.mount("/static", StaticFiles(directory=str(DIST_DIR)), name="static")


@app.get("/")
def serve_index():
    if DIST_DIR.exists():
        index_path = DIST_DIR / "index.html"
        if index_path.exists():
            return FileResponse(str(index_path))
    return {"error": "Frontend not found"}


@app.head("/")
def serve_index_head():
    if DIST_DIR.exists():
        index_path = DIST_DIR / "index.html"
        if index_path.exists():
            return FileResponse(str(index_path))
    return {"error": "Frontend not found"}


@app.get("/api/health")
def healthcheck() -> dict[str, str]:
    return {"status": "ok"}


@app.get("/api/catalog")
def catalog() -> dict[str, Any]:
    return {
        "models": SUPPORTED_MODELS,
        "environments_by_model": ENVIRONMENTS_BY_MODEL,
        "env_backends": ENV_BACKENDS,
        "default_model_configs": DEFAULT_MODEL_CONFIGS,
        "default_training_configs": DEFAULT_TRAINING_CONFIGS,
    }


@app.get("/api/environments")
def environments(
    model: str | None = Query(default=None), backend: str | None = Query(default=None)
) -> dict[str, Any]:
    if backend and backend in ENV_BACKENDS:
        return {
            "backend": backend,
            "items": ENV_BACKENDS[backend].get("environments", []),
        }
    try:
        items = controller.available_environments(model)
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc
    return {"model": model, "items": items}


@app.post("/api/load-model")
def load_model(payload: LoadModelRequest) -> dict[str, Any]:
    logger.info(f"load_model called: model={payload.model}, config={payload.config}")
    try:
        controller.load_model(payload.model, payload.config)
    except Exception as exc:
        logger.error(f"load_model failed: {exc}", exc_info=True)
        raise HTTPException(status_code=400, detail=str(exc)) from exc
    return controller.snapshot_state()


@app.post("/api/load-environment")
def load_environment(payload: LoadEnvironmentRequest) -> dict[str, Any]:
    logger.info(
        f"load_environment called: env={payload.environment}, backend={payload.backend}"
    )
    try:
        controller.load_environment(
            payload.environment, payload.backend, payload.config
        )
    except Exception as exc:
        logger.error(f"load_environment failed: {exc}", exc_info=True)
        raise HTTPException(status_code=400, detail=str(exc)) from exc
    return controller.snapshot_state()


@app.post("/api/train/start")
def start_training(payload: StartTrainingRequest) -> dict[str, Any]:
    logger.info(f"start_training called: config={payload.config}")
    try:
        controller.start_training(payload.config)
    except Exception as exc:
        logger.error(f"start_training failed: {exc}", exc_info=True)
        raise HTTPException(status_code=400, detail=str(exc)) from exc
    return controller.snapshot_state()


@app.post("/api/train/stop")
def stop_training() -> dict[str, Any]:
    stopped = controller.stop_training()
    return {"stop_requested": stopped, **controller.snapshot_state()}


@app.get("/api/state")
def state() -> dict[str, Any]:
    return controller.snapshot_state()


@app.get("/api/metrics")
def metrics(limit: int = Query(default=400, ge=1, le=5000)) -> dict[str, Any]:
    return controller.snapshot_metrics(limit=limit)


@app.get("/api/frame")
def frame() -> dict[str, str | None]:
    return controller.snapshot_frame()


TRAINING_DEPENDENCIES = [
    {"name": "torch", "label": "PyTorch", "required": True},
    {"name": "numpy", "label": "NumPy", "required": True},
    {"name": "einops", "label": "EinOps", "required": True},
    {"name": "cv2", "label": "OpenCV", "required": True},
    {"name": "PIL", "label": "Pillow", "required": True},
    {"name": "gym", "label": "Gym", "required": True},
    {"name": "gymnasium", "label": "Gymnasium", "required": True},
    {"name": "dm_control", "label": "DM Control", "required": False},
    {"name": "mujoco", "label": "MuJoCo", "required": False},
    {"name": "ale_py", "label": "ALE Py", "required": False},
    {"name": "mlagents", "label": "ML Agents", "required": False},
]


def _check_dependency(name: str) -> bool:
    # Prefer a lightweight presence check before importing to avoid triggering
    # heavy runtime errors during import (e.g. mujoco expecting GL config).
    try:
        # Map some friendly names to actual module specs
        special_map = {"mlagents": "mlagents_envs", "PIL": "PIL"}
        real_name = special_map.get(name, name)
        spec = importlib.util.find_spec(real_name)
        return spec is not None
    except Exception:
        return False


@app.get("/api/dependencies")
def get_dependencies() -> dict[str, Any]:
    deps = []
    for dep in TRAINING_DEPENDENCIES:
        installed = _check_dependency(dep["name"])
        deps.append(
            {
                "name": dep["name"],
                "label": dep["label"],
                "required": dep["required"],
                "installed": installed,
            }
        )
    return {"dependencies": deps}


@app.post("/api/visualize")
async def visualize_latents(
    request: Request,
    latents: str | None = Query(
        None, description="Base64 encoded numpy array of latents"
    ),
    method: str = Query("tsne", description="Visualization method: 'tsne' or 'umap'"),
    labels: str | None = Query(
        None, description="Base64 encoded numpy array of labels"
    ),
    perplexity: int = Query(30, description="Perplexity for t-SNE"),
    n_neighbors: int = Query(15, description="n_neighbors for UMAP"),
    shape: str | None = Query(
        None,
        description="Optional latents shape as comma-separated ints, e.g. '100,64'",
    ),
):
    try:
        import base64 as _base64

        # Allow clients to POST a JSON body with the latents to avoid URL length
        # limits. If latents not provided as query param, try to read JSON body.
        if not latents:
            try:
                body = await request.json()
                # body may contain keys: latents, labels, method, shape
                latents = body.get("latents") or body.get("latents_b64")
                if labels is None:
                    labels = body.get("labels")
                if method is None:
                    method = body.get("method", method)
                if shape is None:
                    shape = body.get("shape")
                # allow numeric params in body
                try:
                    perplexity = int(body.get("perplexity", perplexity))
                except Exception:
                    pass
                try:
                    n_neighbors = int(body.get("n_neighbors", n_neighbors))
                except Exception:
                    pass
            except Exception:
                # ignore body parsing failures here and fall through
                pass

        if not latents:
            raise HTTPException(status_code=400, detail="Missing 'latents' parameter")

        latents_data = _base64.b64decode(latents)
        latents_array = np.frombuffer(latents_data, dtype=np.float32)

        # If client passed a shape param, use it to reshape the 1D buffer into (N, D)
        req_shape = None
        if shape:
            try:
                req_shape = [int(x) for x in str(shape).split(",") if x.strip()]
            except Exception:
                req_shape = None

        if req_shape is not None and len(req_shape) > 0:
            try:
                latents_array = latents_array.reshape(tuple(req_shape))
            except Exception:
                raise HTTPException(status_code=400, detail="Invalid shape parameter")
        elif latents_array.ndim == 1:
            # If shape is missing, treat as 'not ready' case; return 204 No Content so client keeps polling silently
            raise HTTPException(status_code=204, detail="Latents not ready")

        labels_array = None
        if labels:
            labels_data = _base64.b64decode(labels)
            labels_array = np.frombuffer(labels_data, dtype=np.int32)

        if method.lower() == "tsne":
            fig = visualize_latent_tsne(
                latents_array, labels_array, perplexity=perplexity
            )
        elif method.lower() == "umap":
            fig = visualize_latent_umap(
                latents_array, labels_array, n_neighbors=n_neighbors
            )
        else:
            raise HTTPException(status_code=400, detail="Invalid method")

        if fig is None:
            # Likely missing visualization dependencies (scikit-learn / umap)
            raise HTTPException(
                status_code=500,
                detail=(
                    "Visualization unavailable: missing python dependencies. "
                    "Install scikit-learn (for t-SNE) and umap-learn (for UMAP): "
                    "pip install scikit-learn umap-learn"
                ),
            )

        html = fig.to_html(full_html=False)
        return {"html": html}
    except HTTPException:
        # Propagate FastAPI HTTPExceptions (204/400/etc.) unchanged
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/api/latents")
def get_latents():
    """Get recent latents from the current model (if available)."""
    # Prefer controller-held latents reference (updated by training/eval routines).
    latents_obj = None
    try:
        if "controller" in globals() and getattr(controller, "last_latents_ref", None):
            latents_obj = controller.last_latents_ref[0]
    except Exception:
        latents_obj = None

    # Fall back to module-level last_latents if present
    if latents_obj is None:
        latents_obj = last_latents

    if latents_obj is None:
        # No latents yet: return 204 so client keeps polling silently
        raise HTTPException(status_code=204, detail="Latents not ready")

    # Ensure it's a numpy array-like with bytes
    try:
        import base64

        # Convert to numpy array if possible
        if hasattr(latents_obj, "tobytes") and hasattr(latents_obj, "shape"):
            latents_b64 = base64.b64encode(latents_obj.tobytes()).decode("utf-8")
            return {"latents": latents_b64, "shape": list(latents_obj.shape)}
        else:
            raise HTTPException(status_code=204, detail="Latents not ready")
    except Exception:
        raise HTTPException(status_code=204, detail="Latents not ready")
