#!/usr/bin/env python3
import sys
import json
import socket
import asyncio
import base64
import uuid
import os
import argparse
from io import BytesIO
from PIL import Image, ImageDraw, ImageFont, ImageColor

# Third-party dependencies
import tornado.web
import tornado.websocket
import tornado.ioloop

# Global stdout reference for MCP protocol responses
mcp_stdout = None

# Helper to parse color values
def parse_color(color_val, default=(255, 255, 255, 255)):
    if isinstance(color_val, (int, float)):
        # Handle Waveshare legacy screen logic: 0 = White, 1 = Black
        if color_val == 0:
            return (255, 255, 255, 255)
        else:
            return (0, 0, 0, 255)
            
    if not color_val:
        return default
        
    try:
        rgb = ImageColor.getrgb(str(color_val))
        if len(rgb) == 3:
            return rgb + (255,)
        return rgb
    except Exception:
        return default

# Class to manage the state of the virtual companion
class VirtualDeviceState:
    def __init__(self):
        # 800x600 virtual screen buffer
        self.canvas = Image.new("RGBA", (800, 600), (15, 23, 42, 255)) # slate-900 initial color
        self.draw = ImageDraw.Draw(self.canvas)
        
        # LED Ring parameters
        self.led_r = 0
        self.led_g = 0
        self.led_b = 0
        self.led_mode = "off"
        
        # Virtual sensors telemetry
        self.sensor_temp = 22.5
        self.sensor_hum = 45.0
        self.sensor_light = 350.0
        
        # Pending microphone recording futures
        self.active_recordings = {}

    def clear_screen(self, color):
        fill_color = parse_color(color, (15, 23, 42, 255))
        self.draw.rectangle([0, 0, 800, 600], fill=fill_color)

    def draw_text(self, text, x, y, size=24, color="white"):
        fill_color = parse_color(color, (255, 255, 255, 255))
        # Load system fonts safely
        font = None
        for font_name in ["Arial.ttf", "Arial.ttc", "Helvetica.ttf", "Helvetica.ttc", "Courier.ttf"]:
            try:
                # Check typical paths or load by name
                font = ImageFont.truetype(font_name, size)
                break
            except IOError:
                continue
                
        if font is None:
            # Try macOS system fonts location
            macos_font_path = f"/System/Library/Fonts/Supplemental/Arial.ttf"
            if os.path.exists(macos_font_path):
                try:
                    font = ImageFont.truetype(macos_font_path, size)
                except IOError:
                    pass
                    
        if font is None:
            font = ImageFont.load_default()
            
        self.draw.text((x, y), text, fill=fill_color, font=font)

    def draw_shape(self, shape, x, y, w, h, color="white", fill=False):
        draw_color = parse_color(color, (255, 255, 255, 255))
        
        if shape == "rect":
            if fill:
                self.draw.rectangle([x, y, x + w, y + h], fill=draw_color)
            else:
                self.draw.rectangle([x, y, x + w, y + h], outline=draw_color, width=2)
        elif shape == "circle":
            if fill:
                self.draw.ellipse([x, y, x + w, y + h], fill=draw_color)
            else:
                self.draw.ellipse([x, y, x + w, y + h], outline=draw_color, width=2)
        elif shape == "line":
            self.draw.line([x, y, x + w, y + h], fill=draw_color, width=2)

    def draw_image(self, image_base64, x=0, y=0, w=None, h=None):
        if "," in image_base64:
            image_base64 = image_base64.split(",")[1]
            
        img_data = base64.b64decode(image_base64)
        img = Image.open(BytesIO(img_data))
        
        # Handle scaling compatibility
        try:
            resample_method = Image.Resampling.LANCZOS
        except AttributeError:
            resample_method = Image.ANTIALIAS
            
        if w and h:
            img = img.resize((w, h), resample_method)
        elif w or h:
            orig_w, orig_h = img.size
            if w:
                ratio = w / orig_w
                img = img.resize((w, int(orig_h * ratio)), resample_method)
            else:
                ratio = h / orig_h
                img = img.resize((int(orig_w * ratio), h), resample_method)
                
        # Paste image on canvas
        self.canvas.paste(img, (x, y), img if img.mode in ('RGBA', 'LA') else None)

    def get_screenshot_base64(self):
        buffered = BytesIO()
        self.canvas.save(buffered, format="PNG")
        return base64.b64encode(buffered.getvalue()).decode("utf-8")

# Instantiate device state
state = VirtualDeviceState()

# Tornado Web Application Route Handlers
class MainHandler(tornado.web.RequestHandler):
    def get(self):
        # Serve the index.html from the same directory
        self.render("index.html")

