#!/usr/bin/env python3
"""
Comprehensive real-time testing tool for SamCloud Music.
Tests search, downloads, and WebSocket events with live updates.
"""
import os
import sys
import time
import json
import logging
import threading
import requests
import websocket
from typing import Dict, List, Optional, Any, Callable
from urllib.parse import urljoin, urlparse
from dataclasses import dataclass, asdict
from enum import Enum
from queue import Queue, Empty

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler('test_app.log')
    ]
)
logger = logging.getLogger(__name__)

class EventType(Enum):
    SEARCH = "search"
    DOWNLOAD = "download"
    PROGRESS = "progress"
    COMPLETE = "complete"
    ERROR = "error"
    CONNECT = "connect"
    DISCONNECT = "disconnect"

@dataclass
class TestEvent:
    event_type: EventType
    data: Dict[str, Any]
    timestamp: float = None
    
    def __post_init__(self):
        if self.timestamp is None:
            self.timestamp = time.time()
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            'event_type': self.event_type.value,
            'data': self.data,
            'timestamp': self.timestamp
        }

class SamCloudTester:
    def __init__(self, base_url: str = "http://localhost:6543"):
        self.base_url = base_url.rstrip('/')
        self.ws_url = self._get_websocket_url()
        self.ws = None
        self.ws_thread = None
        self.running = False
        self.event_queue = Queue()
        self.active_downloads: Dict[str, Dict] = {}
        self.last_search_results: List[Dict] = []
        self.event_handlers = {
            EventType.SEARCH: [],
            EventType.DOWNLOAD: [],
            EventType.PROGRESS: [],
            EventType.COMPLETE: [],
            EventType.ERROR: [],
            EventType.CONNECT: [],
            EventType.DISCONNECT: []
        }
    
    def _get_websocket_url(self) -> str:
        """Convert HTTP URL to WebSocket URL."""
        parsed = urlparse(self.base_url)
        scheme = 'wss' if parsed.scheme == 'https' else 'ws'
        return f"{scheme}://{parsed.netloc}/socket.io/"
    
    def on_event(self, event_type: EventType) -> Callable:
        """Decorator to register event handlers."""
        def decorator(handler):
            self.event_handlers[event_type].append(handler)
            return handler
        return decorator
    
    def _emit_event(self, event: TestEvent):
        """Emit an event to all registered handlers."""
        for handler in self.event_handlers[event.event_type]:
            try:
                handler(event)
            except Exception as e:
                logger.error(f"Error in event handler: {e}")
    
    def connect(self) -> bool:
        """Connect to the WebSocket server."""
        try:
            # First, get the WebSocket URL from the server
            response = requests.get(f"{self.base_url}/socket.io/?EIO=4&transport=polling")
            sid = json.loads(response.text[1:])['sid']
            
            # Connect to the WebSocket
            self.ws = websocket.WebSocketApp(
                f"{self.ws_url}?EIO=4&transport=websocket&sid={sid}",
                on_open=self._on_ws_open,
                on_message=self._on_ws_message,
                on_error=self._on_ws_error,
                on_close=self._on_ws_close,
                header={"Origin": self.base_url}
            )
            
            # Start WebSocket in a separate thread
            self.running = True
            self.ws_thread = threading.Thread(target=self._ws_run_forever)
            self.ws_thread.daemon = True
            self.ws_thread.start()
            
            # Wait for connection to be established
            max_wait = 5  # seconds
            start_time = time.time()
            while not hasattr(self, 'ws_connected') and (time.time() - start_time) < max_wait:
                time.sleep(0.1)
            
            if not hasattr(self, 'ws_connected'):
                logger.error("Failed to establish WebSocket connection")
                return False
                
            logger.info("WebSocket connection established")
            self._emit_event(TestEvent(EventType.CONNECT, {}))
            return True
            
        except Exception as e:
            logger.error(f"Connection failed: {e}")
            return False
    
    def _ws_run_forever(self):
        """Run WebSocket client in a loop."""
        while self.running:
            try:
                self.ws.run_forever(ping_interval=25, ping_timeout=10)
            except Exception as e:
                logger.error(f"WebSocket error: {e}")
                time.sleep(1)
    
    def _on_ws_open(self, ws):
        """Handle WebSocket connection open."""
        logger.debug("WebSocket connection opened")
        self.ws_connected = True
    
    def _on_ws_message(self, ws, message):
        """Handle incoming WebSocket messages."""
        try:
            if not message:
                return
                
            # Parse Socket.IO message
            if message.startswith('4'):  # Message type 4 is a Socket.IO message
                try:
                    # Extract the JSON part after the message type
                    json_str = message[1:]
                    if json_str.startswith('['):
                        json_str = json_str[1:-1]  # Remove brackets
                    
                    event = json.loads(json_str)
                    event_name = event.get('name')
                    data = event.get('args', [{}])[0]
                    
                    # Emit event based on message type
                    if event_name == 'download_progress':
                        download_id = data.get('id')
                        if download_id:
                            self._update_download(download_id, {
                                'status': 'downloading',
                                'progress': data.get('progress', 0),
                                'message': data.get('message', 'Downloading...')
                            })
                            self._emit_event(TestEvent(
                                EventType.PROGRESS,
                                {'download_id': download_id, **data}
                            ))
                            
                    elif event_name == 'download_complete':
                        download_id = data.get('id')
                        if download_id:
                            success = data.get('success', False)
                            self._update_download(download_id, {
                                'status': 'completed' if success else 'failed',
                                'progress': 100 if success else self.active_downloads.get(download_id, {}).get('progress', 0),
                                'message': 'Download completed' if success else 'Download failed'
                            })
                            self._emit_event(TestEvent(
                                EventType.COMPLETE if success else EventType.ERROR,
                                {'download_id': download_id, **data}
                            ))
                            
                except json.JSONDecodeError as e:
                    logger.warning(f"Failed to parse WebSocket message: {e}")
                except Exception as e:
                    logger.error(f"Error processing WebSocket message: {e}")
    
    def _on_ws_error(self, ws, error):
        """Handle WebSocket errors."""
        logger.error(f"WebSocket error: {error}")
        self._emit_event(TestEvent(EventType.ERROR, {'message': str(error)}))
    
    def _on_ws_close(self, ws, close_status_code, close_msg):
        """Handle WebSocket connection close."""
        logger.info("WebSocket connection closed")
        self.ws_connected = False
        self._emit_event(TestEvent(EventType.DISCONNECT, {
            'code': close_status_code,
            'message': close_msg or 'No message'
        }))
    
    def disconnect(self):
        """Disconnect from the WebSocket server."""
        self.running = False
        if hasattr(self, 'ws') and self.ws:
            self.ws.close()
        if self.ws_thread and self.ws_thread.is_alive():
            self.ws_thread.join(timeout=2)
    
    def search(self, query: str, search_type: str = 'track', limit: int = 10) -> List[Dict]:
        """Search for tracks, albums, or artists."""
        try:
            response = requests.get(
                f"{self.base_url}/api/search",
                params={
                    'q': query,
                    'type': search_type,
                    'limit': limit
                },
                timeout=10
            )
            response.raise_for_status()
            
            results = response.json()
            self.last_search_results = results.get('tracks', {}).get('items', [])
            
            event = TestEvent(
                EventType.SEARCH,
                {
                    'query': query,
                    'type': search_type,
                    'results_count': len(self.last_search_results)
                }
            )
            self._emit_event(event)
            
            return self.last_search_results
            
        except Exception as e:
            error_msg = f"Search failed: {str(e)}"
            logger.error(error_msg)
            self._emit_event(TestEvent(EventType.ERROR, {'message': error_msg}))
            return []
    
    def download_track(self, track_url: str) -> Optional[str]:
        """Start downloading a track."""
        try:
            response = requests.post(
                f"{self.base_url}/api/download",
                json={"url": track_url},
                timeout=10
            )
            response.raise_for_status()
            
            data = response.json()
            download_id = data.get('id')
            if download_id:
                self.active_downloads[download_id] = {
                    'status': 'queued',
                    'progress': 0,
                    'message': 'Queued for download',
                    'url': track_url,
                    'start_time': time.time()
                }
                
                event = TestEvent(
                    EventType.DOWNLOAD,
                    {
                        'download_id': download_id,
                        'url': track_url,
                        'status': 'queued'
                    }
                )
                self._emit_event(event)
                
                return download_id
                
        except Exception as e:
            error_msg = f"Download failed to start: {str(e)}"
            logger.error(error_msg)
            self._emit_event(TestEvent(EventType.ERROR, {'message': error_msg}))
            return None
    
    def _update_download(self, download_id: str, updates: Dict[str, Any]):
        """Update download status."""
        if download_id in self.active_downloads:
            self.active_downloads[download_id].update(updates)
    
    def get_download_status(self, download_id: str) -> Optional[Dict]:
        """Get the status of a download."""
        return self.active_downloads.get(download_id)
    
    def get_active_downloads(self) -> Dict[str, Dict]:
        """Get all active downloads."""
        return self.active_downloads

