server.py

#!/usr/bin/env python3
"""
p5.nvim asyncio-based live server with SSE console streaming.
Replaces blocking http.server with asyncio for better performance.
"""
import asyncio
import os
import re
import json
import signal
import sys
import webbrowser
from collections import deque
from datetime import datetime
from pathlib import Path
from typing import Optional
import websockets
import websockets.exceptions

# Configuration
CONFIG = {
    "port": int(sys.argv[1]) if len(sys.argv) > 1 else 8000,
    "live_reload": {
        "enabled": True,
        "port": 12002,
        "debounce_ms": 300,
        "watch_extensions": [".js", ".css", ".html", ".json"],
        "exclude_dirs": [".git", "node_modules", "dist", "build"],
    },
    "console": {
        "buffer_size": 1000,
        "heartbeat_interval": 15,
    }
}

# ANSI color codes for log formatting
ANSI_COLORS = {
    'reset': '\033[0m',
    'error': '\033[1;31m',
    'warn': '\033[1;33m',
    'info': '\033[1;36m',
    'log': '\033[0;37m',
    'timestamp': '\033[0;90m',
    'source': '\033[0;90m',
}

EMOJI_MAP = {
    'ERROR': '❌',
    'WARN': '⚠️ ',
    'INFO': 'ℹ️ ',
    'LOG': '📝',
    'HEARTBEAT': '💓',
}


def format_log_entry(level: str, message: str, source: str = "browser") -> str:
    """Format a log entry with ANSI colors and emojis for terminal display."""
    timestamp = datetime.now().strftime("%H:%M:%S")
    level = level.lower()
    
    emoji = EMOJI_MAP.get(level, '📝')
    level_color = ANSI_COLORS.get(level.lower(), ANSI_COLORS['log'])
    time_color = ANSI_COLORS['timestamp']
    source_color = ANSI_COLORS['source']
    reset = ANSI_COLORS['reset']
    
    return (
        f"{time_color}[{timestamp}]{reset} "
        f"{emoji} "
        f"{level_color}{level:5}{reset} "
        f"{source_color}[{source}]{reset}: {message}"
    )


class ConsoleBuffer:
    """Ring buffer for console logs with configurable size."""
    
    def __init__(self, max_size: int = 1000):
        self.buffer = deque(maxlen=max_size)
        self.max_size = max_size
    
    def append(self, entry: dict):
        """Add entry to buffer."""
        self.buffer.append(entry)
    
    def get_all(self) -> list:
        """Get all entries and clear buffer."""
        entries = list(self.buffer)
        self.buffer.clear()
        return entries
    
    def __len__(self):
        return len(self.buffer)


class LiveReloadServer:
    """WebSocket server for live reload using asyncio."""
    
    def __init__(self, port: int, directory: str, file_watcher):
        self.port = port
        self.directory = directory
        self.file_watcher = file_watcher
        self.clients = set()
        self.server = None
    
    async def handler(self, websocket):
        """Handle WebSocket client connection."""
        self.clients.add(websocket)
        try:
            await websocket.send(json.dumps({"type": "connected", "message": "Live reload connected"}))
            await websocket.wait_closed()
        except (websockets.exceptions.ConnectionClosedOK, websockets.exceptions.ConnectionClosedError):
            pass
        finally:
            self.clients.discard(websocket)
    
    async def start(self):
        """Start the WebSocket server."""
        try:
            self.server = await websockets.serve(self.handler, 'localhost', self.port)
            print(f"Live reload WebSocket running on ws://localhost:{self.port}")
        except OSError as e:
            for offset in range(1, 10):
                try:
                    alt_port = self.port + offset
                    self.server = await websockets.serve(self.handler, 'localhost', alt_port)
                    self.port = alt_port
                    print(f"Live reload WebSocket running on ws://localhost:{self.port}")
                    return
                except OSError:
                    continue
            print(f"Warning: Could not start live reload server: {e}")
    
    async def broadcast(self, message: dict):
        """Broadcast message to all connected clients."""
        data = json.dumps(message)
        
        # Take a snapshot to avoid concurrent modification during iteration
        clients_snapshot = set(self.clients)
        disconnected = set()
        
        for client in clients_snapshot:
            try:
                await client.send(data)
            except (websockets.exceptions.ConnectionClosedOK, websockets.exceptions.ConnectionClosedError, websockets.exceptions.InvalidState):
                disconnected.add(client)
            except Exception:
                disconnected.add(client)
        
        for client in disconnected:
            self.clients.discard(client)
            try:
                await client.close()
            except Exception:
                pass
    
    async def close(self):
        """Close the server and all connections."""
        for client in list(self.clients):
            try:
                await client.close()
            except Exception:
                pass
        self.clients.clear()
        if self.server:
            self.server.close()
            await self.server.wait_closed()


