from __future__ import annotations

import asyncio
import logging
from uuid import uuid4

from yuxi.channel.application.service.inbound_service import InboundService
from yuxi.channel.domain.port.channel_adapter_port import ChannelAdapterPort
from yuxi.channel.domain.port.metrics_port import MetricsPort
from yuxi.channel.domain.port.ws_connection_port import WsConnectionPort

logger = logging.getLogger(__name__)


class WsConnectionManager:
    def __init__(self, metrics: MetricsPort | None = None) -> None:
        self._connections: dict[str, WsConnectionPort] = {}
        self._adapters: dict[str, ChannelAdapterPort] = {}
        self._loop: asyncio.AbstractEventLoop | None = None
        self._inbound_service: InboundService | None = None
        self._metrics = metrics

    def register(self, connection: WsConnectionPort) -> None:
        self._connections[connection.channel_type] = connection

    def set_adapters(self, adapters: dict[str, ChannelAdapterPort]) -> None:
        self._adapters = adapters

    async def start_all(
        self,
        loop: asyncio.AbstractEventLoop,
        inbound_service: InboundService,
    ) -> None:
        self._loop = loop
        self._inbound_service = inbound_service
        if self._connections:
            await asyncio.gather(
                *(conn.start(loop, self._on_message) for conn in self._connections.values()),
                return_exceptions=True,
            )

    async def stop_all(self) -> None:
        for conn in self._connections.values():
            try:
                await conn.stop()
            except Exception:
                logger.warning("ws stop failed for %s", conn.channel_type)

    async def _on_message(self, raw: dict) -> None:
        channel_type = raw.get("channel_type", "")
        adapter = self._adapters.get(channel_type)
        if not adapter:
            logger.warning("no adapter for ws message from %s", channel_type)
            return

        try:
            message = await adapter.receive_message(raw)
        except Exception:
            logger.exception("ws receive_message failed for %s", channel_type)
            return

        trace_id = raw.get("header", {}).get("event_id", str(uuid4()))
        message.metadata["trace_id"] = trace_id
        message.metadata["source"] = "websocket"
        message.metadata["idempotency_key"] = trace_id

        if self._inbound_service:
            try:
                await self._inbound_service.submit(message, channel_type=channel_type, trace_id=trace_id)
            except Exception:
                logger.exception("ws inbound submit failed for %s", channel_type)

    @property
    def connections_status(self) -> dict[str, bool]:
        return {ct: conn.is_connected for ct, conn in self._connections.items()}