import asyncio
import websockets
import json
import sys
import time
import logging

# Enable debug logging
logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(sys.stdout)
    ]
)
logger = logging.getLogger(__name__)

class SocketIOTest:
    def __init__(self):
        self.websocket = None
        self.sid = None
        self.namespace = "/"
        self.ping_interval = 25  # seconds
        self.last_ping = 0
        self.connected = False
        self.ping_task = None
        self.heartbeat_task = None
        self.message_queue = asyncio.Queue()
        self.callbacks = {}
        
    async def _message_handler(self):
        """Handle incoming messages from the WebSocket."""
        try:
            while self.connected:
                try:
                    message = await asyncio.wait_for(
                        self.websocket.recv(),
                        timeout=1.0
                    )
                    logger.debug(f"<<< {message}")
                    
                    # Handle different message types
                    if message == '2':
                        # Ping from server, respond with pong
                        await self.websocket.send('3')  # pong
                        logger.debug(">>> 3 (pong)")
                    elif message.startswith('0'):
                        # Open packet
                        data = json.loads(message[1:])
                        self.sid = data.get('sid')
                        self.ping_interval = data.get('pingInterval', 25000) / 1000
                        logger.info(f"Connected with SID: {self.sid}, ping interval: {self.ping_interval}s")
                    elif message.startswith('40'):
                        # Namespace connection accepted
                        logger.info("Namespace connection accepted")
                        self.connected = True
                    elif message.startswith('42'):
                        # Event message
                        self._handle_event(message[2:])  # Remove '42' prefix
                    
                except asyncio.TimeoutError:
                    # Check if we need to send a ping
                    if time.time() - self.last_ping > self.ping_interval:
                        await self._send_ping()
        except Exception as e:
            logger.error(f"Message handler error: {e}", exc_info=True)
            self.connected = False
    
    def _handle_event(self, message):
        """Handle an incoming event message."""
        try:
            # Parse the message as JSON
            data = json.loads(message)
            
            # Check if it's a standard event (array with event name and data)
            if isinstance(data, list) and len(data) >= 1:
                event_name = data[0]
                event_data = data[1] if len(data) > 1 else None
                
                logger.debug(f"Received event: {event_name} - {event_data}")
                
                # Call any registered callbacks
                if event_name in self.callbacks:
                    for callback in self.callbacks[event_name]:
                        asyncio.create_task(callback(event_data))
                
                # Special handling for search results
                if event_name == 'search_results':
                    self._handle_search_results(event_data)
                    
        except json.JSONDecodeError:
            logger.error(f"Failed to parse event message: {message}")
        except Exception as e:
            logger.error(f"Error handling event: {e}", exc_info=True)
    
    def _handle_search_results(self, data):
        """Handle search results."""
        if not data or not isinstance(data, dict):
            logger.error("Invalid search results format")
            return
            
    async def test_search(self):
        """Test search functionality."""
        print("\n🔍 Sending search request for 'The Beatles'...")
        await self.emit('search', {
            'query': 'The Beatles',
            'type': 'track',
            'limit': 5
        })
        
        # Wait for search results
        print("⏳ Waiting for search results (timeout: 10 seconds)...")
        search_results = await self.wait_for_event('search_results', timeout=10)
        
        if search_results and 'results' in search_results and search_results['results'].get('track'):
            print("\n✅ Search Results Received:")
            print("-" * 50)
            tracks = search_results['results']['track']
            for i, track in enumerate(tracks[:3], 1):  # Show first 3 tracks
                print(f"{i}. {track['name']} - {track['artist']} ({track['album']})")
            if len(tracks) > 3:
                print(f"... and {len(tracks) - 3} more")
            return tracks[0]  # Return first track for download test
        else:
            print("\n❌ No search results received or invalid format")
            if search_results:
                print(json.dumps(search_results, indent=2))
            return None
            
    async def test_download(self, track_uri, track_name):
        """Test download functionality."""
        print(f"\n⬇️  Starting download for track: {track_name} ({track_uri})")
        
        # Listen for download events
        download_events = []
        
        def download_callback(data):
            event_type = data.get('type', 'unknown')
            print(f"\n📥 Download Event: {event_type}")
            if 'progress' in data:
                print(f"   Progress: {data['progress']}%")
            if 'status' in data:
                print(f"   Status: {data['status']}")
            if 'error' in data:
                print(f"   ❌ Error: {data['error']}")
            download_events.append(data)
        
        # Register callback for download events
        self.on('download_progress', download_callback)
        self.on('download_complete', download_callback)
        self.on('download_error', download_callback)
        
        # Start the download
        await self.emit('download', {
            'uri': track_uri,
            'type': 'track',
            'quality': 'high'
        })
        
        # Wait for download to complete or timeout
        print("⏳ Waiting for download to complete (timeout: 120 seconds)...")
        try:
            await asyncio.wait_for(
                self.wait_for_event('download_complete', matcher=lambda x: x.get('uri') == track_uri),
                timeout=120
            )
            print("\n✅ Download completed successfully!")
        except asyncio.TimeoutError:
            print("\n❌ Download timed out")
        
        # Clean up
        self.off('download_progress', download_callback)
        self.off('download_complete', download_callback)
        self.off('download_error', download_callback)
        
        return download_events
    
    async def wait_for_event(self, event_name, timeout=10, matcher=None):
        """Wait for a specific event with optional matcher function."""
        try:
            result_queue = asyncio.Queue()
            
            def event_handler(data):
                if matcher is None or matcher(data):
                    result_queue.put_nowait(data)
            
            self.on(event_name, event_handler)
            
            try:
                result = await asyncio.wait_for(result_queue.get(), timeout=timeout)
                return result
            except asyncio.TimeoutError:
                return None
            finally:
                self.off(event_name, event_handler)
                
        except Exception as e:
            logger.error(f"Error in wait_for_event: {e}", exc_info=True)
            return None
    
    async def _send_ping(self):
        """Send a ping to the server."""
        if not self.connected:
            return
            
        try:
            await self.websocket.send('2')  # ping
            logger.debug(">>> 2 (ping)")
            self.last_ping = time.time()
        except Exception as e:
            logger.error(f"Error sending ping: {e}")
            self.connected = False
    
    def on(self, event_name, callback):
        """Register an event callback."""
        if event_name not in self.callbacks:
            self.callbacks[event_name] = []
        self.callbacks[event_name].append(callback)
    
    def off(self, event_name, callback):
        """Unregister an event callback."""
        if event_name in self.callbacks:
            if callback in self.callbacks[event_name]:
                self.callbacks[event_name].remove(callback)
    
    async def emit(self, event_name, data=None):
        """Send an event to the server."""
        if not self.connected:
            logger.error("Not connected to WebSocket server")
            return False
            
        try:
            await self.websocket.send('42' + json.dumps([event_name, data]))
            logger.debug(f"Sent event: {event_name}")
            return True
        except Exception as e:
            logger.error(f"Error sending event {event_name}: {e}", exc_info=True)
            return False
    
    async def connect(self, uri):
        """Connect to the Socket.IO server."""
        logger.info(f"Connecting to {uri}")
        
        try:
            # Connect to the WebSocket
            self.websocket = await websockets.connect(
                uri,
                ping_interval=None,
                ping_timeout=None,
                close_timeout=1,
                max_queue=2**32
            )
            
            # Start the message handler
            self.connected = True
            asyncio.create_task(self._message_handler())
            
            # Wait for the initial handshake
            await asyncio.sleep(0.5)
            
            # Connect to the default namespace
            await self.websocket.send('40')
            logger.debug(">>> 40 (connect to namespace)")
            
            # Wait for connection to be established
            await asyncio.sleep(0.5)
            
            if not self.connected:
                logger.error("Failed to connect to namespace")
                return False
                
            logger.info("Successfully connected to Socket.IO server")
            return True
            
        except Exception as e:
            logger.error(f"Connection error: {e}", exc_info=True)
            self.connected = False
            return False
        
    async def send_ping(self):
        if not self.websocket:
            return
            
        current_time = time.time()
        if current_time - self.last_ping >= self.ping_interval:
            print("Sending ping")
            await self.websocket.send('2')
            self.last_ping = current_time
    
    async def send_connect(self):
        """Send a connect message to the namespace (for compatibility)."""
        return self.connected
    
    async def send_search(self, query, search_type='track', limit=5):
        """Send a search request to the server."""
        if not self.connected:
            logger.error("Not connected to WebSocket server")
            return False
            
        try:
            # Prepare the search message
            search_msg = {
                'query': query,
                'type': search_type,
                'limit': limit
            }
            
            # Format as Socket.IO message
            message = f'42["search", {json.dumps(search_msg)}]'
            logger.debug(f"Sending search: {message}")
            
            # Send the message
            await self.websocket.send(message)
            logger.debug("Search request sent")
            
            # Wait for results (handled by message handler)
            return True
            
        except Exception as e:
            logger.error(f"Error sending search: {e}", exc_info=True)
            return False
    
    async def close(self):
        """Close the WebSocket connection."""
        self.connected = False
        if self.websocket:
            try:
                await self.websocket.close()
                logger.info("WebSocket connection closed")
            except Exception as e:
                logger.error(f"Error closing WebSocket: {e}")
        
        # Cancel any running tasks
        if hasattr(self, 'ping_task') and self.ping_task:
            self.ping_task.cancel()
            try:
                await self.ping_task
            except asyncio.CancelledError:
                pass
                
        if hasattr(self, 'heartbeat_task') and self.heartbeat_task:
            self.heartbeat_task.cancel()
            try:
                await self.heartbeat_task
            except asyncio.CancelledError:
                pass