class FileWatcher:
    """Async file watcher using asyncio."""
    
    def __init__(self, directory: str, extensions: list, exclude_dirs: list, debounce_ms: int):
        self.directory = directory
        self.extensions = extensions
        self.exclude_dirs = exclude_dirs
        self.debounce_ms = debounce_ms / 1000
        self.last_trigger = 0
        self.running = False
        self._task: Optional[asyncio.Task] = None
    
    def should_watch(self, path: str) -> bool:
        """Check if file should be watched."""
        path_obj = Path(path)
        
        # Check exclusion dirs
        for part in path_obj.parts:
            if part in self.exclude_dirs:
                return False
        
        # Check extensions
        return any(str(path).endswith(ext) for ext in self.extensions)
    
    async def watch(self):
        """Watch for file changes - only trigger on actual file saves."""
        self.running = True
        last_mtimes = {}
        pending_changes = {}  # Track files that have changed but not yet stable
        
        while self.running:
            try:
                current_mtimes = {}
                for root, dirs, files in os.walk(self.directory):
                    dirs[:] = [d for d in dirs if d not in self.exclude_dirs]
                    
                    for file in files:
                        path = os.path.join(root, file)
                        if self.should_watch(path):
                            try:
                                current_mtimes[path] = os.path.getmtime(path)
                            except OSError:
                                continue
                
                now = datetime.now().timestamp()
                
                # Check for actual file changes (mtime differs from last known)
                for path, mtime in current_mtimes.items():
                    last_mtime = last_mtimes.get(path)
                    
                    if last_mtime is None:
                        # First time seeing this file - skip
                        continue
                    
                    if mtime != last_mtime:
                        # File has changed - mark as pending
                        if path not in pending_changes:
                            pending_changes[path] = now
                
                # Check pending changes for stability
                for path, change_time in list(pending_changes.items()):
                    current_mtime = current_mtimes.get(path)
                    last_mtime = last_mtimes.get(path)
                    
                    if current_mtime is None:
                        # File was deleted
                        del pending_changes[path]
                        continue
                    
                    if current_mtime == last_mtime and (now - change_time) >= self.debounce_ms:
                        # File is stable (not changing) and has been stable long enough
                        del pending_changes[path]
                        if now - self.last_trigger > self.debounce_ms:
                            self.last_trigger = now
                            yield path
                
                last_mtimes = current_mtimes
                await asyncio.sleep(0.3)
                
            except Exception as e:
                print(f"File watcher error: {e}")
                await asyncio.sleep(1)
    
    def start(self, callback):
        """Start the file watcher."""
        self._task = asyncio.create_task(self._run_watcher(callback))
    
    async def _run_watcher(self, callback):
        """Run the watcher loop."""
        async for path in self.watch():
            await callback(path)
    
    async def stop(self):
        """Stop the file watcher."""
        self.running = False
        if self._task:
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass


