"""
HOLD Live Mesh Server

Streams HOLD decision data as glTF meshes in real-time.
Games connect via HTTP or WebSocket, receive updated meshes each inference.

Endpoints:
  GET /hold.glb          - Latest HOLD snapshot as glTF
  GET /hold.json         - Latest HOLD data as JSON
  WS  /stream            - WebSocket stream of hold.glb updates
  
Integration:
  Unity:   UnityWebRequest to poll /hold.glb, import with GLTFast
  Unreal:  RuntimeMeshImporter plugin
  Godot:   HTTPRequest + GLTFDocument.append_from_buffer
"""

import asyncio
import json
import struct
import base64
from typing import Dict, Any, Optional, Callable
from pathlib import Path
import threading
import time

# HTTP server
try:
    from aiohttp import web
    HAS_AIOHTTP = True
except ImportError:
    HAS_AIOHTTP = False

# WebSocket
try:
    import websockets
    HAS_WEBSOCKETS = True
except ImportError:
    HAS_WEBSOCKETS = False

from .gltf_export import export_hold_to_gltf, _numpy_to_list


class HoldMeshServer:
    """
    Live mesh server for HOLD data.
    
    Usage:
        server = HoldMeshServer(port=8766)
        server.start()
        
        # During inference loop:
        server.update(hold_data)  # Pushes to connected clients
        
        # In game:
        # GET http://localhost:8766/hold.glb
    """
    
    def __init__(self, port: int = 8766, host: str = "localhost"):
        self.port = port
        self.host = host
        self._latest_hold: Optional[Dict[str, Any]] = None
        self._latest_glb: Optional[bytes] = None
        self._update_count = 0
        self._ws_clients = set()
        self._running = False
        self._thread = None
        self._loop = None
        
    def update(self, hold_data: Dict[str, Any]):
        """
        Update with new HOLD data. Broadcasts to all connected clients.
        """
        self._latest_hold = _numpy_to_list(hold_data)
        self._update_count += 1
        
        # Export to GLB in memory
        import tempfile
        import os
        
        with tempfile.NamedTemporaryFile(suffix='.glb', delete=False) as f:
            temp_path = f.name
        
        try:
            export_hold_to_gltf(hold_data, temp_path, binary=True)
            with open(temp_path, 'rb') as f:
                self._latest_glb = f.read()
        finally:
            if os.path.exists(temp_path):
                os.unlink(temp_path)
        
        # Notify WebSocket clients
        if self._loop and self._ws_clients:
            asyncio.run_coroutine_threadsafe(
                self._broadcast_update(),
                self._loop
            )
    
    async def _broadcast_update(self):
        """Broadcast update to all WebSocket clients."""
        if not self._ws_clients or not self._latest_glb:
            return
        
        # Send binary GLB
        dead_clients = set()
        for ws in self._ws_clients:
            try:
                await ws.send(self._latest_glb)
            except Exception:
                dead_clients.add(ws)
        
        self._ws_clients -= dead_clients
    
    def start(self):
        """Start the server in a background thread."""
        if self._running:
            return
        
        self._running = True
        self._thread = threading.Thread(target=self._run_server, daemon=True)
        self._thread.start()
        
        # Wait for server to be ready
        time.sleep(0.5)
        print(f"\n🎮 HOLD Mesh Server running:")
        print(f"   HTTP: http://{self.host}:{self.port}/hold.glb")
        print(f"   JSON: http://{self.host}:{self.port}/hold.json")
        if HAS_WEBSOCKETS:
            print(f"   WS:   ws://{self.host}:{self.port}/stream")
        print()
    
    def stop(self):
        """Stop the server."""
        self._running = False
        if self._loop:
            self._loop.call_soon_threadsafe(self._loop.stop)
    
    def _run_server(self):
        """Run the async server."""
        self._loop = asyncio.new_event_loop()
        asyncio.set_event_loop(self._loop)
        
        if HAS_AIOHTTP:
            self._loop.run_until_complete(self._run_aiohttp())
        else:
            # Fallback to basic HTTP server
            self._loop.run_until_complete(self._run_basic_http())
    
    async def _run_aiohttp(self):
        """Run aiohttp server."""
        app = web.Application()
        app.router.add_get('/hold.glb', self._handle_glb)
        app.router.add_get('/hold.json', self._handle_json)
        app.router.add_get('/stream', self._handle_ws)
        app.router.add_get('/', self._handle_index)
        
        runner = web.AppRunner(app)
        await runner.setup()
        site = web.TCPSite(runner, self.host, self.port)
        await site.start()
        
        while self._running:
            await asyncio.sleep(0.1)
        
        await runner.cleanup()
    
    async def _handle_index(self, request):
        """Serve 3D viewer page."""
        html = f"""<!DOCTYPE html>
<html>
<head>
    <title>HOLD Mesh Viewer</title>
    <style>
        * {{ margin: 0; padding: 0; box-sizing: border-box; }}
        body {{ 
            background: #0a0a12; 
            color: #eee; 
            font-family: 'Consolas', monospace;
            overflow: hidden;
        }}
        #info {{
            position: absolute;
            top: 10px;
            left: 10px;
            background: rgba(0,0,0,0.7);
            padding: 15px;
            border-radius: 8px;
            z-index: 100;
            max-width: 300px;
        }}
        #info h1 {{ font-size: 18px; color: #0ff; margin-bottom: 10px; }}
        #info p {{ font-size: 12px; margin: 5px 0; }}
        #stats {{
            position: absolute;
            top: 10px;
            right: 10px;
            background: rgba(0,0,0,0.7);
            padding: 15px;
            border-radius: 8px;
            z-index: 100;
            font-size: 14px;
        }}
        .action {{ margin: 3px 0; }}
        .bar {{ 
            display: inline-block;
            height: 12px;
            background: linear-gradient(90deg, #0f0, #ff0);
            margin-left: 5px;
            vertical-align: middle;
        }}
        #canvas {{ width: 100vw; height: 100vh; display: block; }}
        .selected {{ color: #0ff; font-weight: bold; }}
    </style>
</head>
<body>
    <div id="info">
        <h1>🎮 HOLD Mesh Viewer</h1>
        <p>Live AI decision visualization</p>
        <p style="color:#888; font-size:10px;">Drag to rotate • Scroll to zoom</p>
        <p id="update-count" style="color:#0f0;">Updates: 0</p>
    </div>
    <div id="stats">
        <div style="color:#0ff; margin-bottom:10px;">Action Probabilities</div>
        <div id="action-list"></div>
        <div style="margin-top:10px; color:#888;">
            Value: <span id="value">0.00</span>
        </div>
    </div>
    <canvas id="canvas"></canvas>

    <script type="importmap">
    {{
        "imports": {{
            "three": "https://cdn.jsdelivr.net/npm/three@0.160.0/build/three.module.js",
            "three/addons/": "https://cdn.jsdelivr.net/npm/three@0.160.0/examples/jsm/"
        }}
    }}
    </script>
    
    <script type="module">
        import * as THREE from 'three';
        import {{ OrbitControls }} from 'three/addons/controls/OrbitControls.js';
        import {{ GLTFLoader }} from 'three/addons/loaders/GLTFLoader.js';
        
        // Scene setup
        const scene = new THREE.Scene();
        scene.background = new THREE.Color(0x0a0a12);
        
        const camera = new THREE.PerspectiveCamera(60, window.innerWidth / window.innerHeight, 0.1, 1000);
        camera.position.set(5, 4, 5);
        
        const renderer = new THREE.WebGLRenderer({{ canvas: document.getElementById('canvas'), antialias: true }});
        renderer.setSize(window.innerWidth, window.innerHeight);
        renderer.setPixelRatio(window.devicePixelRatio);
        
        const controls = new OrbitControls(camera, renderer.domElement);
        controls.enableDamping = true;
        controls.dampingFactor = 0.05;
        
        // Lighting
        const ambient = new THREE.AmbientLight(0x404040, 2);
        scene.add(ambient);
        
        const directional = new THREE.DirectionalLight(0xffffff, 1);
        directional.position.set(5, 10, 5);
        scene.add(directional);
        
        // Grid
        const grid = new THREE.GridHelper(10, 20, 0x333333, 0x222222);
        scene.add(grid);
        
        // Origin marker
        const originGeo = new THREE.SphereGeometry(0.1);
        const originMat = new THREE.MeshBasicMaterial({{ color: 0x00ffff }});
        const origin = new THREE.Mesh(originGeo, originMat);
        scene.add(origin);
        
        // HOLD mesh container
        let holdGroup = new THREE.Group();
        scene.add(holdGroup);
        
        const loader = new GLTFLoader();
        let updateCount = 0;
        let lastJson = null;
        
        // Fetch and update mesh
        async function updateMesh() {{
            try {{
                // Get JSON for stats display
                const jsonRes = await fetch('/hold.json');
                if (jsonRes.ok) {{
                    lastJson = await jsonRes.json();
                    updateStats(lastJson);
                }}
                
                // Get GLB mesh
                const glbRes = await fetch('/hold.glb');
                if (glbRes.ok) {{
                    const buffer = await glbRes.arrayBuffer();
                    
                    loader.parse(buffer, '', (gltf) => {{
                        // Clear old
                        while (holdGroup.children.length) {{
                            holdGroup.remove(holdGroup.children[0]);
                        }}
                        
                        // Add new
                        holdGroup.add(gltf.scene);
                        
                        updateCount++;
                        document.getElementById('update-count').textContent = 'Updates: ' + updateCount;
                    }});
                }}
            }} catch (e) {{
                console.log('Update error:', e);
            }}
        }}
        
        function updateStats(data) {{
            const list = document.getElementById('action-list');
            const probs = data.action_probs || [];
            const names = data.action_names || probs.map((_, i) => 'action_' + i);
            const selected = data.selected_action || 0;
            
            let html = '';
            for (let i = 0; i < probs.length; i++) {{
                const pct = (probs[i] * 100).toFixed(1);
                const width = Math.max(2, probs[i] * 150);
                const isSelected = i === selected;
                html += '<div class="action' + (isSelected ? ' selected' : '') + '">';
                html += names[i] + ': ' + pct + '%';
                html += '<span class="bar" style="width:' + width + 'px"></span>';
                html += '</div>';
            }}
            list.innerHTML = html;
            
            document.getElementById('value').textContent = (data.value || 0).toFixed(4);
        }}
        
        // Animation loop
        function animate() {{
            requestAnimationFrame(animate);
            controls.update();
            
            // Gentle rotation of the hold visualization
            holdGroup.rotation.y += 0.002;
            
            renderer.render(scene, camera);
        }}
        
        // Start
        animate();
        updateMesh();
        setInterval(updateMesh, 200);  // Poll at 5Hz
        
        // Resize handling
        window.addEventListener('resize', () => {{
            camera.aspect = window.innerWidth / window.innerHeight;
            camera.updateProjectionMatrix();
            renderer.setSize(window.innerWidth, window.innerHeight);
        }});
    </script>
</body>
</html>"""
        return web.Response(text=html, content_type='text/html')
    
    async def _handle_glb(self, request):
        """Serve latest GLB."""
        if self._latest_glb:
            return web.Response(
                body=self._latest_glb,
                content_type='model/gltf-binary',
                headers={
                    'Access-Control-Allow-Origin': '*',
                    'X-Hold-Update': str(self._update_count),
                }
            )
        return web.Response(status=404, text="No HOLD data yet")
    
    async def _handle_json(self, request):
        """Serve latest JSON."""
        if self._latest_hold:
            return web.Response(
                text=json.dumps(self._latest_hold, indent=2),
                content_type='application/json',
                headers={'Access-Control-Allow-Origin': '*'}
            )
        return web.Response(status=404, text="No HOLD data yet")
    
    async def _handle_ws(self, request):
        """Handle WebSocket connection."""
        ws = web.WebSocketResponse()
        await ws.prepare(request)
        
        self._ws_clients.add(ws)
        print(f"🔌 WebSocket client connected ({len(self._ws_clients)} total)")
        
        try:
            # Send current state immediately
            if self._latest_glb:
                await ws.send_bytes(self._latest_glb)
            
            # Keep alive
            async for msg in ws:
                if msg.type == web.WSMsgType.TEXT:
                    if msg.data == 'ping':
                        await ws.send_str('pong')
                elif msg.type == web.WSMsgType.ERROR:
                    break
        finally:
            self._ws_clients.discard(ws)
            print(f"🔌 WebSocket client disconnected ({len(self._ws_clients)} total)")
        
        return ws
    
    async def _run_basic_http(self):
        """Fallback HTTP server without aiohttp."""
        from http.server import HTTPServer, BaseHTTPRequestHandler
        
        server_self = self
        
        class Handler(BaseHTTPRequestHandler):
            def do_GET(self):
                if self.path == '/hold.glb':
                    if server_self._latest_glb:
                        self.send_response(200)
                        self.send_header('Content-Type', 'model/gltf-binary')
                        self.send_header('Access-Control-Allow-Origin', '*')
                        self.end_headers()
                        self.wfile.write(server_self._latest_glb)
                    else:
                        self.send_error(404)
                elif self.path == '/hold.json':
                    if server_self._latest_hold:
                        self.send_response(200)
                        self.send_header('Content-Type', 'application/json')
                        self.send_header('Access-Control-Allow-Origin', '*')
                        self.end_headers()
                        self.wfile.write(json.dumps(server_self._latest_hold).encode())
                    else:
                        self.send_error(404)
                else:
                    self.send_response(200)
                    self.send_header('Content-Type', 'text/plain')
                    self.end_headers()
                    self.wfile.write(b"HOLD Mesh Server - GET /hold.glb or /hold.json")
            
            def log_message(self, format, *args):
                pass  # Suppress logs
        
        server = HTTPServer((self.host, self.port), Handler)
        server.timeout = 0.1
        
        while self._running:
            server.handle_request()


