import asyncio
import json
import logging
import os
import time

import websockets
from urllib.parse import urlparse, parse_qs

from qkernel import QKernel, EventBus, MemoryEngine, PolicyEngine, ObservabilityBrain, SelfHealingEngine
from qkernel.services import BrainService, VoiceService, NetworkService
from qkernel.auth import verify_token

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(name)s] %(levelname)s: %(message)s",
)
logger = logging.getLogger("ws_server")

RATE_LIMIT_WINDOW = 60
RATE_LIMIT_MAX = 30
_rate_limit: dict[str, list[float]] = {}


def check_rate_limit(peer: str) -> bool:
    now = time.time()
    if peer not in _rate_limit:
        _rate_limit[peer] = []
    _rate_limit[peer] = [t for t in _rate_limit[peer] if now - t < RATE_LIMIT_WINDOW]
    if len(_rate_limit[peer]) >= RATE_LIMIT_MAX:
        return False
    _rate_limit[peer].append(now)
    return True


kernel: QKernel | None = None


_ws_auth: dict = {}


async def chat_handler(websocket):
    auth = _ws_auth.get(id(websocket), {})
    async for message in websocket:
        try:
            data = json.loads(message)
            msg_type = data.get("type", "")

            if msg_type == "message":
                event = {
                    "type": "chat",
                    "message": data.get("text", ""),
                    "user_id": auth.get("user_id", "anonymous"),
                    "session_id": data.get("session_id", "main"),
                    "context": data.get("context", []),
                    "model": data.get("model", "deep"),
                }
                await websocket.send(json.dumps({"type": "typing"}))
                result = await kernel.handle(event)
                await websocket.send(json.dumps({
                    "type": "response",
                    "data": result,
                }))

            elif msg_type == "system":
                await websocket.send(json.dumps({
                    "type": "ack",
                    "text": "system_config_received",
                }))

            elif msg_type == "ping":
                await websocket.send(json.dumps({"type": "pong"}))

            elif msg_type == "trace":
                trace = kernel.obs.get_trace(data.get("trace_id", ""))
                await websocket.send(json.dumps({
                    "type": "trace_result",
                    "trace": trace,
                }))

            elif msg_type == "stats":
                stats = kernel.get_stats()
                await websocket.send(json.dumps({
                    "type": "stats",
                    "stats": stats,
                }))

        except json.JSONDecodeError:
            await websocket.send(json.dumps({"type": "error", "text": "invalid_json"}))
        except Exception as e:
            logger.exception("Chat handler error")
            await websocket.send(json.dumps({"type": "error", "text": str(e)}))


async def voice_handler(websocket):
    auth = _ws_auth.get(id(websocket), {})
    async for message in websocket:
        try:
            data = json.loads(message)
            msg_type = data.get("type", "")

            if msg_type == "command":
                text = data.get("text", "")
                action = data.get("action", "")

                event = {
                    "type": "voice",
                    "message": text,
                    "action": action,
                    "user_id": auth.get("user_id", "anonymous"),
                    "session_id": data.get("session_id", "main"),
                }

                result = await kernel.handle(event)

                await websocket.send(json.dumps({
                    "type": "command_result",
                    "action": action,
                    "text": text,
                    "data": result,
                }))

            elif msg_type == "ping":
                await websocket.send(json.dumps({"type": "pong"}))

        except json.JSONDecodeError:
            await websocket.send(json.dumps({"type": "error", "text": "invalid_json"}))


async def handler(websocket):
    path = websocket.path
    peer = websocket.remote_address

    params = parse_qs(urlparse(path).query)
    token = params.get("token", [None])[0]
    auth = verify_token(token) if token else None

    if not auth and token is not None:
        await websocket.send(json.dumps({"type": "error", "text": "invalid_token"}))
        return

    if auth:
        _ws_auth[id(websocket)] = auth

    if not check_rate_limit(peer):
        await websocket.send(json.dumps({"type": "error", "text": "rate_limited"}))
        return

    try:
        if path.endswith("/ws/voice"):
            await voice_handler(websocket)
        else:
            await chat_handler(websocket)
    finally:
        _ws_auth.pop(id(websocket), None)


async def init_kernel():
    global kernel

    memory = MemoryEngine()
    policy = PolicyEngine()
    obs = ObservabilityBrain()
    bus = EventBus()
    healer = SelfHealingEngine()

    kernel = QKernel(
        memory=memory,
        policy=policy,
        obs=obs,
        event_bus=bus,
        healer=healer,
    )

    policy.add_default_rules()

    kernel.register_service("brain", BrainService())
    kernel.register_service("voice", VoiceService())
    kernel.register_service("network", NetworkService())

    logger.info("Q-Kernel initialized with all subsystems")


async def main():
    await init_kernel()
    logger.info("Starting Q WebSocket Proxy on ws://0.0.0.0:8502")
    async with websockets.serve(
        handler,
        "0.0.0.0",
        8502,
        ping_interval=25,
        ping_timeout=10,
    ):
        await asyncio.Future()


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