# osc_daemon.py
import asyncio
from pythonosc.udp_client import SimpleUDPClient
from pythonosc.osc_server import AsyncIOOSCUDPServer
from pythonosc.dispatcher import Dispatcher
import json
from typing import Optional, Dict, Any
import argparse
import sys

class AbletonOSCDaemon:
    def __init__(self, 
                 socket_host='127.0.0.1', socket_port=65432,
                 ableton_host='127.0.0.1', ableton_port=11000,
                 receive_port=11001):
        self.socket_host = socket_host
        self.socket_port = socket_port
        self.ableton_host = ableton_host
        self.ableton_port = ableton_port
        self.receive_port = receive_port
        
        # Initialize OSC client for Ableton
        self.osc_client = SimpleUDPClient(ableton_host, ableton_port)
        
        # Store active connections waiting for responses
        self.pending_responses: Dict[str, asyncio.Future] = {}
        
        # Initialize OSC server dispatcher
        self.dispatcher = Dispatcher()
        self.dispatcher.set_default_handler(self.handle_ableton_message)
        
    def handle_ableton_message(self, address: str, *args):
        """Handle incoming OSC messages from Ableton."""
        print(f"[ABLETON MESSAGE] Address: {address}, Args: {args}")
        
        # If this address has a pending response, resolve it
        if address in self.pending_responses:
            future = self.pending_responses[address]
            if not future.done():
                future.set_result({
                    'status': 'success',
                    'address': address,
                    'data': args
                })
            del self.pending_responses[address]
            
    async def start(self):
        """Start both the socket server and OSC server."""
        try:
            # Start OSC server to receive Ableton messages
            self.osc_server = AsyncIOOSCUDPServer(
                (self.socket_host, self.receive_port),
                self.dispatcher,
                asyncio.get_event_loop()
            )
            
            transport, protocol = await self.osc_server.create_serve_endpoint()
            print(f"OSC server listening for Ableton Live messages on {self.socket_host}:{self.receive_port}")
            
            # Start socket server for MCP communication
            server = await asyncio.start_server(
                self.handle_socket_client,
                self.socket_host,
                self.socket_port
            )
            print(f"Ableton OSC Daemon socket server running on {self.socket_host}:{self.socket_port}")
            print(f"Sending OSC messages to Ableton Live on {self.ableton_host}:{self.ableton_port}")
            print("Waiting for MCP server connection...")
            print("Press Ctrl+C to exit")
            
            async with server:
                await server.serve_forever()
        except Exception as e:
            print(f"ERROR starting the OSC daemon: {e}")
            print("Check if the ports are already in use or if Ableton Live is running.")
            raise

    async def handle_socket_client(self, reader, writer):
        """Handle incoming socket connections from MCP server."""
        client_address = writer.get_extra_info('peername')
        print(f"[NEW CONNECTION] Client connected from {client_address}")
        
        try:
            while True:
                data = await reader.read(1024)
                if not data:
                    break
                    
                try:
                    message = json.loads(data.decode())
                    print(f"[RECEIVED MESSAGE] From {client_address}: {message}")
                    
                    command = message.get('command')

                    
                    if command == 'send_message':
                        # Extract OSC message details
                        address = message.get('address')
                        args = message.get('args', [])
                        
                        # For commands that expect responses, set up a future
                        if address.startswith(('/live/device/get', '/live/scene/get', '/live/view/get', '/live/clip/get', '/live/clip_slot/get', '/live/track/get', '/live/song/get', '/live/api/get', '/live/application/get', '/live/test', '/live/error')):
                            # Create response future with timeout
                            future = asyncio.Future()
                            self.pending_responses[address] = future

                            # Send to Ableton
                            self.osc_client.send_message(address, args)

                            try:
                                # Wait for response with timeout
                                response = await asyncio.wait_for(future, timeout=5.0)
                                print(f"[OSC RESPONSE] Received: {response}")
                                writer.write(json.dumps(response).encode())
                            except asyncio.TimeoutError:
                                response = {
                                    'status': 'error',
                                    'message': f'Timeout waiting for response to {address}'
                                }
                                print(f"[OSC TIMEOUT] {response}")
                                writer.write(json.dumps(response).encode())
                                
                        else:
                            # For commands that don't expect responses
                            self.osc_client.send_message(address, args)
                            response = {'status': 'sent'}
                            writer.write(json.dumps(response).encode())
                            
                    elif command == 'get_status':
                        response = {
                            'status': 'ok',
                            'ableton_port': self.ableton_port,
                            'receive_port': self.receive_port
                        }
                        print(f"[STATUS REQUEST] Responding with: {response}")
                        writer.write(json.dumps(response).encode())
                    else:
                        response = {'status': 'error', 'message': 'Unknown command'}
                        print(f"[UNKNOWN COMMAND] Received: {message}")
                        writer.write(json.dumps(response).encode())
                    
                    await writer.drain()
                    
                except json.JSONDecodeError:
                    print(f"[JSON ERROR] Could not decode message: {data}")
                    response = {'status': 'error', 'message': 'Invalid JSON'}
                    writer.write(json.dumps(response).encode())
                    
        except Exception as e:
            print(f"[CONNECTION ERROR] Error handling client: {e}")
        finally:
            writer.close()
            await writer.wait_closed()
            print(f"[CONNECTION CLOSED] Client {client_address} disconnected")

def main():
    """Entry point for the OSC daemon when run as a command-line application."""
    parser = argparse.ArgumentParser(description="Ableton OSC Daemon")
    parser.add_argument("--socket-host", default="127.0.0.1", help="Host for the socket server")
    parser.add_argument("--socket-port", type=int, default=65432, help="Port for the socket server")
    parser.add_argument("--ableton-host", default="127.0.0.1", help="Host for Ableton Live")
    parser.add_argument("--ableton-port", type=int, default=11000, help="Port for sending to Ableton Live")
    parser.add_argument("--receive-port", type=int, default=11001, help="Port for receiving from Ableton Live")
    
    args = parser.parse_args()
    
    daemon = AbletonOSCDaemon(
        socket_host=args.socket_host,
        socket_port=args.socket_port,
        ableton_host=args.ableton_host,
        ableton_port=args.ableton_port,
        receive_port=args.receive_port
    )
    
    print(f"Starting Ableton OSC Daemon...")
    print(f"NOTE: You also need to start the AbletonOSC control surface in Ableton Live.")
    print(f"Make sure AbletonOSC is installed and enabled in Ableton Live's preferences.")
    
    try:
        asyncio.run(daemon.start())
    except KeyboardInterrupt:
        print("\nShutting down OSC daemon...")
    except Exception as e:
        print(f"Error: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()
