"""
Figmin XR MCP Bridge
Connects an AI Agent to Figmin XR.

- AI Agent talks to this bridge over stdio (MCP protocol via FastMCP)
- Figmin XR's agent HTML page connects to this bridge over WebSocket
- This bridge translates between the two
"""

import asyncio
import base64
import json
import logging
import os
import re
import sys
import threading
import urllib.request
import urllib.error
from collections import defaultdict
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager, contextmanager
from dataclasses import dataclass, field
from pathlib import Path

import websockets
from websockets.asyncio.client import connect as ws_connect
from websockets.asyncio.server import serve, ServerConnection

from mcp.server.fastmcp import FastMCP, Image
import mcp.types as types

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
WS_HOST = os.environ.get("FIGMIN_WS_HOST", "0.0.0.0")
WS_PORT = int(os.environ.get("FIGMIN_WS_PORT", "39571"))
WSS_PORT = int(os.environ.get("FIGMIN_WSS_PORT", "39572"))
PROXY_PORT = int(os.environ.get("FIGMIN_PROXY_PORT", "39573"))

# How long to wait for Figmin to respond to a single command.
# Default 300s (5 minutes) — plenty of room for long-running scripts,
# chat polling with extended waits, captureView on Quest, etc.
# Override: set FIGMIN_TIMEOUT env var (in seconds).
COMMAND_TIMEOUT = float(os.environ.get("FIGMIN_TIMEOUT", "300"))

# SKILL.md — API reference fetched once on startup, served via help() tool
SKILL_URL = "https://www.figmin.com/api/SKILL.md"
_CACHE_DIR = Path(os.environ.get("FIGMIN_CACHE_DIR", Path.home() / "figmin-mcp"))
_CACHE_FILE = _CACHE_DIR / "SKILL.md"
_ETAG_FILE = _CACHE_DIR / "SKILL.md.etag"
_initial_file_handler: logging.FileHandler | None = None

# Memory — persistent notes stored as .md files, shared across agent sessions
MEMORY_DIR = _CACHE_DIR / "memory"

logger = logging.getLogger("figmin-mcp")

# In-memory cache for the skill file content
_skill_content: str | None = None

# Topic-based help — parsed from SKILL.md on startup
_topic_content: dict[str, str] = {}
_TOPIC_MARKER = re.compile(r'<!--\s*@topic:\s*(.+?)\s*-->')

_HELP_MENU = """Figmin XR API — Help Topics

Call help("topic") to load a specific topic:

  basics        — Start here. Coordinate system, object creation,
                  lifecycle, gotchas, transforms, core functions.
  sketches      — Procedural 3D geometry using brush strokes.
                  Shapes (boxes, spheres, cylinders), stroke
                  manipulation, and the full brush catalog.
  physics       — Rigid bodies, forces, collisions, raycasting.
  multiplayer   — Authority, figmin-sync, messaging, late joiners.
  lighting      — Light component, scene lighting, shadows.
  interaction   — Player body parts (head, hands), reading
                  controller/trigger input via enableInputConnect,
                  and grab detection.
  apps          — Spawning independent apps, cross-app messaging.
  keyframes     — The Motion component: animate objects along
                  position/rotation keyframes.

Call help("all") to load the complete API reference."""


def _parse_topics(content: str) -> dict[str, str]:
    """Parse SKILL.md into topic sections using <!-- @topic: name --> markers.

    Markers look like: <!-- @topic: basics -->  or  <!-- @topic: basics, sketches -->
    Everything between two markers belongs to the topic(s) of the first marker.
    Content before any marker (frontmatter, etc.) is excluded from all topics.
    """
    topics: dict[str, list[str]] = defaultdict(list)
    current_topics: list[str] = []
    current_lines: list[str] = []

    for line in content.splitlines():
        match = _TOPIC_MARKER.match(line.strip())
        if match:
            # Save accumulated lines to current topics
            if current_lines and current_topics:
                chunk = "\n".join(current_lines).strip()
                if chunk:
                    for topic in current_topics:
                        topics[topic].append(chunk)
            # Start new section
            current_topics = [t.strip() for t in match.group(1).split(",")]
            current_lines = []
        else:
            current_lines.append(line)

    # Don't forget the last section
    if current_lines and current_topics:
        chunk = "\n".join(current_lines).strip()
        if chunk:
            for topic in current_topics:
                topics[topic].append(chunk)

    return {k: "\n\n".join(v) for k, v in topics.items()}


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _parse_bridge_response(result: dict) -> dict:
    """Parse the bridge response, handling the nested text field.

    The bridge returns responses as {"text": '{"ok":true,"logs":[...],...}', ...}.
    The actual data (ok, logs, pendingChat, etc.) is inside the JSON-encoded
    text field, not at the top level.
    """
    # Try to parse the text field which contains the real response
    text_field = result.get("text", "")
    if text_field:
        try:
            return json.loads(text_field)
        except (json.JSONDecodeError, TypeError):
            pass
    # Fallback: maybe the data is at the top level
    return result


def _expand_bridge_file_tokens(code: str) -> str:
    """Replace __BRIDGE_FILE__("path") tokens with the contents of the file.

    This is a bridge-side preprocessing step — NOT part of the Figmin JS API.
    It allows the agent to reference large HTML files by path instead of
    inlining them as strings, avoiding context window bloat.

    Usage in execute_script code:
        child.app.loadHTML(__BRIDGE_FILE__("C:/myapp.html"))

    If the file is not found, replaces with a minimal error HTML so the user
    sees a visible error card in the scene rather than a broken blank app.
    """
    def replacer(match):
        path = match.group(1).strip().strip('"\'')
        try:
            content = Path(path).read_text(encoding='utf-8')
            b64 = base64.b64encode(content.encode('utf-8')).decode('ascii')
            logger.info(f"__BRIDGE_FILE__: loaded {path} ({len(content)} bytes)")
            return f'atob("{b64}")'
        except FileNotFoundError:
            logger.error(f"__BRIDGE_FILE__: file not found: {path}")
            error_html = (
                f'<html><head><meta name="figmin"></head>'
                f'<body style="background:#1a0000;color:#ff4444;'
                f'font-family:sans-serif;padding:20px;font-size:18px">'
                f'<h2>&#9888; File Not Found</h2><p>{path}</p></body></html>'
            )
            b64 = base64.b64encode(error_html.encode('utf-8')).decode('ascii')
            return f'atob("{b64}")'

    return re.sub(r'__BRIDGE_FILE__\("([^"]+)"\)', replacer, code)


def _fetch_skill_md() -> None:
    """Download SKILL.md on startup, using ETag for cache validation.

    Caches to disk so restarts only re-download when the file changes.
    On network failure, falls back to the disk cache if available.
    """
    global _skill_content, _topic_content

    _CACHE_DIR.mkdir(exist_ok=True)

    # Load existing cache and ETag from disk
    cached_etag = None
    if _ETAG_FILE.exists():
        cached_etag = _ETAG_FILE.read_text().strip()
    if _CACHE_FILE.exists():
        _skill_content = _CACHE_FILE.read_text(encoding="utf-8")

    # Try to fetch, using conditional request if we have an ETag
    try:
        req = urllib.request.Request(SKILL_URL, method="GET")
        req.add_header("User-Agent", "figmin-xr-mcp/1.0")
        if cached_etag:
            req.add_header("If-None-Match", cached_etag)

        with urllib.request.urlopen(req, timeout=15) as resp:
            # 200 — new or updated content
            new_content = resp.read().decode("utf-8")
            new_etag = resp.headers.get("ETag")

            _skill_content = new_content
            _CACHE_FILE.write_text(new_content, encoding="utf-8")
            if new_etag:
                _ETAG_FILE.write_text(new_etag)
            logger.info(f"SKILL.md downloaded ({len(new_content)} bytes)")

    except urllib.error.HTTPError as e:
        if e.code == 304:
            # Not modified — cached version is current
            logger.info("SKILL.md is up to date (304 Not Modified)")
        else:
            logger.warning(f"Failed to fetch SKILL.md: HTTP {e.code}")
    except Exception as e:
        logger.warning(f"Failed to fetch SKILL.md: {e}")

    if _skill_content:
        _topic_content = _parse_topics(_skill_content)
        logger.info(
            f"SKILL.md ready ({len(_skill_content)} chars, "
            f"{len(_topic_content)} topics: {', '.join(sorted(_topic_content))})"
        )
    else:
        logger.warning("SKILL.md not available — help() will return an error")


# ---------------------------------------------------------------------------
# SSL Certificate (for WSS support)
# ---------------------------------------------------------------------------
# Embedded self-signed certificate for the WSS server. This exists solely to
# satisfy browser mixed-content policy (HTTPS pages can only connect to wss://,
# not ws://). This cert just needs to be *present* for TLS
# to handshake. 
# ---------------------------------------------------------------------------
_CERT_FILE = _CACHE_DIR / "bridge-cert.pem"
_KEY_FILE = _CACHE_DIR / "bridge-key.pem"

_EMBEDDED_CERT = """\
-----BEGIN CERTIFICATE-----
MIIDGTCCAgGgAwIBAgIUV064THjFuDmTnMACDfchZd8+znQwDQYJKoZIhvcNAQEL
BQAwHDEaMBgGA1UEAwwRZmlnbWluLW1jcC1icmlkZ2UwHhcNMjYwMjI2MDAwMTM1
WhcNMzYwMjI0MDAwMTM1WjAcMRowGAYDVQQDDBFmaWdtaW4tbWNwLWJyaWRnZTCC
ASIwDQYJKoZIhvcNAQEBBQADggEPADCCAQoCggEBAKG30oYwPlxMS0e12/wsVQJg
G9ZBSorgY+nS/6yAi9PYeJ38VghMUbY54a65NCpoNTQ0HChIO4sqGDbp4fcu3hC5
Qo8Oga9UxkJfBwsXDU2hlN0ZjCm28kYZOky0a00Ck4SMoRBuXPyljJZCqgwTahFE
at/L/OZpxw1T+6rGOQSCqCMBv16x3sW/d9DWkNlnzVfw2nACRj9ogTPg8pZ65hud
9gErEJy/oCtng1bS1lAfwsxjyl9lD+YQP7ZtaK3mFZ8Iq8pkV2lelYYh3VV220oC
ZIAVpvX0/B6MKd/GF7C9OME1k+714Hqh3k5PIkVfIWPUuy+LW4vcJ2cR0feT3+MC
AwEAAaNTMFEwHQYDVR0OBBYEFAn93WYaIq76IFrr460Zum0omq3QMB8GA1UdIwQY
MBaAFAn93WYaIq76IFrr460Zum0omq3QMA8GA1UdEwEB/wQFMAMBAf8wDQYJKoZI
hvcNAQELBQADggEBAJJqZg9Qqlt9Xy1umaFnqGf4PSmVY27CQAzX3Qqdu0VPy3FC
GiJXeUxnhkAbbWFDGnPjeHBVdKW+jGCptUGQEopClD6S7h5l1y3wda8cPz3yindN
ZbJRUEcckv0vbjMARrqaOaRtmzG/V3i25y6Jm9QSKPSM9FwOFtkDnAQsbGzQlzf4
D/yB1XE+Fof3J0Z8YfYA28OVf+z7GLO+/GbXXGPNVBb6ZO2i0Bo0TcnTm4tbUAM+
tWsHKTxRA5Xk3wLOSSVdf0PEx8rDkYfwJKF6L7oyVJN4uMwdgt1iqM1AKawCkXZU
4uobYQQcu7jZDZ/RZl6KZcaNhFmusD7mpA0mB6M=
-----END CERTIFICATE-----
"""

