import asyncio
import websockets
import json
import uuid
from quart import Quart, jsonify, request # Import necessary Quart components

# --- WebSocket Server Components ---
# Dictionary to store active clients: {client_id: websocket}
active_clients = {}

async def register_client(websocket, client_id):
    """Registers a new client connection."""
    active_clients[client_id] = websocket
    print(f"Client {client_id} connected.")

async def unregister_client(client_id):
    """Unregisters a disconnected client."""
    if client_id in active_clients:
        del active_clients[client_id]
        print(f"Client {client_id} disconnected.")

async def send_command_to_client(client_id, command_name, *args):
    """Sends a command to a specific client."""
    if client_id in active_clients:
        websocket = active_clients[client_id]
        message = json.dumps({"type": "command", "name": command_name, "args": args})
        try:
            await websocket.send(message)
            print(f"Sent command '{command_name}' to client {client_id}")
        except websockets.exceptions.ConnectionClosedError:
            print(f"Failed to send to client {client_id}: connection closed.")
            await unregister_client(client_id)
    else:
        print(f"Client {client_id} not found.")

async def send_command_to_all_clients(command_name, *args):
    """Sends a command to all connected clients."""
    message = json.dumps({"type": "command", "name": command_name, "args": args})
    # Use a list of items to avoid issues if clients disconnect during iteration
    clients_to_send = list(active_clients.items())
    for client_id, websocket in clients_to_send:
        try:
            await websocket.send(message)
            print(f"Sent command '{command_name}' to client {client_id}")
        except websockets.exceptions.ConnectionClosedError:
            print(f"Failed to send to client {client_id}: connection closed.")
            await unregister_client(client_id)


async def websocket_server_handler(websocket):
    """Handles incoming WebSocket connections and messages from clients."""
    client_id = None
    try:
        # Handshake: Receive the first message which should contain the client ID
        handshake_message = await websocket.recv()
        handshake_data = json.loads(handshake_message)

        if handshake_data.get("type") == "handshake" and "id" in handshake_data:
            client_id = handshake_data["id"]
            await register_client(websocket, client_id)
            await websocket.send(json.dumps({"type": "handshake_ack", "status": "success"})) # Acknowledge handshake

            # Keep the connection alive and listen for potential messages from the client
            # (Though in this server-commanded model, clients might not send much)
            # We primarily wait for the client to close the connection
            await websocket.wait_closed()

        else:
            print(f"Received invalid handshake from {websocket.remote_address}. Closing connection.")
            await websocket.close()

    except websockets.exceptions.ConnectionClosedError:
        print(f"Client connection closed unexpectedly for {client_id}.")
    except json.JSONDecodeError:
        print(f"Received invalid JSON from {client_id or 'an unknown client'}.")
    except Exception as e:
        print(f"An error occurred with client {client_id}: {e}")
    finally:
        if client_id:
            await unregister_client(client_id)

# --- Radio Channel Data (Placeholder) ---
# In a real app, this would likely come from a database
channels = {
    "channel_1": {
        "id": "channel_1",
        "name": "Local News Radio",
        "frequency_list_khz": [98500],
        "decode_mode": "p25",
        "location": "Cityville",
        "tags": None,
        "tag_whitelist": [1,2,3],
        "avail_on_nodes": ["client-abc123", "client-def456"], # List of client IDs that can tune this channel
        "description": "Your source for local news and weather."
    },
    "channel_2": {
        "id": "channel_2",
        "name": "Music Mix FM",
        "frequency_list_khz": [101300],
        "decode_mode": "p25",
        "location": "Townsville",
        "tags": None,
        "tag_whitelist": [6,7,8],
        "avail_on_nodes": ["client-def456", "client-ghi789"],
        "description": "Playing the hits, all day long."
    }
    # Add more channels here
}

# --- Quart API Components ---
app = Quart(__name__)

# Store the websocket server instance
websocket_server_instance = None

@app.before_serving
async def startup_websocket_server():
    """Starts the WebSocket server when the Quart app starts."""
    global websocket_server_instance
    websocket_server_address = "localhost"
    websocket_server_port = 8765

    # Start the WebSocket server task
    websocket_server_instance = await websockets.serve(
        websocket_server_handler,
        websocket_server_address,
        websocket_server_port
    )
    print(f"WebSocket server started on ws://{websocket_server_address}:{websocket_server_port}")

@app.after_serving
async def shutdown_websocket_server():
    """Shuts down the WebSocket server when the Quart app stops."""
    global websocket_server_instance
    if websocket_server_instance:
        websocket_server_instance.close()
        await websocket_server_instance.wait_closed()
        print("WebSocket server shut down.")


@app.route('/')
async def index():
    return "Welcome to the Radio App Server API!"

@app.route('/channels', methods=['GET'])
async def get_channels():
    """API endpoint to get a list of all channels."""
    # Return a list of channel IDs and names for a summary view
    channel_summary = [{"id": ch["id"], "name": ch["name"]} for ch in channels.values()]
    return jsonify(channel_summary)

@app.route('/channels/<channel_id>', methods=['GET'])
async def get_channel_details(channel_id):
    """API endpoint to get details for a specific channel."""
    channel = channels.get(channel_id)
    if channel:
        return jsonify(channel)
    else:
        return jsonify({"error": "Channel not found"}), 404

@app.route('/clients', methods=['GET'])
async def list_clients():
    """API endpoint to list currently connected client IDs."""
    return jsonify(list(active_clients.keys()))

# --- Example API endpoint to trigger a WebSocket command ---
# This is a simple example. More complex logic might be needed
# to validate commands or arguments before sending.
@app.route('/command/<client_id>/<command_name>', methods=['POST'])
async def api_send_command(client_id, command_name):
    """
    API endpoint to send a command to a specific client via WebSocket.
    Expects JSON body with 'args': [...]
    e.g., POST to /command/client-abc123/print_message with body {"args": ["Hello!"]}
    """
    if client_id not in active_clients:
        return jsonify({"error": f"Client {client_id} not found"}), 404

    if command_name not in ["print_message", "set_status", "run_task"]: # Basic validation
         return jsonify({"error": f"Unknown command: {command_name}"}), 400

    try:
        request_data = await request.get_json()
        args = request_data.get("args", [])
        if not isinstance(args, list):
             return jsonify({"error": "'args' must be a list"}), 400

        # Send the command asynchronously
        await send_command_to_client(client_id, command_name, *args)

        return jsonify({"status": "command sent", "client_id": client_id, "command": command_name}), 200

    except Exception as e:
        return jsonify({"error": f"Failed to send command: {e}"}), 500

# --- Main Execution ---
if __name__ == "__main__":
    # Quart's app.run() will start the asyncio event loop and manage it.
    # The @app.before_serving decorator ensures the websocket server starts within that loop.
    # We removed asyncio.run(main()) and the main() function itself.
    print("Starting Quart API server...")
    app.run(
        host="localhost",
        port=5000,
        debug=False # Set to True for development
    )
    print("Quart API server stopped.")