import asyncio
import websockets
import json
import numpy as np
import traceback # Import the traceback module
from music_generator import MusicGenerator

# --- WebSocket Server Setup ---

# In-memory storage for connected clients
clients = {
    "webapp": set()
}
audio_source = None

# --- Main Application Logic ---

def initialize_dependencies():
    """Loads all necessary files and initializes objects."""
    global notes, generator
    print("Initializing dependencies...")
    try:
        # It's better to load files relative to the script's location
        import os
        dir_path = os.path.dirname(os.path.realpath(__file__))
        
        with open(os.path.join(dir_path, 'consonance_matrix.json')) as f:
            consonance_matrix = np.array(json.load(f))

        notes = ['C', 'C#', 'D', 'D#', 'E', 'F', 'F#', 'G', 'G#', 'A', 'A#', 'B']
        generator = MusicGenerator(len(notes))
        print("Dependencies initialized successfully.")
        return True
    except Exception as e:
        print("--- CRITICAL: FAILED TO INITIALIZE DEPENDENCIES ---")
        print(traceback.format_exc())
        print("----------------------------------------------------")
        return False

async def broadcast_to_webapps(message):
    """Sends a message to all connected webapp clients."""
    if clients["webapp"]:
        tasks = [client.send(message) for client in clients["webapp"]]
        await asyncio.gather(*tasks, return_exceptions=True)

async def handle_audio_data(data):
    """Mocks the audio analysis."""
    import random
    detected_chord = random.choice(notes)
    predicted_chord = random.choice(notes)
    key = "C Major"

    analysis_result = {
        "type": "analysis_update",
        "current_chord": detected_chord,
        "predicted_chord": predicted_chord,
        "musical_key": key
    }
    await broadcast_to_webapps(json.dumps(analysis_result))

# --- WebSocket Connection Management ---

async def connection_handler(websocket, path):
    """Handles incoming WebSocket connections with robust error logging."""
    global audio_source
    print(f"New client connected: {websocket.remote_address}")

    try:
        initial_message = await websocket.recv()
        message_data = json.loads(initial_message)
        client_type = message_data.get("type")

        if client_type == "extension_hello":
            if audio_source is not None:
                await audio_source.close(code=1012, reason="New extension connected.")
            audio_source = websocket
            print("Audio capture extension connected.")
            await websocket.send(json.dumps({"status": "connected", "role": "audio_source"}))
            await broadcast_to_webapps(json.dumps({"type": "status_update", "message": "Audio source connected."}))

        elif client_type == "webapp_hello":
            clients["webapp"].add(websocket)
            print("Web app client connected.")
            await websocket.send(json.dumps({"status": "connected", "role": "viewer"}))
            status_msg = "Audio source connected." if audio_source else "Waiting for audio source..."
            await websocket.send(json.dumps({"type": "status_update", "message": status_msg}))
        
        else:
            print(f"Unknown client type: {client_type}. Disconnecting.")
            return

        async for message in websocket:
            if websocket == audio_source:
                await handle_audio_data(message)

    except websockets.exceptions.ConnectionClosed as e:
        print(f"Connection closed normally for {websocket.remote_address}. Code: {e.code}, Reason: {e.reason}")
    except Exception as e:
        # THIS IS THE CRITICAL PART
        print(f"--- UNEXPECTED SERVER ERROR in connection_handler ---")
        print(f"Error Type: {type(e).__name__}")
        print(f"Error Message: {e}")
        print("Traceback:")
        print(traceback.format_exc())
        print("----------------------------------------------------")
        # Close the connection with a specific error code if it's not already closed
        if not websocket.closed:
            await websocket.close(code=1011, reason="Internal Server Error")
    finally:
        if websocket in clients["webapp"]:
            clients["webapp"].remove(websocket)
        if websocket == audio_source:
            audio_source = None
            print("Audio capture extension disconnected.")
            await broadcast_to_webapps(json.dumps({"type": "status_update", "message": "Audio source disconnected."}))

async def main():
    """Initializes dependencies and starts the WebSocket server."""
    if not initialize_dependencies():
        print("Server cannot start due to initialization failure.")
        return

    websocket_port = 7860
    print(f"Starting WebSocket server on port {websocket_port}...")
    async with websockets.serve(connection_handler, "0.0.0.0", websocket_port):
        await asyncio.Future()  # run forever

if __name__ == "__main__":
    asyncio.run(main())