class APIHandler(tornado.web.RequestHandler):
    def set_default_headers(self):
        self.set_header("Access-Control-Allow-Origin", "*")
        self.set_header("Access-Control-Allow-Headers", "x-requested-with, content-type")
        self.set_header("Access-Control-Allow-Methods", "POST, GET, OPTIONS")

    def options(self):
        self.set_status(204)
        self.finish()

    async def post(self):
        try:
            req = json.loads(self.request.body.decode('utf-8'))
            method = req.get("method")
            rpc_id = req.get("id")
            
            if method == "tools/call":
                params = req.get("params", {})
                tool_name = params.get("name")
                arguments = params.get("arguments", {})
                
                result = await execute_mcp_tool(tool_name, arguments)
                if isinstance(result, list):
                    content = result
                else:
                    content = [{"type": "text", "text": str(result)}]
                    
                resp = {
                    "jsonrpc": "2.0",
                    "id": rpc_id,
                    "result": {
                        "content": content
                    }
                }
            elif method == "initialize":
                resp = {
                    "jsonrpc": "2.0",
                    "id": rpc_id,
                    "result": {
                        "protocolVersion": "2024-11-05",
                        "capabilities": {
                            "tools": {}
                        },
                        "serverInfo": {
                            "name": "virtual-screen-mcp",
                            "version": "1.0.0"
                        }
                    }
                }
            else:
                resp = {
                    "jsonrpc": "2.0",
                    "id": rpc_id,
                    "error": {
                        "code": -32601,
                        "message": f"Method {method} not found"
                    }
                }
            self.write(json.dumps(resp))
        except Exception as e:
            resp = {
                "jsonrpc": "2.0",
                "id": None,
                "error": {
                    "code": -32000,
                    "message": str(e)
                }
            }
            self.write(json.dumps(resp))

class ClientWebSocketHandler(tornado.websocket.WebSocketHandler):
    clients = set()

    def check_origin(self, origin):
        return True  # Enable local network browser connections

    def open(self):
        ClientWebSocketHandler.clients.add(self)
        sys.stderr.write(f"[Web Server] Client connected from: {self.request.remote_ip}\n")
        sys.stderr.flush()
        
        # Send current screen content and active LED state to client on join
        self.send_initial_state()

    def send_initial_state(self):
        try:
            b64_png = state.get_screenshot_base64()
            # Draw the current canvas image to the canvas of the newly joined client
            self.write_message(json.dumps({
                "action": "draw_image",
                "image_base64": b64_png,
                "x": 0,
                "y": 0,
                "w": 800,
                "h": 600
            }))
            # Send current LED settings
            self.write_message(json.dumps({
                "action": "set_led",
                "r": state.led_r,
                "g": state.led_g,
                "b": state.led_b,
                "mode": state.led_mode
            }))
        except Exception as e:
            sys.stderr.write(f"[Web Server] Error syncing state: {e}\n")
            sys.stderr.flush()

    def on_close(self):
        ClientWebSocketHandler.clients.remove(self)
        sys.stderr.write("[Web Server] Client disconnected\n")
        sys.stderr.flush()

    def on_message(self, message):
        try:
            data = json.loads(message)
            event = data.get("event")

            if event == "sensor_update":
                state.sensor_temp = data.get("temperature", 22.5)
                state.sensor_hum = data.get("humidity", 45.0)
                state.sensor_light = data.get("light", 350.0)
                
            elif event == "voice_recording":
                rec_id = data.get("recording_id")
                if rec_id in state.active_recordings:
                    fut = state.active_recordings[rec_id]
                    if not fut.done():
                        error_msg = data.get("error")
                        if error_msg:
                            fut.set_exception(RuntimeError(error_msg))
                        else:
                            fut.set_result(data.get("audio_base64"))
                            
            elif event == "ping":
                # Respond to browser ping with pong latency calculation
                self.write_message(json.dumps({"event": "pong", "t": data.get("t")}))
                
        except Exception as e:
            sys.stderr.write(f"[Web Server] Error processing message: {e}\n")
            sys.stderr.flush()

    @classmethod
    def broadcast(cls, payload):
        msg = json.dumps(payload)
        for client in cls.clients:
            try:
                client.write_message(msg)
            except Exception:
                pass

# Send responses back to the LLM Client via stdio
def send_rpc_response(resp):
    payload = json.dumps(resp) + "\n"
    mcp_stdout.write(payload)
    mcp_stdout.flush()

# Get local LAN IP address
def get_local_ip():
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    try:
        s.connect(('10.255.255.255', 1))
        IP = s.getsockname()[0]
    except Exception:
        IP = '127.0.0.1'
    finally:
        s.close()
    return IP