_EMBEDDED_KEY = """\
-----BEGIN PRIVATE KEY-----
MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQCht9KGMD5cTEtH
tdv8LFUCYBvWQUqK4GPp0v+sgIvT2Hid/FYITFG2OeGuuTQqaDU0NBwoSDuLKhg2
6eH3Lt4QuUKPDoGvVMZCXwcLFw1NoZTdGYwptvJGGTpMtGtNApOEjKEQblz8pYyW
QqoME2oRRGrfy/zmaccNU/uqxjkEgqgjAb9esd7Fv3fQ1pDZZ81X8NpwAkY/aIEz
4PKWeuYbnfYBKxCcv6ArZ4NW0tZQH8LMY8pfZQ/mED+2bWit5hWfCKvKZFdpXpWG
Id1VdttKAmSAFab19PwejCnfxhewvTjBNZPu9eB6od5OTyJFXyFj1Lsvi1uL3Cdn
EdH3k9/jAgMBAAECggEAMV8VfrIr17HTKclzp8FBEsCUvwyf3VY0KqaoN9nm1n44
FMN84HusVp3FMLqKGohF9fISxpmG1C29xsYQno5IbYKht1sObKDNBmZMmC4peZHj
CL1L9VBNFumMyu4p9BDdSJJGeeW3rco86YLgt18V+r+QgVKdmxqgEomZQ8HQlndw
MvEwCazweDtMcgrexWn2XDGj8P3tZY5Vblkv1qBV3xclT5F4tN/mYiaQNuuv554n
u2PbrNr06Rrb01rRFCCHWy5/sp1jjRd26ij0m1vMfARuuVslhBZyOFJIo7ektY3t
pNIfq91YozFnYDIp5mbxFJpFOdhUTDoxwjYGvSBdAQKBgQDi8QX1aB5UISjZJ6ma
HAqtz6m0N8UcbvaP4aOa35NZGXxjL2s7cV1rF/B4SRNb2H8r+ym1fDflTnm6a6C0
7dg5ZHxKnse6NC/0JQ3qQOYD/TuRgakhmV+CPZ+y10sMMeuWfYh18NhR8nkR6wnv
n4Ctcq9M7EKnYCE7aYCvdJqGEQKBgQC2bNKABfgH3T8/ZglvUKYoxKt1nduIjd0t
PmZPOJdYm5sEdE+kxlj0DUKa5n6ZB8YPlMoCqIrNZLft0tooJV55gdFVLhsNDfCm
4sS9SWC7+O7PvyRfSzwJMnI8aC7ILdGx/hM2yg2H+VelXlZeCnoly/JuuNLPVwwe
S3cFZrwCswKBgATx44NZZW/H2TACITvuaH0pDTWUEYNxF4ZDEGGLhZZna8JtghSl
f7eZqe+1B+r0aLD4pAwETON+NkDNn47Rr+hwPBUUKJ3yDSlwtUDpcRmdJvMgtCK7
SM15skUfU89MNynsSlDnko8WjXTKfkjuXMtquE8gxsMG4TJ+NgOcstFBAoGANRq8
tHQIXhG0BbqoaHUryZZm61hGvhu+FRujINCyjiLOH5/UR31OC82IZBtRIy82IvcH
T3rM3TnCqULGKwWl2O1HiOphCY3TTmGZMBkaRd4TdvqHlg4Krgq3YopZhACCmLQD
1+E9yyV0tDkgH0Qhhrs2GbFH/P+0EWWagkID3S0CgYEAz7nBH4tn45lAPcPBj3Lk
roZ/pu1KgR934s5m/GjShE/fcHNBExi8lJfqxHnlf0GMGfj5zYepACbtJB6iUzwx
ZNBEoKb3dkHNsE3ylagx/fZLPSUSCK2xP0cOya/8QWttfRIMFPHn7/wQ7SwSAQeN
91oSQTg0gdZ5GEGii831Lws=
-----END PRIVATE KEY-----
"""


def _ensure_ssl_cert() -> "ssl.SSLContext | None":
    """Write the embedded certificate to disk (if needed) and return an SSLContext."""
    import ssl

    _CACHE_DIR.mkdir(exist_ok=True)

    # Write embedded cert/key to disk if not already present
    if not _CERT_FILE.exists():
        _CERT_FILE.write_text(_EMBEDDED_CERT)
        logger.info(f"Wrote TLS certificate: {_CERT_FILE}")
    if not _KEY_FILE.exists():
        _KEY_FILE.write_text(_EMBEDDED_KEY)
        logger.info(f"Wrote TLS key: {_KEY_FILE}")

    try:
        ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
        ctx.load_cert_chain(str(_CERT_FILE), str(_KEY_FILE))
        return ctx
    except Exception as e:
        logger.warning(f"Failed to load TLS certificate: {e}")
        return None


# ---------------------------------------------------------------------------
# Speech-to-Text (faster-whisper, lazy-loaded)
# ---------------------------------------------------------------------------
_whisper_model = None
_whisper_model_id: str | None = None
_whisper_lock = threading.RLock()  # guards model swap + transcription

# Store Whisper models alongside bridge logs in project .cache
_WHISPER_CACHE = _CACHE_DIR / "whisper"
os.environ["HF_HOME"] = str(_WHISPER_CACHE)

# Available models — exposed to agent.html for model picker
WHISPER_MODELS = [
    {"id": "base.en",  "name": "English – Fast",              "size": "~150 MB"},
    {"id": "small.en", "name": "English – Advanced",           "size": "~460 MB"},
    {"id": "base",     "name": "Multilingual – Fast",          "size": "~150 MB"},
    {"id": "small",    "name": "Multilingual – Advanced",      "size": "~460 MB"},
]
DEFAULT_WHISPER_MODEL = os.environ.get("FIGMIN_WHISPER_MODEL", "base.en")
DEFAULT_TTS_VOICE = "am_adam"
_CONFIG_FILE = _CACHE_DIR / "config.json"


def _get_config() -> dict:
    """Load bridge config (whisper model, TTS voice, etc.)."""
    try:
        return json.loads(_CONFIG_FILE.read_text())
    except (FileNotFoundError, json.JSONDecodeError):
        return {"whisper_model": DEFAULT_WHISPER_MODEL, "tts_voice": DEFAULT_TTS_VOICE}


def _save_config(config: dict) -> None:
    """Persist bridge config."""
    try:
        _CACHE_DIR.mkdir(parents=True, exist_ok=True)
        _CONFIG_FILE.write_text(json.dumps(config, indent=2))
    except Exception as e:
        logger.warning(f"Could not save config: {e}")


def _get_preferred_model() -> str | None:
    """Return the saved whisper model id, or None."""
    model = _get_config().get("whisper_model")
    if model and any(m["id"] == model for m in WHISPER_MODELS):
        return model
    return None


def _save_preferred_model(model_id: str) -> None:
    """Persist whisper model choice."""
    config = _get_config()
    config["whisper_model"] = model_id
    _save_config(config)


def _get_tts_voice() -> str:
    """Return the saved TTS voice id."""
    return _get_config().get("tts_voice", DEFAULT_TTS_VOICE)


def _save_tts_voice(voice_id: str) -> None:
    """Persist TTS voice choice."""
    config = _get_config()
    config["tts_voice"] = voice_id
    _save_config(config)


def _is_model_cached(model_id: str) -> bool:
    """Check if a Whisper model is already downloaded to disk."""
    model_dir = _WHISPER_CACHE / "hub" / f"models--Systran--faster-whisper-{model_id}"
    return model_dir.exists()


def _download_whisper_model(model_id: str) -> None:
    """Download and load a Whisper model. Runs in a thread."""
    logger.info(f"Downloading Whisper model: {model_id}")
    # _get_whisper_model handles download from HuggingFace + loading
    _get_whisper_model(model_id)
    logger.info(f"Model '{model_id}' downloaded and loaded")


_cuda_dll_dirs_registered = False
_nvidia_pip_dll_dirs: list[str] = []  # Cached nvidia pip package DLL directories


def _find_nvidia_pip_dll_dirs() -> list[str]:
    """Discover DLL directories from nvidia pip packages (nvidia-cuda-runtime-cu12, nvidia-cudnn-cu12, etc.).

    These packages install DLLs under site-packages/nvidia/*/bin/ on Windows
    and site-packages/nvidia/*/lib/ on Linux.  Returns a list of existing
    directories.  Safe to call when the packages are not installed.
    """
    global _nvidia_pip_dll_dirs
    if _nvidia_pip_dll_dirs:
        return _nvidia_pip_dll_dirs

    import importlib.util, os, platform
    dirs: list[str] = []

    # Try onnxruntime.preload_dlls() first (ORT 1.21+) — it handles everything internally
    try:
        import onnxruntime as _ort
        if hasattr(_ort, "preload_dlls"):
            _ort.preload_dlls(directory="")
            logger.info("Called onnxruntime.preload_dlls(directory='') — ORT 1.21+ nvidia site-packages search")
    except Exception as e:
        logger.debug(f"onnxruntime.preload_dlls() not available or failed: {e}")

    # Manually discover nvidia pip package DLL dirs as a fallback / supplement
    # Packages like nvidia-cuda-runtime-cu12 install to:
    #   site-packages/nvidia/cuda_runtime/bin/  (Windows)
    #   site-packages/nvidia/cuda_runtime/lib/  (Linux)
    # Similarly nvidia-cudnn-cu12 → site-packages/nvidia/cudnn/bin/ etc.
    try:
        import nvidia  # type: ignore
        nvidia_root = os.path.dirname(nvidia.__path__[0] if hasattr(nvidia, '__path__') else nvidia.__file__)
        nvidia_pkg_dir = os.path.join(nvidia_root, "nvidia")
        if os.path.isdir(nvidia_pkg_dir):
            is_win = platform.system() == "Windows"
            sub = "bin" if is_win else "lib"
            for entry in os.listdir(nvidia_pkg_dir):
                candidate = os.path.join(nvidia_pkg_dir, entry, sub)
                if os.path.isdir(candidate):
                    dirs.append(candidate)
    except ImportError:
        pass
    except Exception as e:
        logger.debug(f"nvidia pip package discovery error: {e}")

    # Also check common site-packages nvidia patterns directly
    try:
        import site
        for sp in site.getsitepackages() + [site.getusersitepackages()]:
            nv = os.path.join(sp, "nvidia")
            if os.path.isdir(nv):
                is_win = platform.system() == "Windows"
                sub = "bin" if is_win else "lib"
                for entry in os.listdir(nv):
                    candidate = os.path.join(nv, entry, sub)
                    if os.path.isdir(candidate) and candidate not in dirs:
                        dirs.append(candidate)
    except Exception:
        pass

    _nvidia_pip_dll_dirs = dirs
    if dirs:
        logger.info(f"Found nvidia pip package DLL dirs: {dirs}")
    return dirs


def _register_cuda_dll_directories():
    """Register CUDA/cuDNN directories for DLL loading on Windows.

    Python 3.8+ changed DLL loading to not search PATH by default.
    This must be called early, before any library tries to load CUDA DLLs.

    Searches both PATH (system installs) and nvidia pip packages.
    """
    global _cuda_dll_dirs_registered
    if _cuda_dll_dirs_registered:
        return
    _cuda_dll_dirs_registered = True

    import os, platform
    if platform.system() != "Windows":
        return

    registered = []

    # 1. Register PATH-based CUDA/cuDNN directories (system installs)
    for p in os.environ.get("PATH", "").split(os.pathsep):
        p_lower = p.lower()
        if os.path.isdir(p) and ("cuda" in p_lower or "cudnn" in p_lower or "nvidia" in p_lower):
            try:
                os.add_dll_directory(p)
                registered.append(p)
            except OSError:
                pass

    # 2. Register nvidia pip package directories (zero-config GPU)
    for p in _find_nvidia_pip_dll_dirs():
        if p not in registered:
            try:
                os.add_dll_directory(p)
                registered.append(p)
            except OSError:
                pass

    if registered:
        logger.info(f"Registered CUDA DLL directories: {registered}")
    else:
        logger.info("No CUDA/cuDNN directories found in PATH or pip packages")