class HTTPServer:
    """Async HTTP server with SSE console streaming."""
    
    def __init__(self, port: int, directory: str, console_buffer: ConsoleBuffer, live_reload_server: LiveReloadServer):
        self.port = port
        self.directory = directory
        self.console_buffer = console_buffer
        self.live_reload_server = live_reload_server
        self.server: Optional[asyncio.Server] = None
        self.running = True
    
    async def handle_client(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
        """Handle incoming HTTP client request."""
        try:
            # Read request line
            request_line = await reader.readline()
            if not request_line:
                writer.close()
                await writer.wait_closed()
                return
            
            request_line = request_line.decode().strip()
            
            # Parse request
            parts = request_line.split()
            if len(parts) < 2:
                writer.close()
                await writer.wait_closed()
                return
            
            method = parts[0]
            path = parts[1]
            
            # Read headers
            headers = {}
            while True:
                line = await reader.readline()
                if not line or line == b'\r\n':
                    break
                header = line.decode().strip()
                if ':' in header:
                    key, value = header.split(':', 1)
                    headers[key.strip().lower()] = value.strip()
            
            # Route request
            if method == 'POST' and path == '/api/console/log':
                await self.handle_console_log(reader, writer, headers)
            elif method == 'GET' and path == '/api/console/stream':
                await self.handle_console_stream(reader, writer, headers)
            elif method == 'GET' and path == '/api/health':
                await self.handle_health(writer)
            else:
                await self.handle_static(method, path, reader, writer, headers)
                
        except Exception as e:
            print(f"Error handling client: {e}")
        finally:
            try:
                writer.close()
                await writer.wait_closed()
            except Exception:
                pass
    
    async def handle_console_log(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter, headers: dict):
        """Handle POST /api/console/log - receive logs from browser."""
        content_length = int(headers.get('content-length', 0))
        body = await reader.read(content_length)
        
        try:
            log_data = json.loads(body.decode('utf-8'))
            
            # Handle batch logs
            if log_data.get('type') == 'console_batch' and 'logs' in log_data:
                for entry in log_data['logs']:
                    if 'timestamp' not in entry:
                        entry['timestamp'] = datetime.now().isoformat()
                    self.console_buffer.append(entry)
            else:
                # Individual log
                if 'timestamp' not in log_data:
                    log_data['timestamp'] = datetime.now().isoformat()
                self.console_buffer.append(log_data)
            
            # Send response
            writer.write(b'HTTP/1.1 200 OK\r\n')
            writer.write(b'Content-Type: application/json\r\n')
            writer.write(b'Access-Control-Allow-Origin: *\r\n')
            writer.write(b'Connection: close\r\n')
            writer.write(b'\r\n')
            writer.write(json.dumps({"status": "received"}).encode())
            await writer.drain()
            
        except Exception as e:
            writer.write(b'HTTP/1.1 400 Bad Request\r\n')
            writer.write(b'Content-Type: application/json\r\n')
            writer.write(b'Access-Control-Allow-Origin: *\r\n')
            writer.write(b'Connection: close\r\n')
            writer.write(b'\r\n')
            writer.write(json.dumps({"error": "Invalid request"}).encode())
            await writer.drain()
    
    async def handle_console_stream(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter, headers: dict):
        """Handle GET /api/console/stream - SSE streaming endpoint."""
        # Send SSE headers
        writer.write(b'HTTP/1.1 200 OK\r\n')
        writer.write(b'Content-Type: text/plain; charset=utf-8\r\n')
        writer.write(b'Cache-Control: no-cache\r\n')
        writer.write(b'Access-Control-Allow-Origin: *\r\n')
        writer.write(b'Connection: keep-alive\r\n')
        writer.write(b'X-Accel-Buffering: no\r\n')
        writer.write(b'\r\n')
        await writer.drain()
        
        # Exponential backoff: 15s → 30s → 60s → 120s → max 300s
        base_interval = 15
        current_interval = base_interval
        max_interval = 300
        heartbeat_count = 0
        
        try:
            while self.running:
                # Get buffered logs
                logs = self.console_buffer.get_all()
                
                # Send buffered logs and reset exponential backoff
                if logs:
                    current_interval = base_interval  # Reset to 15s
                    heartbeat_count = 0
                    
                    for log_entry in logs:
                        level = log_entry.get('level', 'log')
                        message = log_entry.get('message', '')
                        source = log_entry.get('source', 'browser')
                        formatted = format_log_entry(level, message, source)
                        writer.write(f"data: {formatted}\n\n".encode())
                        await writer.drain()
                
                # Send silent keepalive (SSE comment) to prevent connection timeout
                heartbeat_count += 1
                if heartbeat_count >= current_interval:
                    heartbeat_count = 0
                    writer.write(b': heartbeat\n\n')
                    await writer.drain()
                    # Exponential backoff
                    current_interval = min(current_interval * 2, max_interval)
                
                # Wait before next check
                await asyncio.sleep(1)
                
        except (BrokenPipeError, ConnectionResetError, asyncio.CancelledError):
            pass
        finally:
            try:
                writer.close()
                await writer.wait_closed()
            except Exception:
                pass
    
    async def handle_health(self, writer: asyncio.StreamWriter):
        """Handle GET /api/health - health check endpoint."""
        writer.write(b'HTTP/1.1 200 OK\r\n')
        writer.write(b'Content-Type: application/json\r\n')
        writer.write(b'Access-Control-Allow-Origin: *\r\n')
        writer.write(b'Connection: close\r\n')
        writer.write(b'\r\n')
        writer.write(json.dumps({
            "status": "ok",
            "server": "p5.nvim asyncio",
            "console_buffer_size": len(self.console_buffer),
        }).encode())
        await writer.drain()
    
    async def handle_static(self, method: str, path: str, reader: asyncio.StreamReader, writer: asyncio.StreamWriter, headers: dict):
        """Handle static file serving."""
        if method != 'GET':
            writer.write(b'HTTP/1.1 405 Method Not Allowed\r\n')
            writer.write(b'Connection: close\r\n')
            writer.write(b'\r\n')
            await writer.drain()
            return
        
        # Handle root path
        if path == '/':
            path = '/index.html'
        
        # Prevent directory traversal - check resolved path stays within allowed directory
        try:
            resolved = Path(self.directory, path.lstrip('/')).resolve()
            base_resolved = Path(self.directory).resolve()
            if not resolved.is_relative_to(base_resolved):
                writer.write(b'HTTP/1.1 403 Forbidden\r\n')
                writer.write(b'Connection: close\r\n')
                writer.write(b'\r\n')
                await writer.drain()
                return
        except (ValueError, OSError):
            writer.write(b'HTTP/1.1 400 Bad Request\r\n')
            writer.write(b'Connection: close\r\n')
            writer.write(b'\r\n')
            await writer.drain()
            return
        
        # Build file path
        file_path = os.path.join(self.directory, path.lstrip('/'))
        
        # Debug: log index.html request
        if path == '/index.html':
            print(f"[DEBUG] index.html request: file_path={file_path}, exists={os.path.isfile(file_path)}")
        
        # Generate index.html on-the-fly if it doesn't exist
        if path == '/index.html' and not os.path.isfile(file_path):
            print("[DEBUG] Generating index.html on-the-fly...")
            content = self.generate_index_html()
            content = self.inject_scripts(content.encode('utf-8'))
            writer.write(b'HTTP/1.1 200 OK\r\n')
            writer.write(b'Content-Type: text/html\r\n')
            writer.write(f'Content-Length: {len(content)}\r\n'.encode())
            writer.write(b'Access-Control-Allow-Origin: *\r\n')
            writer.write(b'Connection: close\r\n')
            writer.write(b'\r\n')
            writer.write(content)
            await writer.drain()
            return
        
        if not os.path.isfile(file_path):
            writer.write(b'HTTP/1.1 404 Not Found\r\n')
            writer.write(b'Content-Type: text/plain\r\n')
            writer.write(b'Access-Control-Allow-Origin: *\r\n')
            writer.write(b'Connection: close\r\n')
            writer.write(b'\r\n')
            writer.write(b'File not found')
            await writer.drain()
            return
        
        # Determine content type
        ext = os.path.splitext(file_path)[1].lower()
        mime_types = {
            '.html': 'text/html',
            '.js': 'text/javascript',
            '.css': 'text/css',
            '.json': 'application/json',
            '.png': 'image/png',
            '.jpg': 'image/jpeg',
            '.gif': 'image/gif',
            '.svg': 'image/svg+xml',
            '.woff': 'application/font-woff',
            '.ttf': 'application/font-ttf',
        }
        content_type = mime_types.get(ext, 'application/octet-stream')
        
        # Read file
        try:
            with open(file_path, 'rb') as f:
                content = f.read()
            
            # Inject scripts for HTML files
            if ext == '.html':
                content = self.inject_scripts(content)
            
            # Send response
            writer.write(b'HTTP/1.1 200 OK\r\n')
            writer.write(f'Content-Type: {content_type}\r\n'.encode())
            writer.write(f'Content-Length: {len(content)}\r\n'.encode())
            writer.write(b'Access-Control-Allow-Origin: *\r\n')
            writer.write(b'Connection: close\r\n')
            writer.write(b'\r\n')
            writer.write(content)
            await writer.drain()
            
        except Exception as e:
            writer.write(b'HTTP/1.1 500 Internal Server Error\r\n')
            writer.write(b'Connection: close\r\n')
            writer.write(b'\r\n')
            writer.write(b'An error occurred while processing your request')
            await writer.drain()
    
    def generate_index_html(self) -> str:
        """Generate index.html on-the-fly based on p5.json libs."""
        import json
        
        p5_json_path = os.path.join(self.directory, 'p5.json')
        config = {}
        if os.path.isfile(p5_json_path):
            try:
                with open(p5_json_path, 'r') as f:
                    config = json.load(f)
            except:
                pass
        
        libs = config.get('libs', {})
        major = config.get('major', 2)
        version = config.get('version', '2.0.0' if major == 2 else '1.9.0')
        
        # Auto-create sketch.js if missing
        sketch_js_path = os.path.join(self.directory, 'sketch.js')
        if not os.path.isfile(sketch_js_path):
            default_sketch = '''function setup() {
  createCanvas(400, 400);
}

function draw() {
  background(220);
  circle(mouseX, mouseY, 50);
}'''
            with open(sketch_js_path, 'w') as f:
                f.write(default_sketch)
            print(f"Created default sketch.js")
        
        # Build script tags for core and contrib libs
        scripts = []
        scripts.append(f'  <script src="https://cdnjs.cloudflare.com/ajax/libs/p5.js/{version}/p5.min.js"></script>')
        scripts.append(f'  <script src="https://cdnjs.cloudflare.com/ajax/libs/p5.js/{version}/addons/p5.sound.min.js"></script>')
        
        for lib_name in libs.keys():
            lib_version = libs[lib_name]
            scripts.append(f'  <script src="assets/libs/{lib_name}.js"></script>')
        
        scripts.append('  <script src="assets/libs/libs.js"></script>')
        
        html = f'''<!DOCTYPE html>
<html lang="en">
<head>
  <meta charset="UTF-8">
  <meta name="viewport" content="width=device-width, initial-scale=1.0">
  <title>p5.js Sketch</title>
  <link rel="icon" type="image/x-icon" href="assets/favicon.ico">
{chr(10).join(scripts)}
</head>
<body>
  <main>
  </main>
  <script src="sketch.js"></script>
</body>
</html>'''
        return html
    
    def inject_scripts(self, html_content: bytes) -> bytes:
        """Inject console and live reload scripts into HTML."""
        content = html_content.decode('utf-8', errors='ignore')
        
        # Console injection script
        console_script = '''
    <script>
      (function() {
        console.log('p5.nvim console integration enabled');
        
        const originalConsole = {
          log: console.log,
          error: console.error,
          warn: console.warn,
          info: console.info
        };
        
        function sendToConsole(level, args) {
          const message = args.map(arg => {
            if (typeof arg === 'object') {
              try {
                return JSON.stringify(arg);
              } catch (e) {
                return String(arg);
              }
            }
            return String(arg);
          }).join(' ');
          
          fetch('/api/console/log', {
            method: 'POST',
            headers: { 'Content-Type': 'application/json' },
            body: JSON.stringify({
              type: 'console',
              level: level,
              message: message,
              source: 'browser',
              timestamp: new Date().toISOString()
            })
          }).catch(() => {});
        }
        
        console.log = function(...args) {
          originalConsole.log.apply(console, args);
          sendToConsole('log', args);
        };
        
        console.error = function(...args) {
          originalConsole.error.apply(console, args);
          sendToConsole('error', args);
        };
        
        console.warn = function(...args) {
          originalConsole.warn.apply(console, args);
          sendToConsole('warn', args);
        };
        
        console.info = function(...args) {
          originalConsole.info.apply(console, args);
          sendToConsole('info', args);
        };
        
        window.onerror = function(msg, source, lineno, colno, error) {
          sendToConsole('error', [msg + ' at ' + source + ':' + lineno + ':' + colno]);
          return false;
        };
      })();
    </script>'''
        
        # Live reload script
        live_reload_script = f'''
    <script>
      (function() {{
        let ws = null;
        let reconnectAttempts = 0;
        const maxReconnectAttempts = 10;
        
        function connect() {{
          ws = new WebSocket('ws://localhost:{self.live_reload_server.port}');
          
          ws.onopen = function() {{
            console.log('Live reload connected');
            reconnectAttempts = 0;
          }};
          
          ws.onclose = function() {{
            // Only attempt reconnect, don't reload
            if (reconnectAttempts < maxReconnectAttempts) {{
              reconnectAttempts++;
              setTimeout(connect, Math.min(1000 * reconnectAttempts, 5000));
            }}
          }};
          
          ws.onerror = function() {{
            // Silently fail, onclose will handle reconnect
          }};
          
          ws.onmessage = function(event) {{
            try {{
              const data = JSON.parse(event.data);
              if (data.type === 'reload') {{
                window.location.reload();
              }}
            }} catch (e) {{}}
          }};
        }}
        
        connect();
      }})();
    </script>'''
        
        # Inject scripts before </body>
        if '</body>' in content.lower():
            content = re.sub(r'</body>', console_script + live_reload_script + '</body>', content, flags=re.IGNORECASE)
        else:
            content += console_script + live_reload_script
        
        return content.encode('utf-8')
    
    async def start(self):
        """Start the HTTP server."""
        try:
            self.server = await asyncio.start_server(
                self.handle_client, 'localhost', self.port
            )
            print(f"Server running at http://localhost:{self.port}/")
        except OSError as e:
            print(f"Error starting server on port {self.port}: {e}")
            # Try alternate ports
            for offset in range(1, 10):
                try:
                    alt_port = self.port + offset
                    self.server = await asyncio.start_server(
                        self.handle_client, 'localhost', alt_port
                    )
                    self.port = alt_port
                    print(f"Server running at http://localhost:{self.port}/")
                    return
                except OSError:
                    continue
            raise
    
    async def close(self):
        """Close the HTTP server."""
        self.running = False
        if self.server:
            self.server.close()
            await self.server.wait_closed()


async def main():
    """Main entry point."""
    directory = os.getcwd()
    
    # Create components
    console_buffer = ConsoleBuffer(max_size=CONFIG['console']['buffer_size'])
    
    lr_config = CONFIG['live_reload']
    file_watcher = FileWatcher(
        directory=directory,
        extensions=lr_config['watch_extensions'],
        exclude_dirs=lr_config['exclude_dirs'],
        debounce_ms=lr_config['debounce_ms']
    )
    
    live_reload_server = LiveReloadServer(
        port=lr_config['port'],
        directory=directory,
        file_watcher=file_watcher
    )
    
    http_server = HTTPServer(
        port=CONFIG['port'],
        directory=directory,
        console_buffer=console_buffer,
        live_reload_server=live_reload_server
    )
    
    # Start servers
    await live_reload_server.start()
    await http_server.start()
    
    # Update live reload port in HTTP server if it changed
    lr_config['port'] = live_reload_server.port
    
    # Start file watcher
    async def on_file_change(path: str):
        message = {
            "type": "reload",
            "file": path,
            "timestamp": datetime.now().isoformat()
        }
        try:
            await live_reload_server.broadcast(message)
            print(f"Reload triggered for: {path}")
        except Exception:
            pass
    
    file_watcher.start(on_file_change)
    
    # Handle shutdown
    shutdown_event = asyncio.Event()
    
    def signal_handler():
        print("\nShutting down server...")
        shutdown_event.set()
    
    loop = asyncio.get_event_loop()
    for sig in (signal.SIGTERM, signal.SIGINT):
        try:
            loop.add_signal_handler(sig, signal_handler)
        except NotImplementedError:
            # Windows doesn't support add_signal_handler
            pass
    
    # Wait for shutdown
    await shutdown_event.wait()
    
    # Cleanup
    print("Closing connections...")
    await file_watcher.stop()
    await live_reload_server.close()
    await http_server.close()
    print("Server stopped")


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\nServer stopped by user")