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

Generic I/O bridge with fixed endpoints:
- Outputs (WS → Dora): text_out, audio_out, data_out
- Inputs (Dora → WS): text_in, audio_in, data_in
"""

from __future__ import annotations

import asyncio
import base64
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:
    """WebSocket server that bridges clients to Dora dataflow."""

    def __init__(self, host: str, port: int):
        self.host = host
        self.port = port
        self.clients: Set[WebSocketServerProtocol] = set()
        # Callbacks for sending to Dora
        self.on_text: Optional[callable] = None
        self.on_audio: Optional[callable] = None
        self.on_data: Optional[callable] = None

    async def handler(self, websocket: WebSocketServerProtocol):
        """Handle WebSocket connection."""
        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",
                        "message": "Invalid JSON"
                    }))
                    continue
                response = await self._route_message(data)
                if response:
                    await websocket.send(json.dumps(response))
        finally:
            self.clients.discard(websocket)

    async def _route_message(self, data: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """Route incoming message to appropriate handler."""
        msg_type = data.get("type")

        if msg_type == "ping":
            return {"type": "pong"}

        if msg_type == "text":
            content = data.get("content", "")
            if self.on_text and content:
                self.on_text(content)
            return None

        if msg_type == "audio":
            content = data.get("content", "")  # base64 WAV
            if self.on_audio and content:
                self.on_audio(content)
            return None

        if msg_type == "data":
            payload = data.get("payload", {})
            if self.on_data:
                self.on_data(payload)
            return None

        return {"type": "error", "message": f"Unknown type: {msg_type}"}

    async def broadcast(self, message: Dict[str, Any]):
        """Broadcast message to all connected clients."""
        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 start(self):
        """Start the WebSocket server."""
        async with serve(self.handler, self.host, self.port):
            await asyncio.Future()


def main() -> None:
    """Main entry point."""
    host = os.getenv("WS_HOST", "0.0.0.0")
    port = int(os.getenv("WS_PORT", "9000"))

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

    # Callbacks: WebSocket → Dora
    def send_text(content: str) -> None:
        node.send_output(
            "text_out",
            pa.array([content]),
            {"encoding": "utf-8", "timestamp_ns": time.time_ns()},
        )

    def send_audio(content_b64: str) -> None:
        audio_bytes = base64.b64decode(content_b64)
        node.send_output(
            "audio_out",
            pa.array([audio_bytes]),
            {"encoding": "wav", "timestamp_ns": time.time_ns()},
        )

    def send_data(payload: Dict[str, Any]) -> None:
        node.send_output(
            "data_out",
            pa.array([json.dumps(payload)]),
            {"encoding": "json", "timestamp_ns": time.time_ns()},
        )

    server.on_text = send_text
    server.on_audio = send_audio
    server.on_data = send_data

    # Start WebSocket server in background thread
    def run_server():
        asyncio.set_event_loop(loop)
        loop.run_until_complete(server.start())

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

    timestamp = time.strftime("%H:%M:%S")
    print(f"[iobridge {timestamp}] WebSocket server at ws://{host}:{port}", flush=True)

    # Dora event loop: Dora → WebSocket
    for event in node:
        if event["type"] != "INPUT":
            continue

        event_id = event["id"]
        raw = event["value"][0].as_py() if len(event["value"]) else None
        if raw is None:
            continue

        if event_id == "text_in":
            # Text from Dora to broadcast
            if isinstance(raw, str):
                try:
                    # Try to parse as JSON for structured response
                    payload = json.loads(raw)
                    message = {"type": "text", "content": payload.get("text", raw)}
                except json.JSONDecodeError:
                    message = {"type": "text", "content": raw}
            else:
                message = {"type": "text", "content": str(raw)}
            asyncio.run_coroutine_threadsafe(server.broadcast(message), loop)

        elif event_id == "audio_in":
            # Audio from Dora to broadcast (binary → base64)
            if isinstance(raw, bytes):
                content_b64 = base64.b64encode(raw).decode("utf-8")
            else:
                content_b64 = raw
            message = {"type": "audio", "content": content_b64, "format": "wav"}
            asyncio.run_coroutine_threadsafe(server.broadcast(message), loop)

        elif event_id == "data_in":
            # Generic JSON data from Dora to broadcast
            try:
                payload = json.loads(raw) if isinstance(raw, str) else raw
            except json.JSONDecodeError:
                payload = {"raw": raw}
            message = {"type": "data", "payload": payload}
            asyncio.run_coroutine_threadsafe(server.broadcast(message), loop)


if __name__ == "__main__":
    main()
