# server.py
import asyncio
import websockets
import json
import uuid
from quart import Quart, jsonify, request
from quart_cors import cors
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
from internal.db_wrappers import SystemDbController, DiscordIdDbController #
from internal.auth_wrappers import UserDbController #
from config.jwt_config import jwt, configure_jwt


# --- WebSocket Server Components ---
active_clients = {}

async def websocket_server_handler(websocket):
    client_id = None
    try:
        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"}))
            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__)
app = cors(app, allow_origin="*")

websocket_server_instance = None
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(app)
jwt.init_app(app)


@app.before_serving
async def startup_tasks(): # Combined startup logic
    """Starts the WebSocket server and prepares other resources."""
    global websocket_server_instance
    websocket_server_address = "0.0.0.0"
    websocket_server_port = 8765
    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}")
    # Database connections are now established on first use by MongoHandler's __aenter__/connect
    # No explicit connect calls needed here unless desired for early failure detection.
    print("Application startup complete. DB connections will be initialized on first use.")

@app.after_serving
async def shutdown_tasks(): # Combined shutdown logic
    """Shuts down services and closes connections."""
    global websocket_server_instance
    if websocket_server_instance:
        websocket_server_instance.close()
        await websocket_server_instance.wait_closed()
        print("WebSocket server shut down.")

    # Close database connections
    if hasattr(app, 'user_db_h') and app.user_db_h:
        print("Closing User DB connection...")
        await app.user_db_h.close_db_connection() #
    if hasattr(app, 'sys_db_h') and app.sys_db_h:
        print("Closing System DB connection...")
        await app.sys_db_h.close_db_connection() #
    if hasattr(app, 'd_id_db_h') and app.d_id_db_h:
        print("Closing Discord ID DB connection...")
        await app.d_id_db_h.close_db_connection() #
    print("All database connections have been signaled to close.")


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")

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

if __name__ == "__main__":
    print("Starting Quart API server...")
    app.run(
        host="0.0.0.0",
        port=5000,
        debug=False
    )
    print("Quart API server stopped.")