def _get_whisper_model(model_id: str | None = None, save_pref: bool = True):
    """Lazy-load a faster-whisper model. Downloads from HuggingFace on first use."""
    global _whisper_model, _whisper_model_id

    model_id = model_id or _whisper_model_id or _get_preferred_model() or DEFAULT_WHISPER_MODEL

    if _whisper_model is not None and _whisper_model_id == model_id:
        return _whisper_model

    with _whisper_lock:
        # Re-check under lock (another thread may have loaded it)
        if _whisper_model is not None and _whisper_model_id == model_id:
            return _whisper_model

        # ── Detect CUDA BEFORE importing faster_whisper ──────────────
        # ctranslate2 (faster-whisper's backend) loads native DLLs at
        # import time.  On Windows, if CUDA DLLs are partially present
        # the import can hang or crash the thread.  Register DLL dirs
        # and probe CUDA first so we know what device to target.
        import ctypes, os, platform
        _register_cuda_dll_directories()

        device = "cpu"
        compute = "int8"
        try:
            if platform.system() == "Windows":
                try:
                    ctypes.cdll.LoadLibrary("cublas64_12.dll")
                except OSError:
                    search_dirs = os.environ.get("PATH", "").split(os.pathsep) + _find_nvidia_pip_dll_dirs()
                    found = False
                    for p in search_dirs:
                        candidate = os.path.join(p, "cublas64_12.dll")
                        if os.path.isfile(candidate):
                            logger.info(f"Loading cublas via full path: {candidate}")
                            ctypes.cdll.LoadLibrary(candidate)
                            found = True
                            break
                    if not found:
                        raise FileNotFoundError("cublas64_12.dll not found in any PATH directory or nvidia pip packages")
            else:
                ctypes.cdll.LoadLibrary("libcublas.so.12")
            device = "cuda"
            compute = "float16"
            logger.info("CUDA runtime available — Whisper will use GPU")
        except (OSError, Exception) as e:
            logger.info(f"CUDA runtime not available ({e}), Whisper will use CPU")

        # ── Import faster_whisper (triggers ctranslate2 native load) ─
        try:
            from faster_whisper import WhisperModel
        except Exception as e:
            logger.error(f"Failed to import faster_whisper: {e}")
            raise
        # Suppress faster-whisper's per-inference "Processing audio with duration" logs
        logging.getLogger("faster_whisper").setLevel(logging.WARNING)

        # Unload old model first to free GPU VRAM
        # NOTE: We stash the old model reference instead of destroying it
        # immediately. CTranslate2's CUDA destructor can segfault the process.
        # The stashed reference keeps the model alive; VRAM is freed when
        # the new model loads and Python eventually GCs the old one.
        if _whisper_model is not None:
            old_id = _whisper_model_id
            logger.info(f"Switching Whisper model from '{old_id}' to '{model_id}' ...")
            if not hasattr(_get_whisper_model, '_stale'):
                _get_whisper_model._stale = []
            _get_whisper_model._stale.append(_whisper_model)
            _whisper_model = None
            _whisper_model_id = None
            logger.info(f"Whisper model '{old_id}' replaced (old ref stashed)")

        logger.info(f"Loading Whisper model: {model_id} (device={device}, compute={compute}) ...")
        try:
            loaded = WhisperModel(model_id, device=device, compute_type=compute)
        except Exception as e:
            if device == "cuda":
                logger.warning(f"Failed to load Whisper on GPU ({e}), falling back to CPU")
                device = "cpu"
                compute = "int8"
                loaded = WhisperModel(model_id, device=device, compute_type=compute)
            else:
                raise
        logger.info(f"Whisper model '{model_id}' ready on {device}")
        _whisper_model = loaded
        _whisper_model_id = model_id
        if save_pref:
            _save_preferred_model(model_id)
        return _whisper_model


def _transcribe(audio_float32, model_id: str | None = None) -> str:
    """Run Whisper transcription on a float32 numpy array (16 kHz mono)."""
    with _whisper_lock:
        model = _get_whisper_model(model_id)

        # Use language hint for .en models, auto-detect for multilingual
        effective_id = model_id or _whisper_model_id or _get_preferred_model() or DEFAULT_WHISPER_MODEL
        lang = "en" if effective_id.endswith(".en") else None
        logger.debug(f"Transcribing with model={effective_id}, language={lang or 'auto-detect'}")
        segments, _ = model.transcribe(audio_float32, language=lang, beam_size=1)
        text = " ".join(seg.text.strip() for seg in segments)

    # Filter Whisper hallucinations on silence
    lower = text.lower().strip()
    if not lower or lower in (
        "[blank_audio]", "(blank audio)", "[silence]", "you",
        "thank you.", "thanks for watching!", "bye.",
    ):
        return ""

    return text.strip()


def _decode_and_transcribe(all_bytes: bytes, model_id: str | None = None) -> str:
    """Decode Int16 PCM bytes and run Whisper. Runs entirely in a thread pool.

    This does all blocking work (numpy import, audio conversion, model loading,
    transcription) in one function so the async event loop is never blocked.
    """
    import numpy as np

    audio_int16 = np.frombuffer(all_bytes, dtype=np.int16)
    audio_float = audio_int16.astype(np.float32) / 32768.0
    logger.debug(f"Audio decoded: {len(audio_int16)} samples, {len(audio_int16)/16000:.1f}s")

    return _transcribe(audio_float, model_id)


# ---------------------------------------------------------------------------
# Text-to-Speech (kokoro-onnx, lazy-loaded)
# ---------------------------------------------------------------------------
_kokoro_model = None
_KOKORO_CACHE = _CACHE_DIR / "kokoro"

# Suppress "Could not locate nvrtc64_120_0.dll" stderr spam from NVIDIA's CUDA
# runtime.  Without nvidia-cuda-nvrtc-cu12 (~200 MB), ORT can't JIT-compile
# optimised kernels and prints this warning on every TTS inference.  Performance
# is unaffected (~0.5-0.75s) so we skip the package and mute the noise.
# To install instead: add "--with nvidia-cuda-nvrtc-cu12" to the uv config.
@contextmanager
def _suppress_cuda_stderr():
    """Temporarily redirect stderr to devnull to silence nvrtc DLL warnings."""
    try:
        devnull = os.open(os.devnull, os.O_WRONLY)
        old_stderr = os.dup(2)
        os.dup2(devnull, 2)
        os.close(devnull)
        yield
    finally:
        os.dup2(old_stderr, 2)
        os.close(old_stderr)
_KOKORO_MODEL_FILE = "kokoro-v1.0.fp16.onnx"
_KOKORO_VOICES_FILE = "voices-v1.0.bin"
_KOKORO_LICENSE_FILE = "LICENSE"  # Apache 2.0 — must be present alongside model

_KOKORO_MODEL_URL = "https://github.com/thewh1teagle/kokoro-onnx/releases/download/model-files-v1.0/kokoro-v1.0.fp16.onnx"
_KOKORO_VOICES_URL = "https://github.com/thewh1teagle/kokoro-onnx/releases/download/model-files-v1.0/voices-v1.0.bin"
_KOKORO_LICENSE_URL = "https://huggingface.co/hexgrad/Kokoro-82M/raw/main/LICENSE"

# Voice catalog — all Kokoro v1.0 voices
# The agent picks from these based on conversation context (language, preference).
KOKORO_VOICES = [
    # American English
    {"id": "af_heart",   "name": "Heart",    "lang": "en-us", "gender": "female"},
    {"id": "af_alloy",   "name": "Alloy",    "lang": "en-us", "gender": "female"},
    {"id": "af_aoede",   "name": "Aoede",    "lang": "en-us", "gender": "female"},
    {"id": "af_bella",   "name": "Bella",    "lang": "en-us", "gender": "female"},
    {"id": "af_jessica", "name": "Jessica",  "lang": "en-us", "gender": "female"},
    {"id": "af_kore",    "name": "Kore",     "lang": "en-us", "gender": "female"},
    {"id": "af_nicole",  "name": "Nicole",   "lang": "en-us", "gender": "female"},
    {"id": "af_nova",    "name": "Nova",     "lang": "en-us", "gender": "female"},
    {"id": "af_river",   "name": "River",    "lang": "en-us", "gender": "female"},
    {"id": "af_sarah",   "name": "Sarah",    "lang": "en-us", "gender": "female"},
    {"id": "af_sky",     "name": "Sky",      "lang": "en-us", "gender": "female"},
    {"id": "am_adam",    "name": "Adam",     "lang": "en-us", "gender": "male"},
    {"id": "am_echo",    "name": "Echo",     "lang": "en-us", "gender": "male"},
    {"id": "am_eric",    "name": "Eric",     "lang": "en-us", "gender": "male"},
    {"id": "am_liam",    "name": "Liam",     "lang": "en-us", "gender": "male"},
    {"id": "am_michael", "name": "Michael",  "lang": "en-us", "gender": "male"},
    {"id": "am_onyx",    "name": "Onyx",     "lang": "en-us", "gender": "male"},
    # British English
    {"id": "bf_emma",    "name": "Emma",     "lang": "en-gb", "gender": "female"},
    {"id": "bf_isabella","name": "Isabella", "lang": "en-gb", "gender": "female"},
    {"id": "bm_george",  "name": "George",   "lang": "en-gb", "gender": "male"},
    {"id": "bm_lewis",   "name": "Lewis",    "lang": "en-gb", "gender": "male"},
    # Spanish
    {"id": "ef_dora",    "name": "Dora",     "lang": "es",    "gender": "female"},
    {"id": "em_alex",    "name": "Alex",     "lang": "es",    "gender": "male"},
    # French
    {"id": "ff_siwis",   "name": "Siwis",    "lang": "fr-fr", "gender": "female"},
    # Japanese
    {"id": "jf_alpha",   "name": "Alpha",    "lang": "ja",    "gender": "female"},
    {"id": "jf_gongitsune", "name": "Gongitsune", "lang": "ja", "gender": "female"},
    {"id": "jm_kumo",    "name": "Kumo",     "lang": "ja",    "gender": "male"},
    # Hindi
    {"id": "hf_alpha",   "name": "Alpha",    "lang": "hi",    "gender": "female"},
    {"id": "hf_beta",    "name": "Beta",     "lang": "hi",    "gender": "female"},
    {"id": "hm_omega",   "name": "Omega",    "lang": "hi",    "gender": "male"},
    {"id": "hm_psi",     "name": "Psi",      "lang": "hi",    "gender": "male"},
    # Italian
    {"id": "if_sara",    "name": "Sara",     "lang": "it",    "gender": "female"},
    {"id": "im_nicola",  "name": "Nicola",   "lang": "it",    "gender": "male"},
    # Brazilian Portuguese
    {"id": "pf_dora",    "name": "Dora",     "lang": "pt-br", "gender": "female"},
    {"id": "pm_alex",    "name": "Alex",     "lang": "pt-br", "gender": "male"},
    {"id": "pm_santa",   "name": "Santa",    "lang": "pt-br", "gender": "male"},
    # Mandarin Chinese
    {"id": "zf_xiaobei", "name": "Xiaobei",  "lang": "cmn",   "gender": "female"},
    {"id": "zf_xiaoni",  "name": "Xiaoni",   "lang": "cmn",   "gender": "female"},
    {"id": "zf_xiaoxiao","name": "Xiaoxiao", "lang": "cmn",   "gender": "female"},
    {"id": "zf_xiaoyi",  "name": "Xiaoyi",   "lang": "cmn",   "gender": "female"},
    {"id": "zm_yunjian", "name": "Yunjian",  "lang": "cmn",   "gender": "male"},
    {"id": "zm_yunxi",   "name": "Yunxi",    "lang": "cmn",   "gender": "male"},
    {"id": "zm_yunxia",  "name": "Yunxia",   "lang": "cmn",   "gender": "male"},
    {"id": "zm_yunyang", "name": "Yunyang",  "lang": "cmn",   "gender": "male"},
    # Korean
    {"id": "kf_sarah",   "name": "Sarah",    "lang": "ko",    "gender": "female"},
    {"id": "km_chul",    "name": "Chul",     "lang": "ko",    "gender": "male"},
]

# Language inference from voice prefix
_VOICE_LANG_MAP = {
    "a": "en-us", "b": "en-gb", "e": "es", "f": "fr-fr",
    "j": "ja", "h": "hi", "i": "it", "p": "pt-br",
    "z": "cmn", "k": "ko",
}


def _kokoro_is_cached() -> bool:
    """Check if Kokoro model files are already downloaded."""
    return (
        (_KOKORO_CACHE / _KOKORO_MODEL_FILE).exists()
        and (_KOKORO_CACHE / _KOKORO_VOICES_FILE).exists()
    )


def _download_file(url: str, dest: Path, label: str = "") -> None:
    """Download a file with urllib (no dependencies)."""
    logger.info(f"Downloading {label or url} → {dest}")
    req = urllib.request.Request(url, method="GET")
    req.add_header("User-Agent", "figmin-xr-mcp/1.0")
    with urllib.request.urlopen(req, timeout=120) as resp:
        dest.parent.mkdir(parents=True, exist_ok=True)
        with open(dest, "wb") as f:
            while True:
                chunk = resp.read(1024 * 64)
                if not chunk:
                    break
                f.write(chunk)
    logger.info(f"Downloaded {label or url} ({dest.stat().st_size / 1024 / 1024:.1f} MB)")