# Define the list of MCP tools
def get_mcp_tools_list():
    return [
        {
            "name": "clear_screen",
            "description": "Clear the virtual screen. Optionally specify color (e.g. standard hex '#1e1e2e', name 'white', or 0 for white / 1 for black).",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "color": {"type": "string", "description": "Color hex, CSS name, or integer (0=white, 1=black)"}
                }
            }
        },
        {
            "name": "draw_text",
            "description": "Draw text at specified (x,y) coordinates on the virtual display.",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "text": {"type": "string", "description": "The text message to draw"},
                    "x": {"type": "integer", "description": "X coordinate (0-780)"},
                    "y": {"type": "integer", "description": "Y coordinate (0-580)"},
                    "size": {"type": "integer", "description": "Font size in pixels (default 24)"},
                    "color": {"type": "string", "description": "Text color (default 'white')"}
                },
                "required": ["text", "x", "y"]
            }
        },
        {
            "name": "draw_shape",
            "description": "Draw a geometric shape (rectangle, circle, or line) on the virtual display.",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "shape": {"type": "string", "enum": ["rect", "circle", "line"], "description": "Type of shape"},
                    "x": {"type": "integer", "description": "Start X coordinate"},
                    "y": {"type": "integer", "description": "Start Y coordinate"},
                    "width": {"type": "integer", "description": "Width (or delta X for line)"},
                    "height": {"type": "integer", "description": "Height (or delta Y for line)"},
                    "color": {"type": "string", "description": "Color of the outline or fill (default 'white')"},
                    "fill": {"type": "boolean", "description": "Whether to fill the shape (default false)"}
                },
                "required": ["shape", "x", "y", "width", "height"]
            }
        },
        {
            "name": "draw_image",
            "description": "Draw a base64 encoded PNG or JPEG image on the virtual screen at (x,y).",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "image_base64": {"type": "string", "description": "Base64 encoded string of the image file"},
                    "x": {"type": "integer", "description": "X coordinate (default 0)", "default": 0},
                    "y": {"type": "integer", "description": "Y coordinate (default 0)", "default": 0},
                    "width": {"type": "integer", "description": "Width to scale the image (optional)"},
                    "height": {"type": "integer", "description": "Height to scale the image (optional)"}
                },
                "required": ["image_base64"]
            }
        },
        {
            "name": "get_screenshot",
            "description": "Capture the current virtual display buffer and return it as a PNG image, allowing you to see what is currently rendered on the screen.",
            "inputSchema": {"type": "object", "properties": {}}
        },
        {
            "name": "set_led",
            "description": "Control the virtual RGB LED ring lights.",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "r": {"type": "integer", "minimum": 0, "maximum": 255, "description": "Red channel (0-255)"},
                    "g": {"type": "integer", "minimum": 0, "maximum": 255, "description": "Green channel (0-255)"},
                    "b": {"type": "integer", "minimum": 0, "maximum": 255, "description": "Blue channel (0-255)"},
                    "mode": {"type": "string", "enum": ["static", "breath", "rainbow", "off"], "description": "LED animation mode"}
                },
                "required": ["r", "g", "b", "mode"]
            }
        },
        {
            "name": "get_sensors",
            "description": "Read telemetry values from virtual temperature, humidity, and ambient light sensors.",
            "inputSchema": {"type": "object", "properties": {}}
        },
        {
            "name": "play_tone",
            "description": "Play a sound tone of a specific frequency and duration on the virtual speaker.",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "frequency": {"type": "integer", "description": "Frequency in Hz (default 440)"},
                    "duration_ms": {"type": "integer", "description": "Duration in milliseconds (default 1000)"},
                    "volume": {"type": "integer", "minimum": 0, "maximum": 100, "description": "Volume percent from 0 to 100 (default 50)"}
                }
            }
        },
        {
            "name": "play_audio",
            "description": "Play a base64 encoded audio track (like WAV or MP3) or an audio URL on the virtual speaker.",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "audio_base64_or_url": {"type": "string", "description": "Base64 audio bytes or accessible audio file URL"},
                    "volume": {"type": "integer", "minimum": 0, "maximum": 100, "description": "Volume percent from 0 to 100 (default 50)"}
                },
                "required": ["audio_base64_or_url"]
            }
        },
        {
            "name": "record_voice",
            "description": "Record voice commands/input from the connected client's microphone for a specified duration.",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "duration_sec": {"type": "integer", "minimum": 1, "maximum": 20, "description": "Duration to record in seconds (default 4)", "default": 4}
                }
            }
        }
    ]

