# server.py
import asyncio
import websockets
import json
import uuid
from quart import Quart, jsonify, request
from routers.systems import systems_bp
from routers.nodes import nodes_bp, register_client, unregister_client
from routers.bot import bot_bp
from routers.auth import auth_bp # ONLY import auth_bp, not jwt instance from auth.py
from internal.db_wrappers import SystemDbController, DiscordIdDbController
from internal.auth_wrappers import UserDbController
# Import the JWTManager instance and its configuration function
from config.jwt_config import jwt, configure_jwt # Import the actual jwt instance and the config function


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


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)

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

# Store the websocket server instance
websocket_server_instance = None

# Make active_clients accessible via the app instance.
app.active_clients = active_clients

# Create and attach the DB wrappers
app.sys_db_h = SystemDbController()
app.d_id_db_h = DiscordIdDbController()
app.user_db_h = UserDbController()

# Configure JWT settings and initialize the JWTManager instance with the app
configure_jwt(app)
jwt.init_app(app) # Crucial: This initializes the global 'jwt' instance with your app


@app.before_serving
async def startup_websocket_server():
    """Starts the WebSocket server when the Quart app starts."""
    global websocket_server_instance
    websocket_server_address = "0.0.0.0"
    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.register_blueprint(systems_bp, url_prefix="/systems")
app.register_blueprint(nodes_bp, url_prefix="/nodes")
app.register_blueprint(bot_bp, url_prefix="/bots")
app.register_blueprint(auth_bp, url_prefix="/auth") # Register the auth blueprint

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

# --- Main Execution ---
if __name__ == "__main__":
    print("Starting Quart API server...")
    app.run(
        host="0.0.0.0",
        port=5000,
        debug=False # Set to True for development
    )
    print("Quart API server stopped.")