async def main():
    # Use the correct WebSocket URL for your Socket.IO server
    uri = "ws://localhost:6544/socket.io/?EIO=4&transport=websocket"
    
    client = SocketIOTest()
    
    async def test_search(self):
        """Test search functionality."""
        print("\n🔍 Sending search request for 'The Beatles'...")
        await self.emit('search', {
            'query': 'The Beatles',
            'type': 'track',
            'limit': 5
        })
        
        # Wait for search results
        print("⏳ Waiting for search results (timeout: 10 seconds)...")
        search_results = await self.wait_for_event('search_results', timeout=10)
        
        if search_results and 'results' in search_results and search_results['results'].get('track'):
            print("\n✅ Search Results Received:")
            print("-" * 50)
            tracks = search_results['results']['track']
            for i, track in enumerate(tracks[:3], 1):  # Show first 3 tracks
                print(f"{i}. {track['name']} - {track['artist']} ({track['album']})")
            if len(tracks) > 3:
                print(f"... and {len(tracks) - 3} more")
            return tracks[0]  # Return first track for download test
        else:
            print("\n❌ No search results received or invalid format")
            print(json.dumps(search_results, indent=2))
            return None
            
    async def test_download(self, track_uri, track_name):
        """Test download functionality."""
        print(f"\n⬇️  Starting download for track: {track_name} ({track_uri})")
        
        # Listen for download events
        download_events = []
        
        def download_callback(data):
            event_type = data.get('type', 'unknown')
            print(f"\n📥 Download Event: {event_type}")
            if 'progress' in data:
                print(f"   Progress: {data['progress']}%")
            if 'status' in data:
                print(f"   Status: {data['status']}")
            if 'error' in data:
                print(f"   ❌ Error: {data['error']}")
            download_events.append(data)
        
        # Register callback for download events
        self.on('download_progress', download_callback)
        self.on('download_complete', download_callback)
        self.on('download_error', download_callback)
        
        # Start the download
        await self.emit('download', {
            'uri': track_uri,
            'type': 'track',
            'quality': 'high'
        })
        
        # Wait for download to complete or timeout
        print("⏳ Waiting for download to complete (timeout: 120 seconds)...")
        try:
            await asyncio.wait_for(
                self.wait_for_event('download_complete', matcher=lambda x: x.get('uri') == track_uri),
                timeout=120
            )
            print("\n✅ Download completed successfully!")
        except asyncio.TimeoutError:
            print("\n❌ Download timed out")
        
        # Clean up
        self.off('download_progress', download_callback)
        self.off('download_complete', download_callback)
        self.off('download_error', download_callback)
        
        return download_events
        