# Execute tool operations
async def execute_mcp_tool(name, args):
    if name == "clear_screen":
        color = args.get("color", 0)
        state.clear_screen(color)
        ClientWebSocketHandler.broadcast({"action": "clear", "color": color})
        return "Screen cleared successfully."
        
    elif name == "draw_text":
        text = str(args.get("text", ""))
        x = int(args.get("x", 10))
        y = int(args.get("y", 10))
        size = int(args.get("size", 24))
        color = str(args.get("color", "white"))
        
        state.draw_text(text, x, y, size, color)
        ClientWebSocketHandler.broadcast({
            "action": "draw_text",
            "text": text,
            "x": x,
            "y": y,
            "size": size,
            "color": color
        })
        return f"Successfully drew text at ({x}, {y})."
        
    elif name == "draw_shape":
        shape = str(args.get("shape"))
        x = int(args.get("x"))
        y = int(args.get("y"))
        w = int(args.get("width"))
        h = int(args.get("height"))
        color = str(args.get("color", "white"))
        fill = bool(args.get("fill", False))
        
        state.draw_shape(shape, x, y, w, h, color, fill)
        ClientWebSocketHandler.broadcast({
            "action": "draw_shape",
            "shape": shape,
            "x": x,
            "y": y,
            "w": w,
            "h": h,
            "color": color,
            "fill": fill
        })
        return f"Successfully drew {shape} shape."
        
    elif name == "draw_image":
        img_b64 = str(args.get("image_base64"))
        x = int(args.get("x", 0))
        y = int(args.get("y", 0))
        w = args.get("width")
        h = args.get("height")
        
        if w is not None: w = int(w)
        if h is not None: h = int(h)
        
        state.draw_image(img_b64, x, y, w, h)
        ClientWebSocketHandler.broadcast({
            "action": "draw_image",
            "image_base64": img_b64,
            "x": x,
            "y": y,
            "w": w,
            "h": h
        })
        return "Successfully rendered image on display."
        
    elif name == "get_screenshot":
        b64_png = state.get_screenshot_base64()
        return [
            {
                "type": "image",
                "data": b64_png,
                "mimeType": "image/png"
            }
        ]
        
    elif name == "set_led":
        r = int(args.get("r", 0))
        g = int(args.get("g", 0))
        b = int(args.get("b", 0))
        mode = str(args.get("mode", "static"))
        
        state.led_r, state.led_g, state.led_b, state.led_mode = r, g, b, mode
        ClientWebSocketHandler.broadcast({
            "action": "set_led",
            "r": r,
            "g": g,
            "b": b,
            "mode": mode
        })
        return f"NeoPixel LED ring set to '{mode}' mode with color ({r}, {g}, {b})."
        
    elif name == "get_sensors":
        return json.dumps({
            "temperature_c": state.sensor_temp,
            "humidity_pct": state.sensor_hum,
            "light_lux": state.sensor_light
        })
        
    elif name == "play_tone":
        freq = int(args.get("frequency", 440))
        dur = int(args.get("duration_ms", 1000))
        vol = int(args.get("volume", 50))
        
        ClientWebSocketHandler.broadcast({
            "action": "play_tone",
            "frequency": freq,
            "duration": dur,
            "volume": vol / 100.0
        })
        return f"Played synthesized tone at {freq}Hz for {dur}ms."
        
    elif name == "play_audio":
        audio_payload = str(args.get("audio_base64_or_url"))
        vol = int(args.get("volume", 50))
        
        ClientWebSocketHandler.broadcast({
            "action": "play_audio",
            "audio": audio_payload,
            "volume": vol / 100.0
        })
        return "Successfully transmitted audio buffer."
        
    elif name == "record_voice":
        dur = int(args.get("duration_sec", 4))
        
        if not ClientWebSocketHandler.clients:
            return "No active browser sessions to capture microphone input. Recording aborted."
            
        rec_id = str(uuid.uuid4())
        loop = asyncio.get_running_loop()
        fut = loop.create_future()
        state.active_recordings[rec_id] = fut
        
        ClientWebSocketHandler.broadcast({
            "action": "record_voice",
            "duration_sec": dur,
            "recording_id": rec_id
        })
        
        try:
            # Wait for client recording with a safety buffer
            audio_b64 = await asyncio.wait_for(fut, timeout=dur + 5.0)
            audio_bytes = base64.b64decode(audio_b64)
            
            # Save the WAV file to local disk
            filename = f"recording_{rec_id[:8]}.wav"
            with open(filename, "wb") as f:
                f.write(audio_bytes)
                
            return f"Successfully recorded {dur} seconds of audio. File saved to '{filename}'."
        except asyncio.TimeoutError:
            return "Microphone recording timed out. Make sure the connected web client has granted mic permissions."
        except Exception as e:
            return f"Microphone recording failed: {str(e)}"
        finally:
            state.active_recordings.pop(rec_id, None)
            
    else:
        raise ValueError(f"Unknown MCP tool: {name}")