class ConsoleUI:
    """Simple console-based UI for testing."""
    
    def __init__(self, tester: SamCloudTester):
        self.tester = tester
        self.running = True
        self.setup_event_handlers()
    
    def setup_event_handlers(self):
        """Set up event handlers for the tester."""
        @self.tester.on_event(EventType.SEARCH)
        def on_search(event):
            print(f"\n🔍 Search results ({event.data['results_count']} items)")
            
        @self.tester.on_event(EventType.DOWNLOAD)
        def on_download(event):
            print(f"\n⬇️  Download started: {event.data['url']}")
            print(f"   ID: {event.data['download_id']}")
            
        @self.tester.on_event(EventType.PROGRESS)
        def on_progress(event):
            download = self.tester.get_download_status(event.data['download_id'])
            if download:
                print(f"\r📥 {download['message']} - {download['progress']}%", end='', flush=True)
                
        @self.tester.on_event(EventType.COMPLETE)
        def on_complete(event):
            download = self.tester.get_download_status(event.data['download_id'])
            if download:
                duration = time.time() - download['start_time']
                print(f"\n✅ Download completed in {duration:.1f}s")
                print(f"   Saved to: {download.get('path', 'unknown')}")
                
        @self.tester.on_event(EventType.ERROR)
        def on_error(event):
            print(f"\n❌ Error: {event.data.get('message', 'Unknown error')}")
    
    def run(self):
        """Run the console UI."""
        print("=== SamCloud Music Tester ===")
        print("Type 'search <query>' to search for tracks")
        print("Type 'download <url>' to download a track")
        print("Type 'status' to view active downloads")
        print("Type 'exit' to quit\n")
        
        # Connect to the WebSocket
        if not self.tester.connect():
            print("Failed to connect to the server. Is it running?")
            return
        
        try:
            while self.running:
                try:
                    command = input("\n> ").strip().split(' ', 1)
                    cmd = command[0].lower()
                    args = command[1] if len(command) > 1 else ''
                    
                    if cmd == 'search' and args:
                        print(f"Searching for: {args}")
                        results = self.tester.search(args)
                        if results:
                            print("\nSearch Results:")
                            for i, item in enumerate(results[:10], 1):
                                print(f"{i}. {item['name']} - {item['artists'][0]['name']} ({item['album']['name']})")
                        else:
                            print("No results found")
                            
                    elif cmd == 'download' and args:
                        print(f"Starting download: {args}")
                        download_id = self.tester.download_track(args)
                        if not download_id:
                            print("Failed to start download")
                            
                    elif cmd == 'status':
                        downloads = self.tester.get_active_downloads()
                        if downloads:
                            print("\nActive Downloads:")
                            for d_id, download in downloads.items():
                                print(f"- {d_id}: {download['status']} - {download['progress']}% - {download['url']}")
                        else:
                            print("No active downloads")
                            
                    elif cmd in ('exit', 'quit'):
                        print("Goodbye!")
                        self.running = False
                        
                    else:
                        print("Unknown command. Try 'search', 'download', 'status', or 'exit'")
                        
                except KeyboardInterrupt:
                    print("\nUse 'exit' to quit")
                except Exception as e:
                    print(f"Error: {e}")
                    
        finally:
            self.tester.disconnect()

if __name__ == "__main__":
    # Default to localhost:6543 or use command line argument
    base_url = sys.argv[1] if len(sys.argv) > 1 else "http://localhost:6543"
    
    tester = SamCloudTester(base_url)
    ui = ConsoleUI(tester)
    ui.run()