async def main():
    client = SocketIOTest()
    
    try:
        # Connect to the WebSocket server
        print("Connecting to WebSocket server...")
        
        connected = await client.connect("ws://localhost:6544/socket.io/?EIO=4&transport=websocket")
        if not connected:
            print("❌ Failed to connect to WebSocket server")
            return
            
        print("✅ Connected to WebSocket server")
        
        # Test search
        print("\n=== TESTING SEARCH ===")
        track = await client.test_search()
        
        if track:
            # Test download
            print("\n=== TESTING DOWNLOAD ===")
            await client.test_download(track['uri'], track['name'])
        else:
            print("❌ No track found for download test")
            return
        await asyncio.sleep(10)
        
    except KeyboardInterrupt:
        print("\n🛑 Operation cancelled by user")
    except Exception as e:
        print(f"\n❌ Error: {e}")
        import traceback
        traceback.print_exc()
    finally:
        print("\n🔌 Closing connection...")
        await client.close()
        print("✅ Disconnected")

if __name__ == "__main__":
    print("🚀 Starting Socket.IO Test Client")
    print("-" * 50)
    
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\n👋 Goodbye!")
    except Exception as e:
        print(f"\n💥 Fatal error: {e}")
        import traceback
        traceback.print_exc()
    
    print("\n✨ Test completed")