"""Dora node bridging WebSocket IO to Dora topics."""

from __future__ import annotations

import asyncio
import json
import os
import threading
import time
from typing import Any, Dict, Optional, Set

import pyarrow as pa
from dora import Node
from websockets.server import serve, WebSocketServerProtocol


class IoBridgeServer:
    def __init__(self, host: str, port: int):
        self.host = host
        self.port = port
        self.clients: Set[WebSocketServerProtocol] = set()
        self.command_handler = None
        self.scene_refresh_handler = None

    async def handler(self, websocket: WebSocketServerProtocol):
        self.clients.add(websocket)
        try:
            async for message in websocket:
                try:
                    data = json.loads(message)
                except json.JSONDecodeError:
                    await websocket.send(
                        json.dumps({"type": "error", "text": "Invalid JSON message"})
                    )
                    continue
                response = await self._route_message(data, websocket)
                if response:
                    await websocket.send(json.dumps(response))
        finally:
            self.clients.discard(websocket)

    async def _route_message(
        self, data: Dict[str, Any], websocket: WebSocketServerProtocol
    ) -> Optional[Dict[str, Any]]:
        msg_type = data.get("type")
        if msg_type == "command":
            text = data.get("text", "")
            if self.command_handler:
                await self.command_handler(text)
                return None
            return {"type": "error", "text": "No command handler registered"}
        if msg_type == "ping":
            return {"type": "pong"}
        if msg_type == "scene_refresh":
            if self.scene_refresh_handler:
                objects = await self.scene_refresh_handler()
                return {"type": "scene_updated", "objects": objects}
            return {"type": "error", "text": "No scene handler registered"}
        return {"type": "error", "text": f"Unknown message type: {msg_type}"}

    async def broadcast(self, message: Dict[str, Any]):
        if not self.clients:
            return
        payload = json.dumps(message)
        await asyncio.gather(
            *[client.send(payload) for client in self.clients], return_exceptions=True
        )

    async def send(self, message: Dict[str, Any], websocket: WebSocketServerProtocol):
        await websocket.send(json.dumps(message))

    async def start(self):
        async with serve(self.handler, self.host, self.port):
            await asyncio.Future()


def main() -> None:
    host = os.getenv("VOICE_HOST", "0.0.0.0")
    port = int(os.getenv("VOICE_PORT", "8765"))
    input_topic = os.getenv("VOICE_IN_OUTPUT", "voice_in")
    response_input = os.getenv("VOICE_OUT_INPUT", "voice_out")
    scene_input = os.getenv("SCENE_INPUT", "scene_update")

    node = Node()
    server = IoBridgeServer(host, port)
    loop = asyncio.new_event_loop()

    def push_command(text: str) -> None:
        node.send_output(
            input_topic,
            pa.array([text]),
            metadata={"encoding": "utf-8", "timestamp_ns": time.time_ns()},
        )

    async def handle_scene_refresh():
        return []

    def command_handler(text: str):
        push_command(text)
        return None

    server.command_handler = command_handler
    server.scene_refresh_handler = handle_scene_refresh

    def run_server():
        asyncio.set_event_loop(loop)
        loop.run_until_complete(server.start())

    threading.Thread(target=run_server, daemon=True).start()

    for event in node:
        if event["type"] != "INPUT":
            continue

        if event["id"] == response_input:
            raw = event["value"][0].as_py() if len(event["value"]) else ""
            if not raw:
                continue
            try:
                payload = json.loads(raw)
                message = {
                    "type": "response",
                    "text": payload.get("text", ""),
                    "status": payload.get("status", "ok"),
                }
            except Exception:
                message = {"type": "response", "text": raw, "status": "ok"}
            asyncio.run_coroutine_threadsafe(server.broadcast(message), loop)
            continue

        if event["id"] == scene_input:
            raw = event["value"][0].as_py() if len(event["value"]) else ""
            if not raw:
                continue
            try:
                payload = json.loads(raw)
                objects = payload.get("objects", [])
                message = {"type": "scene_updated", "objects": objects}
            except Exception:
                message = {"type": "scene_updated", "objects": []}
            asyncio.run_coroutine_threadsafe(server.broadcast(message), loop)


if __name__ == "__main__":
    main()