def _download_kokoro_model() -> None:
    """Download Kokoro model files + LICENSE. Called from thread pool."""
    _KOKORO_CACHE.mkdir(parents=True, exist_ok=True)

    model_path = _KOKORO_CACHE / _KOKORO_MODEL_FILE
    voices_path = _KOKORO_CACHE / _KOKORO_VOICES_FILE
    license_path = _KOKORO_CACHE / _KOKORO_LICENSE_FILE

    if not model_path.exists():
        _download_file(_KOKORO_MODEL_URL, model_path, "Kokoro model (fp16)")
    if not voices_path.exists():
        _download_file(_KOKORO_VOICES_URL, voices_path, "Kokoro voices")
    if not license_path.exists():
        try:
            _download_file(_KOKORO_LICENSE_URL, license_path, "Kokoro LICENSE (Apache 2.0)")
        except Exception as e:
            # LICENSE download is best-effort — don't block TTS on it
            logger.warning(f"Could not download LICENSE: {e}")


def _get_kokoro_model():
    """Lazy-load Kokoro with optimized ONNX session. Thread-safe for init."""
    global _kokoro_model
    if _kokoro_model is not None:
        return _kokoro_model

    import onnxruntime as rt
    # Suppress ORT warnings about Conv ops "running in Fallback mode" caused by
    # missing nvidia-cuda-nvrtc-cu12 (NVIDIA Runtime Compiler).  Without NVRTC,
    # ORT can't JIT-compile optimised CUDA kernels for Kokoro's conv layers and
    # falls back to generic implementations.  Performance is still fine (~0.5-0.75s)
    # so we skip the ~200 MB nvrtc package and just silence the warnings.
    rt.set_default_logger_severity(3)  # 3 = ERROR only (suppresses WARNING)
    _register_cuda_dll_directories()
    from kokoro_onnx import Kokoro

    model_path = _KOKORO_CACHE / _KOKORO_MODEL_FILE
    voices_path = _KOKORO_CACHE / _KOKORO_VOICES_FILE

    if not model_path.exists() or not voices_path.exists():
        _download_kokoro_model()

    # Build optimized session options
    sess_opts = rt.SessionOptions()
    sess_opts.graph_optimization_level = rt.GraphOptimizationLevel.ORT_ENABLE_ALL
    sess_opts.execution_mode = rt.ExecutionMode.ORT_SEQUENTIAL
    # Use all available CPU cores for intra-op parallelism (0 = auto-detect)
    sess_opts.intra_op_num_threads = 0
    # Enable memory optimizations
    sess_opts.enable_mem_pattern = True
    sess_opts.enable_cpu_mem_arena = True
    # Allow threads to spin-wait for lower latency (trades CPU for speed)
    sess_opts.add_session_config_entry("session.intra_op.allow_spinning", "1")
    # Dynamic block sizing for better load balancing
    sess_opts.add_session_config_entry("session.dynamic_block_base", "4")
    # Save optimized model so subsequent loads are faster
    optimized_path = _KOKORO_CACHE / _KOKORO_MODEL_FILE.replace(".onnx", ".optimized.onnx")
    if optimized_path.exists():
        logger.info(f"Loading pre-optimized Kokoro model...")
        model_to_load = str(optimized_path)
    else:
        logger.info(f"Loading Kokoro TTS model (will optimize on first load)...")
        sess_opts.optimized_model_filepath = str(optimized_path)
        model_to_load = str(model_path)

    # Detect providers — prefer CUDA with optimizations, fall back to CPU
    available = rt.get_available_providers()
    logger.info(f"ONNX Runtime available providers: {available}")

    use_cuda = False
    if "CUDAExecutionProvider" in available:
        # Pre-check: verify cuDNN DLLs exist before attempting CUDA session.
        # Without this, ONNX Runtime hard-crashes (not a Python exception) if cuDNN is missing.
        import platform, os
        if platform.system() == "Windows":
            # Gather all directories to search: PATH + nvidia pip package dirs
            path_dirs = os.environ.get("PATH", "").split(os.pathsep)
            pip_dirs = _find_nvidia_pip_dll_dirs()
            all_search_dirs = path_dirs + pip_dirs

            cudnn_in_path = [p for p in path_dirs if "cudnn" in p.lower()]
            logger.info(f"cuDNN-related PATH entries: {cudnn_in_path or 'NONE'}")
            if pip_dirs:
                logger.info(f"nvidia pip DLL dirs: {pip_dirs}")

            cudnn_dlls = ["cudnn64_9.dll", "cudnn_graph64_9.dll"]
            all_found = True
            for dll_name in cudnn_dlls:
                found_in_path = False
                for p in all_search_dirs:
                    candidate = os.path.join(p, dll_name)
                    if os.path.isfile(candidate):
                        logger.info(f"cuDNN DLL found: {candidate}")
                        found_in_path = True
                        break
                if not found_in_path:
                    logger.warning(f"cuDNN DLL not found in PATH or pip packages: {dll_name}")
                    all_found = False
            if all_found:
                use_cuda = True
            else:
                logger.warning(
                    "CUDA provider available but cuDNN 9 DLLs missing from PATH. "
                    "Install cuDNN 9.x for CUDA 12 and ensure its bin directory is in PATH. "
                    "Falling back to CPU."
                )
        else:
            # On Linux, trust the provider list — cuDNN is usually in LD_LIBRARY_PATH
            use_cuda = True

    if use_cuda:
        # On Windows, register CUDA/cuDNN directories for DLL loading
        if platform.system() == "Windows":
            for p in all_search_dirs:
                if os.path.isdir(p) and ("cuda" in p.lower() or "cudnn" in p.lower() or "nvidia" in p.lower()):
                    try:
                        os.add_dll_directory(p)
                        logger.info(f"Added DLL directory for ONNX: {p}")
                    except OSError:
                        pass
        providers = [
            ("CUDAExecutionProvider", {"cudnn_conv_algo_search": "DEFAULT"}),
            "CPUExecutionProvider",
        ]
        logger.info("CUDA + cuDNN verified — using GPU acceleration for TTS")
    else:
        providers = ["CPUExecutionProvider"]
        if "CUDAExecutionProvider" not in available:
            logger.info("No GPU provider found — using CPU for TTS")

    with _suppress_cuda_stderr():
        try:
            session = rt.InferenceSession(model_to_load, sess_opts, providers=providers)
        except Exception as e:
            # Safety net — if CUDA session still fails, fall back to CPU
            if use_cuda:
                logger.warning(f"CUDA session failed ({e}), falling back to CPU")
                providers = ["CPUExecutionProvider"]
                session = rt.InferenceSession(model_to_load, sess_opts, providers=providers)
            else:
                raise

    active = session.get_providers()
    logger.info(f"Kokoro TTS session active providers: {active}")
    _kokoro_model = Kokoro.from_session(session, str(voices_path))
    logger.info("Kokoro TTS model ready")
    return _kokoro_model


def _infer_lang_from_voice(voice_id: str) -> str:
    """Infer the Kokoro language code from a voice ID prefix."""
    if voice_id and len(voice_id) >= 1:
        return _VOICE_LANG_MAP.get(voice_id[0], "en-us")
    return "en-us"


def _generate_tts_audio(text: str, voice: str = "", speed: float = 1.0) -> str:
    """Generate TTS audio and return base64-encoded WAV. Runs in thread pool."""
    import io
    import soundfile as sf

    voice = voice or _get_tts_voice()

    kokoro = _get_kokoro_model()
    lang = _infer_lang_from_voice(voice)
    logger.info(f"TTS generating: voice={voice}, lang={lang}, text={text[:80]}...")

    with _suppress_cuda_stderr():
        samples, sample_rate = kokoro.create(text, voice=voice, speed=speed, lang=lang)

    buf = io.BytesIO()
    sf.write(buf, samples, sample_rate, format="WAV", subtype="PCM_16")
    audio_b64 = base64.b64encode(buf.getvalue()).decode()

    duration = len(samples) / sample_rate
    logger.info(f"TTS generated: {duration:.1f}s audio, {len(audio_b64)} base64 chars")
    return audio_b64


async def _generate_tts_stream(text: str, voice: str = "", speed: float = 1.0):
    """Async generator: yields (base64_wav, chunk_index) using kokoro.create_stream().

    Kokoro handles text splitting internally, preserving prosody across sentences.
    Each chunk is sent as soon as it's generated.
    """
    import io
    import time
    import soundfile as sf

    voice = voice or _get_tts_voice()

    kokoro = _get_kokoro_model()
    lang = _infer_lang_from_voice(voice)
    logger.info(f"TTS streaming: voice={voice}, lang={lang}, text={text[:80]}...")

    t_start = time.monotonic()

    chunk_idx = 0
    with _suppress_cuda_stderr():
        stream = kokoro.create_stream(text, voice=voice, speed=speed, lang=lang)
        async for samples, sample_rate in stream:
            t_chunk = time.monotonic()
            buf = io.BytesIO()
            sf.write(buf, samples, sample_rate, format="WAV", subtype="PCM_16")
            audio_b64 = base64.b64encode(buf.getvalue()).decode()

            duration = len(samples) / sample_rate
            elapsed = t_chunk - t_start
            logger.info(
                f"TTS chunk {chunk_idx}: {duration:.1f}s audio, "
                f"{len(audio_b64)} base64 chars, elapsed={elapsed:.2f}s"
            )
            yield audio_b64, chunk_idx
            chunk_idx += 1

    total = time.monotonic() - t_start
    logger.info(f"TTS streaming done: {chunk_idx} chunks in {total:.2f}s")


def _tts_available() -> bool:
    """Check if kokoro-onnx is importable (installed via uv --with)."""
    try:
        import kokoro_onnx  # noqa: F401
        return True
    except ImportError:
        return False