# Convenience function to run with an agent
def serve_hold_meshes(agent, port: int = 8766, poll_interval: float = 0.1):
    """
    Start mesh server that polls an agent's HOLD data.
    
    Args:
        agent: CapsuleAgent or similar with .brain.forward()
        port: Server port
        poll_interval: Seconds between polls
    """
    server = HoldMeshServer(port=port)
    server.start()
    
    print("Press Ctrl+C to stop")
    
    try:
        import numpy as np
        while True:
            # Generate observation (in real use, this comes from game)
            obs = np.random.randn(64 * 64 * 3).astype(np.float32)
            
            # Get HOLD data from agent
            output = agent.brain.forward({"obs": obs})
            
            hold_data = {
                "action_probs": output.get('action_probs', []),
                "action": output.get('action', 0),
                "value": output.get('value', 0.0),
                "latent": output.get('latent', []),
            }
            
            server.update(hold_data)
            time.sleep(poll_interval)
    except KeyboardInterrupt:
        print("\nStopping...")
        server.stop()


if __name__ == "__main__":
    # Demo server with fake data
    server = HoldMeshServer(port=8766)
    server.start()
    
    import numpy as np
    
    print("Sending demo HOLD updates...")
    print("Open http://localhost:8766 in browser")
    print("Press Ctrl+C to stop")
    
    try:
        step = 0
        while True:
            # Simulate changing decisions
            probs = np.random.dirichlet(np.ones(8) * 2)
            hold_data = {
                "action_probs": probs.tolist(),
                "action_names": ["forward", "left", "right", "back", "jump", "attack", "use", "wait"],
                "expected_values": (np.random.randn(8) * 0.5).tolist(),
                "model_card": {"name": "demo", "step": step},
            }
            server.update(hold_data)
            step += 1
            time.sleep(0.5)
    except KeyboardInterrupt:
        print("\nStopping...")
        server.stop()