# Handles parsed MCP stdin request objects
async def handle_mcp_request(req):
    method = req.get("method")
    rpc_id = req.get("id")
    
    is_notification = rpc_id is None
    
    if method == "initialize":
        resp = {
            "jsonrpc": "2.0",
            "id": rpc_id,
            "result": {
                "protocolVersion": "2024-11-05",
                "capabilities": {
                    "tools": {}
                },
                "serverInfo": {
                    "name": "virtual-screen-mcp",
                    "version": "1.0.0"
                }
            }
        }
        send_rpc_response(resp)
        return
        
    elif method == "tools/list":
        resp = {
            "jsonrpc": "2.0",
            "id": rpc_id,
            "result": {
                "tools": get_mcp_tools_list()
            }
        }
        send_rpc_response(resp)
        return
        
    elif method == "tools/call":
        params = req.get("params", {})
        tool_name = params.get("name")
        arguments = params.get("arguments", {})
        
        try:
            result = await execute_mcp_tool(tool_name, arguments)
            if isinstance(result, list):
                content = result
            else:
                content = [{"type": "text", "text": str(result)}]
                
            resp = {
                "jsonrpc": "2.0",
                "id": rpc_id,
                "result": {
                    "content": content
                }
            }
        except Exception as e:
            resp = {
                "jsonrpc": "2.0",
                "id": rpc_id,
                "error": {
                    "code": -32000,
                    "message": str(e)
                }
            }
        send_rpc_response(resp)
        return
        
    if not is_notification:
        resp = {
            "jsonrpc": "2.0",
            "id": rpc_id,
            "error": {
                "code": -32601,
                "message": f"Method {method} not found"
            }
        }
        send_rpc_response(resp)

# Read stdin line-by-line asynchronously
async def stdio_reader_loop():
    loop = asyncio.get_running_loop()
    reader = asyncio.StreamReader()
    protocol = asyncio.StreamReaderProtocol(reader)
    await loop.connect_read_pipe(lambda: protocol, sys.stdin)
    
    while True:
        line = await reader.readline()
        if not line:
            break
            
        try:
            req = json.loads(line.decode('utf-8').strip())
            await handle_mcp_request(req)
        except Exception as e:
            sys.stderr.write(f"[MCP Stdio] Error handling line input: {e}\n")
            sys.stderr.flush()

# Strong references to prevent garbage collection of background tasks
background_tasks = set()

# Async main runtime loop
async def main_async(port):
    app = tornado.web.Application([
        (r"/", MainHandler),
        (r"/api/mcp", APIHandler),
        (r"/ws", ClientWebSocketHandler),
    ], template_path=os.path.dirname(os.path.abspath(__file__)))
    
    app.listen(port)
    
    local_ip = get_local_ip()
    sys.stderr.write(f"--------------------------------------------------\n")
    sys.stderr.write(f"Virtual Screen & Speaker MCP Server Initialized.\n")
    sys.stderr.write(f"Connect local devices in your browser to:\n")
    sys.stderr.write(f"==> http://{local_ip}:{port}/\n")
    sys.stderr.write(f"==> http://localhost:{port}/\n")
    sys.stderr.write(f"--------------------------------------------------\n")
    sys.stderr.flush()
    
    # Spawn the stdin line reader task and store a strong reference
    task = asyncio.create_task(stdio_reader_loop())
    background_tasks.add(task)
    task.add_done_callback(background_tasks.discard)
    
    # Run indefinitely
    while True:
        await asyncio.sleep(3600)

def main():
    parser = argparse.ArgumentParser(description="Virtual Screen & Speakers MCP Server")
    parser.add_argument("--port", type=int, default=8080, help="Web port (default 8080)")
    args = parser.parse_args()
    
    # Set the MCP stdio channel
    global mcp_stdout
    mcp_stdout = sys.stdout
    # Route default stdout to stderr so prints don't interrupt standard JSON-RPC communications
    sys.stdout = sys.stderr
    
    try:
        asyncio.run(main_async(args.port))
    except KeyboardInterrupt:
        sys.stderr.write("[Server] Stopped by KeyboardInterrupt.\n")
        sys.stderr.flush()

if __name__ == "__main__":
    main()