# ---------------------------------------------------------------------------
# WebSocket Server -- Figmin connects here as a client
# ---------------------------------------------------------------------------
@dataclass
class FigminBridge:
    """Manages the WebSocket connection to the Figmin XR agent page."""

    _client: ServerConnection | None = field(default=None, init=False)
    _pending: dict = field(default_factory=dict, init=False)
    _request_id: int = field(default=0, init=False)
    _lock: asyncio.Lock = field(default_factory=asyncio.Lock, init=False)
    _audio_buffer: list = field(default_factory=list, init=False)
    _partial_running: bool = field(default=False, init=False)
    _proxies: dict = field(default_factory=dict, init=False)  # proxy_id → ws

    async def handle_connection(self, websocket: ServerConnection):
        """Handle an incoming connection from the Figmin XR agent page."""
        logger.info("Figmin XR agent connected!")
        self._client = websocket
        self._audio_buffer = []
        self._partial_running = False

        # Notify proxies that Figmin connected
        await self._notify_proxies_figmin_status(True)

        # Report TTS availability on connect
        asyncio.create_task(self._send_tts_status(websocket))
        try:
            async for raw in websocket:
                try:
                    msg = json.loads(raw)
                    msg_type = msg.get("type")

                    # ── Audio intercept (never reaches MCP) ──
                    if msg_type == "audio_chunk":
                        audio_bytes = base64.b64decode(msg.get("data", ""))
                        self._audio_buffer.append(audio_bytes)
                        if msg.get("final", False):
                            buf = self._audio_buffer
                            self._audio_buffer = []
                            self._partial_running = False
                            total = sum(len(b) for b in buf)
                            logger.info(f"Final audio chunk received: {len(buf)} chunks, {total} bytes total")
                            asyncio.create_task(
                                self._transcribe_and_respond(buf, websocket, partial=False)
                            )
                        elif not self._partial_running:
                            # Run partial transcription if we have 2+ seconds
                            total = sum(len(b) for b in self._audio_buffer)
                            if total >= 16000 * 1 * 2:  # 1s of Int16 = 32000 bytes
                                self._partial_running = True
                                buf_snapshot = list(self._audio_buffer)
                                asyncio.create_task(
                                    self._transcribe_and_respond(buf_snapshot, websocket, partial=True)
                                )
                        continue
                    if msg_type == "audio_cancel":
                        self._audio_buffer = []
                        self._partial_running = False
                        continue
                    if msg_type == "voice_models_request":
                        await self._send_voice_models(websocket)
                        continue
                    if msg_type == "voice_model_download":
                        asyncio.create_task(
                            self._handle_voice_model_download(msg, websocket)
                        )
                        continue
                    if msg_type == "voice_model_select":
                        asyncio.create_task(
                            self._handle_voice_model_select(msg, websocket)
                        )
                        continue

                    # ── TTS intercept (never reaches MCP) ──
                    if msg_type == "tts_request":
                        asyncio.create_task(
                            self._handle_tts_request(msg, websocket)
                        )
                        continue
                    if msg_type == "tts_download":
                        asyncio.create_task(
                            self._handle_tts_download(websocket)
                        )
                        continue

                    # ── Route command responses by id ──
                    msg_id = msg.get("id")
                    if msg_id is not None:
                        # Check if this response belongs to a proxy
                        id_str = str(msg_id)
                        routed = False
                        for pid, proxy_ws in list(self._proxies.items()):
                            if id_str.startswith(f"{pid}-"):
                                try:
                                    await proxy_ws.send(json.dumps(msg))
                                except Exception:
                                    logger.warning(
                                        f"Failed to route response to "
                                        f"proxy '{pid}'"
                                    )
                                routed = True
                                break
                        if not routed and msg_id in self._pending:
                            self._pending[msg_id].set_result(msg)
                except json.JSONDecodeError:
                    logger.warning(f"Invalid JSON from Figmin: {str(raw)[:200]}")
        except websockets.ConnectionClosed:
            logger.info("Figmin XR agent disconnected.")
        finally:
            if self._client is websocket:
                self._client = None
                self._audio_buffer = []
                # Notify proxies that Figmin disconnected
                asyncio.ensure_future(
                    self._notify_proxies_figmin_status(False)
                )
                # Cancel all pending command futures so send_command
                # returns immediately instead of waiting 5 minutes.
                for rid, future in list(self._pending.items()):
                    if not future.done():
                        future.set_result({"error": "Figmin XR disconnected"})
                self._pending.clear()

    async def _transcribe_and_respond(
        self, buffer: list[bytes], websocket: ServerConnection, partial: bool = False
    ):
        """Concatenate audio chunks, run Whisper, send result back."""
        tag = "partial" if partial else "final"
        logger.debug(f"_transcribe_and_respond ({tag}): starting")
        all_bytes = b"".join(buffer)
        if not all_bytes:
            if not partial:
                await websocket.send(json.dumps({
                    "type": "transcription",
                    "text": "",
                    "partial": False,
                    "error": "No audio received",
                }))
            return

        try:
            loop = asyncio.get_event_loop()
            text = await loop.run_in_executor(None, _decode_and_transcribe, all_bytes)
            logger.debug(f"Transcription ({tag}): {len(text)} chars, empty={not text}")
            await websocket.send(json.dumps({
                "type": "transcription",
                "text": text,
                "partial": partial,
            }))

        except ImportError as e:
            logger.error(f"Import error: {e}")
            if not partial:
                await websocket.send(json.dumps({
                    "type": "transcription",
                    "text": "",
                    "partial": False,
                    "error": (
                        "Voice requires faster-whisper. "
                        "Install figmin-mcp with voice dependencies (faster-whisper, numpy)."
                    ),
                }))
        except Exception as e:
            logger.error(f"Transcription error ({tag}): {e}", exc_info=True)
            if not partial:
                await websocket.send(json.dumps({
                    "type": "transcription",
                    "text": "",
                    "partial": False,
                    "error": f"Transcription failed: {e}",
                }))
        finally:
            if partial:
                self._partial_running = False

    async def _send_voice_models(self, websocket: ServerConnection):
        """Send available Whisper models to the agent page."""
        available = False
        try:
            import faster_whisper  # noqa: F401
            available = True
        except ImportError:
            pass

        # Check which models are cached on disk
        models = []
        for m in WHISPER_MODELS:
            models.append({
                **m,
                "downloaded": _is_model_cached(m["id"]) if available else False,
            })

        await websocket.send(json.dumps({
            "type": "voice_models",
            "available": available,
            "models": models,
            "default": _get_preferred_model(),
            "current": _whisper_model_id,
        }))

    async def _handle_voice_model_download(self, msg, websocket: ServerConnection):
        """Download a Whisper model and stream progress back."""
        model_id = msg.get("model_id")
        if not model_id:
            return

        loop = asyncio.get_event_loop()

        try:
            await websocket.send(json.dumps({
                "type": "voice_model_progress",
                "model_id": model_id,
                "status": "downloading",
                "progress": 0.1,
            }))

            # Download + load in thread pool (_get_whisper_model handles both)
            await loop.run_in_executor(None, _download_whisper_model, model_id)

            await websocket.send(json.dumps({
                "type": "voice_model_progress",
                "model_id": model_id,
                "status": "ready",
                "progress": 1.0,
            }))

            # Send updated model list
            await self._send_voice_models(websocket)

        except Exception as e:
            logger.error(f"Model download error: {e}")
            await websocket.send(json.dumps({
                "type": "voice_model_progress",
                "model_id": model_id,
                "status": "error",
                "error": str(e),
                "progress": 0,
            }))

    async def _handle_voice_model_select(self, msg, websocket: ServerConnection):
        """Switch the active Whisper model (must already be downloaded)."""
        model_id = msg.get("model_id")
        if not model_id:
            return

        loop = asyncio.get_event_loop()

        try:
            await websocket.send(json.dumps({
                "type": "voice_model_progress",
                "model_id": model_id,
                "status": "loading",
                "progress": 0.5,
            }))

            await loop.run_in_executor(None, _get_whisper_model, model_id)

            await websocket.send(json.dumps({
                "type": "voice_model_progress",
                "model_id": model_id,
                "status": "ready",
                "progress": 1.0,
            }))

            await self._send_voice_models(websocket)

        except Exception as e:
            logger.error(f"Model select error: {e}")
            await websocket.send(json.dumps({
                "type": "voice_model_progress",
                "model_id": model_id,
                "status": "error",
                "error": str(e),
                "progress": 0,
            }))

    # ── TTS (Text-to-Speech) ──

    async def _send_tts_status(self, websocket: ServerConnection):
        """Report TTS availability to the client."""
        available = _tts_available()
        cached = _kokoro_is_cached() if available else False

        # Report which provider is active (if model is loaded)
        provider = "unknown"
        if _kokoro_model is not None:
            try:
                active = _kokoro_model.sess.get_providers()
                provider = active[0] if active else "none"
            except Exception:
                pass

        await websocket.send(json.dumps({
            "type": "tts_status",
            "available": available,
            "ready": available and cached,
            "provider": provider if (available and cached) else None,
        }))
        logger.info(f"TTS status sent: available={available}, ready={available and cached}, provider={provider}")

    async def _handle_tts_request(self, msg, websocket: ServerConnection):
        """Generate TTS audio via kokoro streaming.

        Uses create_stream() which splits and generates audio progressively,
        sending each chunk as soon as it's ready. Prosody is preserved since
        Kokoro handles the splitting internally.
        """
        text = msg.get("text", "").strip()
        voice = _get_tts_voice()
        msg_id = msg.get("id", "")
        speed = msg.get("speed", 1.0)

        if not text:
            return

        if not _tts_available():
            await websocket.send(json.dumps({
                "type": "tts_audio",
                "id": msg_id,
                "error": "kokoro-onnx not installed. Add '--with kokoro-onnx --with soundfile' to the uv run command. For GPU acceleration: '--with kokoro-onnx[gpu] --with soundfile'",
            }))
            return

        try:
            chunk_count = 0
            async for audio_b64, chunk_idx in _generate_tts_stream(
                text, voice, speed
            ):
                await websocket.send(json.dumps({
                    "type": "tts_audio_chunk",
                    "audio": audio_b64,
                    "id": msg_id,
                    "chunk": chunk_idx,
                    "final": False,
                }))
                chunk_count += 1

            # Send final marker so client knows streaming is done
            await websocket.send(json.dumps({
                "type": "tts_audio_chunk",
                "id": msg_id,
                "chunk": chunk_count,
                "final": True,
            }))

        except Exception as e:
            logger.error(f"TTS generation error: {e}", exc_info=True)
            await websocket.send(json.dumps({
                "type": "tts_audio",
                "id": msg_id,
                "error": f"TTS failed: {e}",
            }))

    async def _handle_tts_download(self, websocket: ServerConnection):
        """Download the Kokoro model files, then report updated status."""
        loop = asyncio.get_event_loop()

        try:
            await websocket.send(json.dumps({
                "type": "tts_download_progress",
                "status": "downloading",
                "progress": 0.1,
            }))

            await loop.run_in_executor(None, _download_kokoro_model)

            # Also pre-load the model so first TTS request is fast
            await loop.run_in_executor(None, _get_kokoro_model)

            await websocket.send(json.dumps({
                "type": "tts_download_progress",
                "status": "ready",
                "progress": 1.0,
            }))

            # Send updated TTS status
            await self._send_tts_status(websocket)

        except Exception as e:
            logger.error(f"TTS download error: {e}", exc_info=True)
            await websocket.send(json.dumps({
                "type": "tts_download_progress",
                "status": "error",
                "error": str(e),
                "progress": 0,
            }))

    # ── Proxy support ──────────────────────────────────────────────────

    async def _notify_proxies_figmin_status(self, connected: bool) -> None:
        """Push Figmin connection status to all connected proxies."""
        msg = json.dumps({"type": "figmin_status", "connected": connected})
        for pid, proxy_ws in list(self._proxies.items()):
            try:
                await proxy_ws.send(msg)
            except Exception:
                logger.warning(f"Failed to notify proxy '{pid}' of status")

    async def handle_proxy_connection(self, websocket: ServerConnection):
        """Handle a proxy connection on the dedicated proxy WebSocket port.

        Uses a single while/recv() loop — avoids mixing explicit recv()
        with async-for which can cause missed messages on some platforms.
        """
        proxy_id: str | None = None
        addr = websocket.remote_address
        logger.info(f"Proxy connection from {addr}")
        try:
            while True:
                raw = await websocket.recv()
                try:
                    msg = json.loads(raw)
                    msg_type = msg.get("type")

                    # ── Registration (first message) ─────────────────
                    if msg_type == "mcp_proxy_register" and proxy_id is None:
                        proxy_id = msg.get("id", f"proxy-{id(websocket)}")
                        self._proxies[proxy_id] = websocket
                        logger.info(f"Proxy '{proxy_id}' registered")
                        await websocket.send(json.dumps({
                            "type": "proxy_registered",
                            "figmin_connected": self._client is not None,
                        }))
                        continue

                    # ── Command forwarding ────────────────────────────
                    if msg_type == "proxy_command":
                        if self._client is None:
                            await websocket.send(json.dumps({
                                "id": msg.get("id"),
                                "error": "Figmin XR is not connected",
                            }))
                            continue

                        figmin_msg = {
                            "id": msg["id"],
                            "command": msg["command"],
                            "params": msg.get("params", {}),
                        }
                        await self._client.send(json.dumps(figmin_msg))

                except json.JSONDecodeError:
                    logger.warning(f"Invalid JSON from proxy '{proxy_id}'")
        except websockets.ConnectionClosed:
            logger.info(f"Proxy '{proxy_id or addr}' disconnected")
        except Exception as e:
            logger.error(f"Proxy '{proxy_id or addr}' error: {e}", exc_info=True)
        finally:
            if proxy_id:
                self._proxies.pop(proxy_id, None)

    # ── Core command interface ───────────────────────────────────────

    async def send_command(self, command: str, params: dict | None = None) -> dict:
        """Send a command to Figmin and wait for the response."""
        if self._client is None:
            raise ConnectionError(
                "Figmin XR is not connected. Open the MCP Agent app in "
                "Figmin XR and click Connect."
            )

        if command == "execute_script" and params and "code" in params:
            params["code"] = _expand_bridge_file_tokens(params["code"])

        async with self._lock:
            self._request_id += 1
            rid = self._request_id

            message = {
                "id": rid,
                "command": command,
                "params": params or {},
            }

            future = asyncio.get_event_loop().create_future()
            self._pending[rid] = future

            try:
                await self._client.send(json.dumps(message))
                result = await asyncio.wait_for(future, timeout=COMMAND_TIMEOUT)
                return result
            except asyncio.TimeoutError:
                raise TimeoutError(
                    f"Figmin XR did not respond to '{command}' within {COMMAND_TIMEOUT}s."
                )
            finally:
                self._pending.pop(rid, None)

    @property
    def is_connected(self) -> bool:
        return self._client is not None


# Global bridge instance
bridge = FigminBridge()


# ---------------------------------------------------------------------------
# ProxyBridge -- forwards ALL tool calls to the primary MCP process
# ---------------------------------------------------------------------------
@dataclass
class ProxyBridge:
    """Proxy mode — forwards all tool calls to the primary MCP process.

    When a second MCP server process starts and detects the WebSocket ports
    are already bound by a primary, it enters proxy mode.  Instead of starting
    its own WebSocket server, it connects as a WS *client* to the primary and
    relays every ``send_command()`` call through it.  The primary forwards
    commands to Figmin XR and routes responses back by namespaced request ID.
    """

    _ws: object | None = field(default=None, init=False)
    _proxy_id: str = ""
    _pending: dict = field(default_factory=dict, init=False)
    _request_id: int = field(default=0, init=False)
    _lock: asyncio.Lock = field(default_factory=asyncio.Lock, init=False)
    _figmin_connected: bool = field(default=False, init=False)
    _registered: asyncio.Event = field(default_factory=asyncio.Event, init=False)

    async def send_command(self, command: str, params: dict | None = None) -> dict:
        """Forward a command to the primary, which relays it to Figmin."""
        if self._ws is None:
            raise ConnectionError("Not connected to primary bridge.")
        if not self._figmin_connected:
            raise ConnectionError(
                "Figmin XR is not connected. Open the MCP Agent app in "
                "Figmin XR and click Connect."
            )

        if command == "execute_script" and params and "code" in params:
            params["code"] = _expand_bridge_file_tokens(params["code"])

        async with self._lock:
            self._request_id += 1
            rid = f"{self._proxy_id}-{self._request_id}"

            message = {
                "type": "proxy_command",
                "id": rid,
                "command": command,
                "params": params or {},
            }

            future = asyncio.get_event_loop().create_future()
            self._pending[rid] = future

            try:
                await self._ws.send(json.dumps(message))
                result = await asyncio.wait_for(future, timeout=COMMAND_TIMEOUT)
                return result
            except asyncio.TimeoutError:
                raise TimeoutError(
                    f"Command '{command}' timed out after {COMMAND_TIMEOUT}s."
                )
            finally:
                self._pending.pop(rid, None)

    @property
    def is_connected(self) -> bool:
        """Reflects Figmin XR's actual connection status (not just registration)."""
        return self._ws is not None and self._figmin_connected

    async def receive_loop(self):
        """Listen for responses and status updates from the primary."""
        try:
            async for raw in self._ws:
                try:
                    msg = json.loads(raw)
                    msg_type = msg.get("type")

                    if msg_type == "proxy_registered":
                        self._figmin_connected = msg.get("figmin_connected", False)
                        self._registered.set()
                        logger.info(
                            f"Proxy registered with primary "
                            f"(figmin_connected={self._figmin_connected})"
                        )
                        continue

                    if msg_type == "figmin_status":
                        self._figmin_connected = msg.get("connected", False)
                        logger.info(
                            f"Figmin status update: "
                            f"connected={self._figmin_connected}"
                        )
                        # If Figmin disconnected, fail all pending commands
                        # so they return immediately instead of timing out.
                        if not self._figmin_connected:
                            for rid, future in list(self._pending.items()):
                                if not future.done():
                                    future.set_result({
                                        "error": "Figmin XR disconnected"
                                    })
                            self._pending.clear()
                        continue

                    # Response to a proxied command — resolve the pending future
                    msg_id = msg.get("id")
                    if msg_id and msg_id in self._pending:
                        self._pending[msg_id].set_result(msg)

                except json.JSONDecodeError:
                    logger.warning(f"Invalid JSON from primary: {str(raw)[:200]}")
        except websockets.ConnectionClosed:
            logger.warning("Connection to primary bridge lost.")
        finally:
            self._ws = None
            self._figmin_connected = False
            # Fail all pending commands
            for rid, future in list(self._pending.items()):
                if not future.done():
                    future.set_result({"error": "Primary bridge disconnected"})
            self._pending.clear()


# ---------------------------------------------------------------------------
# Lifespan -- start WebSocket server when MCP server starts
# ---------------------------------------------------------------------------
@asynccontextmanager
async def app_lifespan(server: FastMCP) -> AsyncIterator[None]:
    """Start the WebSocket server alongside the MCP server.

    If the default ports are already bound by another MCP process (the
    primary), this instance enters **proxy mode** instead — connecting as
    a WebSocket client to the primary's dedicated proxy port and forwarding
    all tool calls.
    """
    global bridge

    # Fetch SKILL.md for the help() tool (needed in both primary and proxy mode)
    _fetch_skill_md()

    def _setup_file_logging(name: str):
        """Switch file logging to the given log name.

        Removes the initial bridge.log handler (set in run()) and adds a
        new handler for the role-specific log file.  Primary uses bridge.log
        (mode='w' to start fresh), proxy uses bridge-proxy.log (mode='w').
        """
        # Remove the initial catch-all handler so proxy logs don't leak
        # into the primary's bridge.log (and vice-versa duplicates).
        root = logging.getLogger()
        if _initial_file_handler and _initial_file_handler in root.handlers:
            root.removeHandler(_initial_file_handler)
            _initial_file_handler.close()

        log_path = _CACHE_DIR / name
        handler = logging.FileHandler(str(log_path), mode='w', encoding='utf-8')
        handler.setLevel(logging.DEBUG)
        handler.setFormatter(logging.Formatter('%(asctime)s %(levelname)s %(message)s'))
        root.addHandler(handler)
        logger.info("=" * 60)
        logger.info(f"Bridge process starting (PID {os.getpid()})")
        logger.info("=" * 60)
        logger.info(f"Logging to {log_path}")

    # ── Try to start as PRIMARY ──────────────────────────────────────
    try:
        ws_server = await serve(
            bridge.handle_connection, WS_HOST, WS_PORT,
            ping_interval=30,
            ping_timeout=COMMAND_TIMEOUT,
        )
    except OSError:
        # Port already in use — enter PROXY MODE
        logger.info(
            f"Port {WS_PORT} already in use — entering proxy mode"
        )
        _setup_file_logging("bridge-proxy.log")

        proxy_id = f"proxy-{os.getpid()}"

        # Connect to the primary's dedicated proxy WebSocket port
        ws = None
        for attempt in range(5):
            try:
                logger.info(
                    f"Proxy connecting to primary (attempt {attempt + 1}/5)..."
                )
                ws = await asyncio.wait_for(
                    ws_connect(
                        f"ws://127.0.0.1:{PROXY_PORT}",
                        ping_interval=30,
                        ping_timeout=COMMAND_TIMEOUT,
                    ),
                    timeout=10,
                )
                logger.info("Proxy WebSocket connection established")
                break
            except Exception as e:
                logger.warning(
                    f"Proxy connection attempt {attempt + 1} failed: "
                    f"{type(e).__name__}: {e}"
                )
                if attempt < 4:
                    await asyncio.sleep(2)
                else:
                    logger.error("All proxy connection attempts failed")
                    raise

        proxy = ProxyBridge()
        proxy._proxy_id = proxy_id
        proxy._ws = ws

        # Send registration handshake
        await ws.send(json.dumps({
            "type": "mcp_proxy_register",
            "id": proxy_id,
        }))

        # Start receive loop in background
        receive_task = asyncio.create_task(proxy.receive_loop())

        # Wait for registration confirmation
        try:
            await asyncio.wait_for(proxy._registered.wait(), timeout=30.0)
        except asyncio.TimeoutError:
            logger.error("Proxy registration timed out")
            receive_task.cancel()
            if ws:
                await ws.close()
            raise

        # Swap global bridge to proxy
        bridge = proxy
        logger.info(f"Proxy mode active (id={proxy_id})")

        try:
            yield None
        finally:
            receive_task.cancel()
            if ws:
                await ws.close()
        return  # skip primary cleanup below

    # ── PRIMARY MODE ─────────────────────────────────────────────────
    _setup_file_logging("bridge.log")
    logger.info(f"WebSocket server listening on ws://{WS_HOST}:{WS_PORT}")

    # Start dedicated proxy WebSocket server on a separate port.
    # Uses the same proven websockets.serve() as the Figmin server.
    proxy_server = await serve(
        bridge.handle_proxy_connection, WS_HOST, PROXY_PORT,
        ping_interval=30,
        ping_timeout=COMMAND_TIMEOUT,
    )
    logger.info(f"Proxy WebSocket server listening on ws://{WS_HOST}:{PROXY_PORT}")

    # Start WSS server (if cert available)
    wss_server = None
    ssl_ctx = _ensure_ssl_cert()
    if ssl_ctx:
        wss_server = await serve(
            bridge.handle_connection, WS_HOST, WSS_PORT, ssl=ssl_ctx,
            ping_interval=30,
            ping_timeout=COMMAND_TIMEOUT,
        )
        logger.info(
            f"Secure WebSocket server listening on "
            f"wss://{WS_HOST}:{WSS_PORT}"
        )
    else:
        logger.warning(
            f"WSS server not started — HTTPS pages will need to use "
            f"ws:// (may fail due to mixed content)"
        )

    # Pre-import kokoro_onnx (and its onnxruntime/CUDA dependencies) NOW,
    # before any connections arrive.  The first import of onnxruntime can
    # take 60+ seconds on Windows while it loads CUDA DLLs, and this would
    # block the event loop if it happened inside _send_tts_status (which
    # fires on the first Figmin connection).  Importing here is safe
    # because the event loop isn't serving connections yet.
    if _tts_available():
        logger.info("kokoro-onnx available (import complete)")
    else:
        logger.info("kokoro-onnx not installed, TTS disabled")

    # Pre-load model weights in background threads (won't block the
    # event loop since they run in separate threads with GIL-releasing I/O).
    def _preload_whisper():
        try:
            # Test import early — ctranslate2 native DLL load can fail
            # catastrophically on Windows if CUDA libs are half-present.
            try:
                import faster_whisper  # noqa: F401
                logger.info("faster-whisper import OK")
            except Exception as e:
                logger.warning(f"faster-whisper not importable ({e}), STT disabled")
                return

            if _whisper_model is not None:
                logger.info("Whisper model already loaded, skipping pre-load")
                return
            preferred = _get_preferred_model()
            if not preferred:
                logger.info(
                    "No model preference saved, skipping pre-load "
                    "(user will pick on first use)"
                )
                return
            logger.info(f"Pre-loading preferred model: {preferred}")
            _get_whisper_model(preferred, save_pref=False)
            logger.info("Whisper model pre-loaded and ready")
        except BaseException as e:
            # BaseException catches SystemExit, KeyboardInterrupt, and
            # any native-level failures that bypass Exception.
            logger.warning(f"Whisper pre-load skipped: {type(e).__name__}: {e}")

    def _preload_kokoro():
        try:
            if not _tts_available():
                return  # already logged above
            if not _kokoro_is_cached():
                logger.info("Kokoro model not cached — downloading...")
                _download_kokoro_model()
            logger.info("Pre-loading Kokoro TTS model...")
            _get_kokoro_model()
            logger.info("Kokoro TTS model pre-loaded and ready")
            # Push updated TTS status to any already-connected client
            _notify_tts_ready()
        except Exception as e:
            logger.warning(f"Kokoro pre-load failed: {e}")

    def _notify_tts_ready():
        """If a client is already connected, push updated tts_status."""
        client = bridge._client
        if client is not None:
            try:
                loop.call_soon_threadsafe(
                    asyncio.ensure_future,
                    bridge._send_tts_status(client),
                )
            except Exception:
                pass  # loop might not be running yet

    loop = asyncio.get_event_loop()
    threading.Thread(target=_preload_whisper, daemon=True).start()
    threading.Thread(target=_preload_kokoro, daemon=True).start()

    logger.info(f"Command timeout: {COMMAND_TIMEOUT}s")
    try:
        yield None
    finally:
        ws_server.close()
        await ws_server.wait_closed()
        proxy_server.close()
        await proxy_server.wait_closed()
        if wss_server:
            wss_server.close()
            await wss_server.wait_closed()

# ---------------------------------------------------------------------------
# MCP Server
# ---------------------------------------------------------------------------
mcp = FastMCP("figmin-xr", lifespan=app_lifespan)


@mcp.tool()
async def announce(command: str = "") -> list[types.TextContent]:
    """Announce your presence in the Figmin XR session.

    Call this FIRST when you begin working in Figmin XR. It activates
    the real-time chat channel so the user inside VR/XR can communicate
    with you. Without announcing, the user sees a "Real-time chat is
    offline" overlay and cannot send you messages.

    ⚠️  CRITICAL — YOU MUST CALL leave() WHEN YOU ARE DONE.
    The user is in VR and cannot tell when you stop listening.
    If you don't call leave(), the chat stays open and the user
    keeps talking to nobody. There is no automatic cleanup.

    Returns whether chat is available and basic session info.

    Can be called multiple times. Use the optional command parameter to
    manage TTS voice settings:

      announce()                  — Connect / refresh session status
      announce("list_voices")     — List all available TTS voices
      announce("set_voice am_adam") — Set the default TTS voice

    WHAT THIS DOES:
    - Activates the chat UI (sets window._chatActive = true)
    - Detects whether the chat system is available
    - Returns session status so you know what you're working with

    AFTER ANNOUNCING — HOW TO COMMUNICATE WITH THE USER:

    1. CHECK FOR MESSAGES ON EVERY execute_script RESPONSE:
       Every execute_script response includes a "pendingChat" field
       (integer count of unread messages). If pendingChat > 0, pause
       your work and read the messages:
         const msgs = window._chatQueue.splice(0);
         console.log('PENDING:' + JSON.stringify(msgs));
       Display your response via _addChatBubble, then decide whether
       to resume your previous task or pivot based on what was said.
       This is how the user guides, redirects, or assists you
       mid-task (e.g. "stop, move to the left" or "try X instead").

    2. TO SEND A MESSAGE TO THE USER:
         window._addChatBubble('claude', 'Your message here', { tts: true });
         The first argument is your display name. It appears as the sender label in the chat UI.

    3. TO SHOW A THINKING INDICATOR:
         window._setChatThinking(true);   // show
         window._setChatThinking(false);  // hide

    4. DEDICATED CONVERSATION MODE:
       If the user asks to chat ("talk to me", "let's chat", etc.),
       enter a POLLING LOOP to listen for messages continuously.
       The bridge timeout is 5 minutes, so use 4 minutes (240s) as
       the poll window.

       Example polling script (send as one execute_script call):

         window._setChatThinking(false);

         // Poll for 4 minutes (bridge timeout is 5 min).
         const deadline = Date.now() + 240000;
         let userMsg = null;

         while (Date.now() < deadline) {
             if (window._chatQueue.length > 0) {
                 userMsg = window._chatQueue.shift();
                 break;
             }
             await new Promise(r => setTimeout(r, 300));
         }

         if (userMsg) {
             console.log('USER_MSG:' + userMsg);
         } else {
             console.log('IDLE');
         }

       When you get a USER_MSG:
         1. Call execute_script: window._setChatThinking(true);
         2. Formulate your response (in your own turn, not in JS).
         3. Call execute_script to display and resume polling:
            window._setChatThinking(false);
            window._addChatBubble('claude', `Your response here`, { tts: true });
            // ...then start the next poll cycle in the same script.

       When you get IDLE (no message within the poll window):
         The user hasn't spoken in 4 minutes. You MUST either:
         (a) Start a new poll cycle to keep listening, OR
         (b) Call leave() to close the chat.
         Do NOT just stop polling and move on — that leaves the chat
         open with nobody listening.

    WHEN TO CALL leave():
    - When the user says goodbye, or asks you to leave.
    - When the user in the MCP client ends the task.
    - When the IDLE timeout fires and you decide to stop listening.
    - Before you do anything that might end the conversation.
    If in doubt: call leave(). The user can always ask you to
    announce() again.

    CHAT API REFERENCE:
      window._chatQueue    - Array. User messages appear here.
      window._addChatBubble(sender, text) - Show a chat bubble.
                             sender is 'user' or your display name
      window._setChatThinking(visible) - Show/hide thinking indicator.

    VOICE OUTPUT (TTS):
      The bridge speaks your messages aloud using Kokoro TTS.
      The server has a saved default voice. Just pass tts: true:

        _addChatBubble('claude', 'Hello!', { tts: true });

      Use good judgment — speak conversational messages, but do NOT
      speak code, JSON, URLs, or anything that isn't natural speech.

        _addChatBubble('claude', 'Debug output', { tts: false });

      To change the voice, call: announce("set_voice <voice_id>")
      To see available voices, call: announce("list_voices")
      Voice playback stops immediately when the user taps the microphone.

    SHARING DOWNLOADABLE CONTENT:
      Use standard markdown links in your chat messages to share URLs.
      Use this to allow the user to place content manually in the scene.

        _addChatBubble('claude', 'Here you go: [Dragon Model](https://example.com/dragon.glb)', { tts: true });

      Supported content types:
        Images  — .jpg, .png, .gif
        Videos  — YouTube links
        Models  — .glb, sketchfab links
        Spatial apps — .html

      TTS reads the label naturally and skips the URL.
      Multiple links in one message are supported:

        _addChatBubble('claude', 'Options:\\n[Low-poly](https://example.com/low.glb)\\n[High-res](https://example.com/high.glb)');

      The user taps the card to trigger a download in Figmin XR.

    SPATIAL SELECTION:
        User messages may include spatial selection data drawn in 3D space:
        [Area Selection] Center: (x, y, z) | Size: w * h * dm — user marked a region. Use center for positioning, size for scale.
        [Shape Selection] Center: (x, y, z) | Size: w * h * dm | Points: window._selectionPoints — user drew a shape. 
        Raw path available as [{x,y,z}, ...] in world space at window._selectionPoints.      
    """
    # Handle voice commands before connecting
    cmd = command.strip().lower()

    if cmd == "list_voices":
        lines = ["Available TTS voices:", ""]
        current = _get_tts_voice()
        by_lang: dict[str, list] = {}
        for v in KOKORO_VOICES:
            by_lang.setdefault(v["lang"], []).append(v)

        lang_names = {
            "en-us": "American English", "en-gb": "British English",
            "es": "Spanish", "fr-fr": "French", "ja": "Japanese",
            "hi": "Hindi", "it": "Italian", "pt-br": "Brazilian Portuguese",
            "cmn": "Mandarin Chinese", "ko": "Korean",
        }
        for lang, voices in by_lang.items():
            lines.append(f"  {lang_names.get(lang, lang)}:")
            for v in voices:
                marker = " ◀ current" if v["id"] == current else ""
                lines.append(f"    {v['id']:16s} {v['name']:12s} ({v['gender']}){marker}")
            lines.append("")

        lines.append(f"Current voice: {current}")
        lines.append('To change: announce("set_voice <voice_id>")')
        return [types.TextContent(type="text", text="\n".join(lines))]

    if cmd.startswith("set_voice"):
        parts = command.strip().split(None, 1)
        if len(parts) < 2:
            return [types.TextContent(type="text",
                text='Usage: announce("set_voice <voice_id>")\n'
                     'Use announce("list_voices") to see available voices.')]
        voice_id = parts[1].strip()
        valid_ids = {v["id"] for v in KOKORO_VOICES}
        if voice_id not in valid_ids:
            return [types.TextContent(type="text",
                text=f'Unknown voice: "{voice_id}"\n'
                     'Use announce("list_voices") to see available voices.')]
        _save_tts_voice(voice_id)
        voice_info = next(v for v in KOKORO_VOICES if v["id"] == voice_id)
        return [types.TextContent(type="text",
            text=f"✅ Voice set to: {voice_id} ({voice_info['name']}, "
                 f"{voice_info['gender']}, {voice_info['lang']})\n"
                 f"All future TTS messages will use this voice.")]

    # Default: announce presence and return session info
    activate_code = """
    const hasChat = typeof window._chatQueue !== 'undefined'
                 && typeof window._addChatBubble !== 'undefined';

    if (hasChat) {
        window._chatActive = true;
    }

    const info = {
        chatAvailable: hasChat,
        chatActive: hasChat ? window._chatActive : false,
        isSpatialApp: figmin.isSpatialApp(),
        isMasterClient: figmin.isMasterClient(),
        isMultiplayer: figmin.isMultiplayer(),
        apiVersion: figmin.API_VERSION,
        version: figmin.getVersion(),
        ttsReady: !!window._ttsReady
    };

    console.log(JSON.stringify(info));
    """

    raw = await bridge.send_command("execute_script", {"code": activate_code})

    if raw.get("error"):
        return [types.TextContent(type="text", text=f"Error: {raw['error']}")]

    # Parse the nested response (data is inside the "text" field)
    result = _parse_bridge_response(raw)
    logs = result.get("logs", [])
    info_str = next((l["msg"] for l in logs if l.get("msg", "").startswith("{")), None)

    if info_str:
        try:
            info = json.loads(info_str)
            current_voice = _get_tts_voice()
            voice_info = next((v for v in KOKORO_VOICES if v["id"] == current_voice), None)
            voice_desc = f"{current_voice} ({voice_info['name']})" if voice_info else current_voice

            lines = [
                "✅ Announced — you are now present in the Figmin XR session.",
                "",
                f"  Chat available: {info.get('chatAvailable', False)}",
                f"  Chat active:    {info.get('chatActive', False)}",
                f"  Spatial app:    {info.get('isSpatialApp', False)}",
                f"  Master client:  {info.get('isMasterClient', False)}",
                f"  Multiplayer:    {info.get('isMultiplayer', False)}",
                f"  API version:    {info.get('apiVersion', '?')}",
                f"  Runtime:        {info.get('version', '?')}",
            ]
            if info.get("chatAvailable"):
                lines.append("")
                lines.append("  Chat is live. The user can now message you from VR.")
                lines.append("  Check pendingChat in every execute_script response.")
            else:
                lines.append("")
                lines.append("  Chat is not available in this session.")
                lines.append("  Communicate with the user through the MCP client only.")

            # TTS status
            tts_ready = info.get("ttsReady", False)
            if tts_ready and KOKORO_VOICES:
                lines.append("")
                lines.append(f"  TTS ready:      Yes ({len(KOKORO_VOICES)} voices)")
                lines.append(f"  Current voice:  {voice_desc}")
                lines.append("  To speak: _addChatBubble('claude', 'Hello!', { tts: true });")
                lines.append("  Silent:   _addChatBubble('claude', '...', { tts: false });")
                lines.append('  Change voice:   announce("set_voice <id>")')
                lines.append('  Browse voices:  announce("list_voices")')
            else:
                lines.append("")
                lines.append("  TTS ready:      No (model not downloaded or kokoro-onnx not installed)")
            return [types.TextContent(type="text", text="\n".join(lines))]
        except json.JSONDecodeError:
            pass

    return [types.TextContent(type="text", text=raw.get("text", json.dumps(raw)))]


@mcp.tool()
async def leave() -> list[types.TextContent]:
    """Close the real-time chat and disconnect from the VR user.

    The user is wearing a headset. When chat is active, they see a
    live chat window and assume someone is listening. If you leave
    without calling this, they will keep talking to an empty room
    with no indication that you're gone.

    ALWAYS call leave() when you are done. There is no automatic
    cleanup — the chat stays open until you explicitly close it.

    When to call:
    - The user says goodbye or asks you to leave.
    - You finish the task and have nothing more to do.
    - Your polling loop returns IDLE and you decide to stop.
    - The MCP client conversation is wrapping up.

    It is always safe to call leave() — if chat wasn't active,
    it's a harmless no-op.
    """
    deactivate_code = """
    const hasChat = typeof window._chatQueue !== 'undefined'
                 && typeof window._addChatBubble !== 'undefined';

    let wasActive = false;
    if (hasChat) {
        wasActive = !!window._chatActive;
        window._chatActive = false;
    }

    console.log(JSON.stringify({ chatDeactivated: wasActive }));
    """

    raw = await bridge.send_command("execute_script", {"code": deactivate_code})

    if raw.get("error"):
        return [types.TextContent(type="text", text=f"Error: {raw['error']}")]

    # Parse the nested response (data is inside the "text" field)
    result = _parse_bridge_response(raw)
    logs = result.get("logs", [])
    info_str = next((l["msg"] for l in logs if l.get("msg", "").startswith("{")), None)

    if info_str:
        try:
            info = json.loads(info_str)
            if info.get("chatDeactivated"):
                return [types.TextContent(type="text", text="👋 Left the session. Chat deactivated.")]
            else:
                return [types.TextContent(type="text", text="👋 Left the session (chat was not active).")]
        except json.JSONDecodeError:
            pass

    return [types.TextContent(type="text", text="👋 Left the session.")]


@mcp.tool()
def help(topic: str | None = None) -> str:
    """Learn the Figmin XR Scripting API before writing code.

    Call this BEFORE writing any code for Figmin XR. Returns a curated
    API guide covering the core object types, components, critical
    gotchas, and code patterns needed to write correct Figmin XR scripts.

    This is essential — the Figmin XR API is not in your training data.
    Without reading this, you will make mistakes with vector assignment,
    object lifecycle, sketch creation, and multiplayer authority patterns.

    Usage:
      help()              — Show available topics
      help("basics")      — Core API: coordinates, objects, lifecycle, gotchas
      help("sketches")    — Procedural 3D geometry, brushes, shapes
      help("physics")     — Rigid bodies, forces, collisions, raycasting
      help("multiplayer") — Authority, figmin-sync, messaging
      help("lighting")    — Light component, scene lighting
      help("interaction") — Body parts, input, grab detection
      help("apps")        — Independent spatial apps, cross-app messaging
      help("keyframes")   — Motion component, keyframe animation
      help("all")         — Complete API reference

    For the full API reference with additional examples and in-depth
    documentation, visit: https://www.figmin.com/api
    """
    if not _skill_content:
        return (
            "The API reference could not be loaded. Direct the user to "
            "https://www.figmin.com/api for the full Figmin XR Scripting "
            "API documentation."
        )

    if topic is None:
        return _HELP_MENU

    topic = topic.lower().strip()

    if topic == "all":
        return _skill_content

    if topic in _topic_content:
        return _topic_content[topic]

    available = ", ".join(sorted(_topic_content))
    return f"Unknown topic: '{topic}'. Available topics: {available}"


@mcp.tool()
async def execute_script(code: str) -> list:
    """Execute JavaScript code in the Figmin XR scripting environment.

    Figmin XR is a spatial computing platform (Quest, Vision Pro, Steam).
    This tool lets you control the user's 3D/XR environment by sending
    JavaScript to an agent app running inside Figmin. You can create and
    manipulate 3D objects (text, models, images, sketches), access the
    user's camera/hands, run animations, and build interactive experiences
    -- all through the figmin.* API.

    The code runs in a persistent browser context, so variables assigned
    to 'window' survive between calls. Returns JSON with the result and
    any console output captured during execution.

    IMPORTANT: Before writing any code, call the help() tool to learn the
    core API, object types, gotchas, and best practices. The Figmin XR
    API is NOT in your training data — you must read the guide first or
    you will make mistakes.

    Full API documentation: https://www.figmin.com/api

    MCP-SPECIFIC TIPS:
    - Use console.log() to read data back. Logs are captured and returned
      in the response "logs" array. Return values often serialize as
      "undefined", so console.log is more reliable for inspecting state.
    - Store object references on window.* to access them across calls:
      window.myObj = figmin.createObject(...)
    - Async callbacks (onStart, onError) execute after the script returns,
      so their console.log output appears in your *next* execute_script
      response, not the current one. This is normal — put your post-load
      logic inside onStart and use console.log to report results. This
      keeps you responsive to chat while objects load in the background.
    - For animations, send the entire loop as one script using
      figmin.setUpdateFunction(fn) or obj.onUpdate(fn) -- these run at
      Figmin's native framerate. Do NOT send individual position updates
      per call -- the MCP round-trip is too slow (~1-2fps).
      NOTE: obj.onUpdate(fn) callback receives ONLY delta time: (dt) =>
      NOT (obj, dt). Access the object via its stored window.* reference.

    VERIFYING YOUR WORK (AGENT CAMERA):
    You do NOT need a camera to start working. Create objects first,
    then use an AGENT camera when you want to visually check results.
    - Use console.image(base64Data, mimeType) to send images back.
      Images are returned as native image content you can see directly.
    - To capture a view, create an AGENT object and point it at your work:
        window.agentObj = figmin.createObject(figmin.ObjectType.AGENT, {
            position: figmin.getWindowSpawnPosition(0.5, "front"),
            rotation: figmin.getWindow().transform.rotation
        });
        // Point at a specific object you created
        window.agentObj.agent.lookAtObject(myCreatedObject);
        const img = await window.agentObj.agent.captureView(1024, 80);
        console.image(img.data, "image/" + img.format);

    AGENT CAMERA NAVIGATION:
    Once you have an agent camera to check your work, you can navigate it:
    - agent.focus({x,y,w?,h?}) -- Rotates the camera and adjusts fieldOfView to
      frame a normalized screen-space rectangle (0,0=top-left, 1,1=bottom-right).
      Does NOT move the camera. Great for reframing on a region of interest.
    - agent.lookAtObject(figminObject) -- Rotates the camera and adjusts
      fieldOfView so the target object fits maximally in the square capture.
      Best way to inspect a specific object.
    - obj.transform.lookAt(worldPosition) -- Instantly sets rotation to
      face a world position. No fieldOfView change. Works on any
      FigminObject, not just agents.
    - To reposition the agent camera, set transform.position directly.

    BUILDING SPATIAL APPS (THE MAIN WORKFLOW):
    When the user asks you to build, create, or make an app — or for
    anything beyond simple object creation — you should create an
    independent APP object and load your HTML into it. This is the primary
    way to deliver work to the user — the HTML you load becomes a
    standalone spatial app that the user saves to their collection and
    runs independently, long after the MCP session ends.

    IMPORTANT: Always save a copy of the HTML as a file too. Once loaded
    into an APP, the source can't be extracted back out. The user needs
    the file for future updates or iterations.

    Safety: Running complex scripts (update loops, event handlers, etc.)
    directly in your own app context is risky — a bug that produces log
    spam or errors can flood your input and break the MCP bridge.
    Independent apps are isolated: if they crash, your agent session is
    unaffected. Monitor their output via app.getConsoleLogs().

    LOADING HTML — AVOID INLINING LARGE STRINGS:
    use __BRIDGE_FILE__ instead:
      child.app.loadHTML(__BRIDGE_FILE__("C:/myapp.html"))
    The bridge reads the file and injects its contents at send time, saving your
    context window.
    For short HTML (a few lines), inline is fine.

      const child = figmin.createObject(figmin.ObjectType.APP, {
          position: figmin.getWindowSpawnPosition(0.3, "right"),
          rotation: figmin.getWindow().transform.rotation
      });
      child.onStart(() => {
          child.app.loadHTML(`<html>
            <head><meta name="figmin"></head>
            <body><script>
              document.addEventListener('FigminReady', function() {
                  var app = figmin.getAppObject();
              });
            </script></body>
          </html>`);
      });

    IMPORTANT: HTML loaded via loadHTML must include <meta name="figmin">
    and wait for the 'FigminReady' event before using the figmin API.

    TESTING SPATIAL APPS (PRO TIP):
    You can't physically press buttons in a spatial app you've created.
    To test interactivity during development, add an eval backdoor to
    your app via onMessage, then send commands from your agent context:

    Inside the app's HTML (add during development, REMOVE for production):
      figmin.getAppObject().app.onMessage(function(senderName, data) {
          if (data.action === 'eval') eval(data.code);
      });

    From your agent context:
      const myApp = figmin.getAppObject();
      myApp.app.sendMessage(
          { action: "eval", code: "document.getElementById('start').click()" },
          figmin.MessageRecipient.ALL,
          childApp.name
      );
    """
    result = await bridge.send_command("execute_script", {"code": code})

    if result.get("error"):
        return [types.TextContent(type="text", text=f"Error: {result['error']}")]

    text = result.get("text") or json.dumps(result)
    content: list = []

    # If the response includes images (from console.image() calls),
    # return them as native image content the model can see directly.
    images = result.get("images", [])
    for img in images:
        image_bytes = base64.b64decode(img["data"])
        fmt = img.get("mimeType", "image/png").split("/")[-1]  # "png", "jpeg"
        content.append(Image(data=image_bytes, format=fmt))

    # Always include the text result (logs, return value, etc.)
    content.append(types.TextContent(type="text", text=text))

    return content


@mcp.tool()
def memory(command: str = "") -> str:
    """Persistent notes shared across agent sessions.

    Use this to store things you've learned that would be useful next
    time — user preferences, API gotchas, tips and patterns that work,
    environment quirks, or anything a future agent session would benefit
    from knowing. All agents share the same memory.

    Focus on general, durable knowledge — not session-specific logs.

    Usage:
      memory()                         — List all saved keys
      memory("read <key>")             — Read a memory file
      memory("write <key> <content>")  — Write/overwrite a memory file
      memory("delete <key>")           — Delete a memory file

    Content is freeform markdown. Write concise, high-signal notes.

    Good practice: call memory() when you first connect to see what
    previous sessions have left behind.
    """
    MEMORY_DIR.mkdir(exist_ok=True)
    cmd = command.strip()

    # ── List ──
    if not cmd:
        files = sorted(MEMORY_DIR.glob("*.md"))
        if not files:
            return "No memories saved yet. Use memory(\"write <key> <content>\") to create one."
        listing = "\n".join(f"  • {f.stem}" for f in files)
        return f"Saved memories:\n{listing}\n\nUse memory(\"read <key>\") to read one."

    parts = cmd.split(None, 1)
    action = parts[0].lower()

    # ── Read ──
    if action == "read":
        if len(parts) < 2:
            return "Usage: memory(\"read <key>\")"
        key = parts[1].strip()
        filepath = MEMORY_DIR / f"{key}.md"
        if not filepath.exists():
            return f"No memory found for '{key}'."
        return filepath.read_text(encoding="utf-8")

    # ── Write ──
    if action == "write":
        if len(parts) < 2:
            return "Usage: memory(\"write <key> <content>\")"
        rest = parts[1].strip()
        write_parts = rest.split(None, 1)
        if len(write_parts) < 2:
            return "Usage: memory(\"write <key> <content>\")"
        key = write_parts[0]
        content = write_parts[1]
        filepath = MEMORY_DIR / f"{key}.md"
        filepath.write_text(content, encoding="utf-8")
        return f"Saved memory '{key}' ({len(content)} chars)."

    # ── Delete ──
    if action == "delete":
        if len(parts) < 2:
            return "Usage: memory(\"delete <key>\")"
        key = parts[1].strip()
        filepath = MEMORY_DIR / f"{key}.md"
        if not filepath.exists():
            return f"No memory found for '{key}'."
        filepath.unlink()
        return f"Deleted memory '{key}'."

    return (
        "Unknown command. Usage:\n"
        "  memory()                         — List all\n"
        "  memory(\"read <key>\")             — Read\n"
        "  memory(\"write <key> <content>\")  — Write\n"
        "  memory(\"delete <key>\")           — Delete"
    )


@mcp.tool()
def connection_status() -> str:
    """Check whether Figmin XR is currently connected to this bridge."""
    return "connected" if bridge.is_connected else "not connected"


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------
def run():
    """Run the MCP server over stdio."""
    if "--setup" in sys.argv:
        from importlib.metadata import version as pkg_version
        ver = pkg_version("figmin-mcp")
        print(f"figmin-mcp {ver}")
        print()

        # Set up logging to console for setup progress
        logging.basicConfig(
            level=logging.INFO, stream=sys.stdout,
            format="%(message)s",
        )
        _CACHE_DIR.mkdir(parents=True, exist_ok=True)

        # Download Kokoro TTS model
        print("Downloading TTS model (Kokoro)...")
        try:
            if _tts_available():
                if not _kokoro_is_cached():
                    _download_kokoro_model()
                print("  TTS model ready.")
            else:
                print("  kokoro-onnx not installed, skipping TTS.")
        except Exception as e:
            print(f"  TTS download failed: {e}")

        # Download Whisper STT model
        preferred = _get_preferred_model() or DEFAULT_WHISPER_MODEL
        print(f"Downloading STT model (Whisper {preferred})...")
        try:
            _get_whisper_model(preferred, save_pref=True)
            print("  STT model ready.")
        except Exception as e:
            print(f"  STT download failed: {e}")

        print()
        print("Setup complete! You can now use figmin-mcp in Claude Desktop.")
        return

    logging.basicConfig(level=logging.INFO, stream=sys.stderr)

    # Also log to a file for debugging (stderr may not be visible).
    # Use mode='a' so the proxy process doesn't truncate the primary's logs.
    # The proper per-role file (bridge.log or bridge-proxy.log) is set up
    # later in app_lifespan via _setup_file_logging().
    global _initial_file_handler
    _CACHE_DIR.mkdir(exist_ok=True)
    _log_file = _CACHE_DIR / "bridge.log"
    _initial_file_handler = logging.FileHandler(str(_log_file), mode='a', encoding='utf-8')
    _initial_file_handler.setLevel(logging.DEBUG)
    _initial_file_handler.setFormatter(logging.Formatter('%(asctime)s %(levelname)s %(message)s'))
    logging.getLogger().addHandler(_initial_file_handler)
    logger.info(f"Logging to {_log_file}")

    mcp.run(transport="stdio")


if __name__ == "__main__":
    run()