"""Async TCP server runtime for CRP."""

from __future__ import annotations

import asyncio
import logging
import socket
from contextlib import suppress

from cliproom.protocol import (
    ClipMessage,
    ErrorMessage,
    HeartbeatMessage,
    DiscoverMessage,
    DiscoverReplyMessage,
    JoinAckMessage,
    JoinMessage,
    LeaveMessage,
    ProtocolValidationError,
    decode_message_line,
    encode_message,
    DEFAULT_DISCOVERY_PORT,
)

from .session_manager import ClientConnection, SessionManager


LOG = logging.getLogger("cliproom.server")


class ClipRoomServer:
    """Stateful CRP server with session routing."""

    def __init__(
        self,
        host: str,
        port: int,
        *,
        fixed_session: str | None = None,
        fixed_auth_hash: str | None = None,
        discovery_port: int = DEFAULT_DISCOVERY_PORT,
        enable_discovery: bool = True,
    ) -> None:
        self.host = host
        self.port = port
        self.discovery_port = discovery_port
        self.enable_discovery = enable_discovery
        self._session_manager = SessionManager(
            fixed_session=fixed_session,
            fixed_auth_hash=fixed_auth_hash,
        )
        self._server: asyncio.AbstractServer | None = None
        self._discovery_transport: asyncio.DatagramTransport | None = None

    async def start(self) -> None:
        self._server = await asyncio.start_server(self._handle_client, self.host, self.port)
        sockets = self._server.sockets or []
        addresses = ", ".join(str(sock.getsockname()) for sock in sockets)
        LOG.info("Server listening on %s", addresses)
        if self.enable_discovery:
            await self._start_discovery()

    async def serve_forever(self) -> None:
        if self._server is None:
            await self.start()
        assert self._server is not None
        async with self._server:
            await self._server.serve_forever()

    async def stop(self) -> None:
        if self._server is None:
            return
        self._server.close()
        await self._server.wait_closed()
        self._server = None
        if self._discovery_transport is not None:
            self._discovery_transport.close()
            self._discovery_transport = None

    async def _handle_client(
        self,
        reader: asyncio.StreamReader,
        writer: asyncio.StreamWriter,
    ) -> None:
        peer = writer.get_extra_info("peername")
        connection: ClientConnection | None = None
        session_name: str | None = None
        LOG.info("Client connected: %s", peer)
        try:
            while True:
                raw_line = await reader.readline()
                if not raw_line:
                    break

                message = decode_message_line(raw_line)
                if connection is None:
                    if not isinstance(message, JoinMessage):
                        raise ProtocolValidationError("join_required")

                    connection = ClientConnection(
                        client_name=message.client_name,
                        writer=writer,
                    )
                    session = await self._session_manager.join(
                        session_name=message.session,
                        auth_hash=message.auth,
                        connection=connection,
                    )
                    session_name = session.name
                    await connection.send(encode_message(JoinAckMessage(session=session.name)))
                    LOG.info("Client %s joined session %s", message.client_name, session_name)
                    continue

                if isinstance(message, ClipMessage):
                    if message.session != session_name:
                        raise ProtocolValidationError("invalid_session")
                    await self._broadcast(session_name, connection, encode_message(message))
                    continue

                if isinstance(message, (HeartbeatMessage,)):
                    continue

                if isinstance(message, LeaveMessage):
                    break

                raise ProtocolValidationError(f"unsupported_after_join:{message.type}")
        except ProtocolValidationError as exc:
            LOG.warning("Protocol error from %s: %s", peer, exc)
            await self._safe_send_error(writer, str(exc))
        except Exception:
            LOG.exception("Unhandled error serving client %s", peer)
            await self._safe_send_error(writer, "internal_error")
        finally:
            if connection is not None and session_name is not None:
                await self._session_manager.leave(session_name, connection)
            writer.close()
            with suppress(Exception):
                await writer.wait_closed()
            LOG.info("Client disconnected: %s", peer)

    async def _broadcast(
        self,
        session_name: str,
        origin: ClientConnection,
        payload: bytes,
    ) -> None:
        targets = [
            client
            for client in await self._session_manager.list_clients(session_name)
            if client is not origin
        ]
        if not targets:
            return
        await asyncio.gather(*(client.send(payload) for client in targets), return_exceptions=True)

    async def _safe_send_error(self, writer: asyncio.StreamWriter, message: str) -> None:
        try:
            writer.write(encode_message(ErrorMessage(message=message)))
            await writer.drain()
        except Exception:
            pass

    async def _start_discovery(self) -> None:
        loop = asyncio.get_running_loop()
        transport, _ = await loop.create_datagram_endpoint(
            lambda: _DiscoveryServerProtocol(self),
            local_addr=("0.0.0.0", self.discovery_port),
            allow_broadcast=True,
        )
        self._discovery_transport = transport
        LOG.info("Discovery listening on UDP %s", self.discovery_port)

    async def _build_discovery_replies(
        self, requester_addr: tuple[str, int]
    ) -> list[bytes]:
        host = self._advertise_host_for(requester_addr[0])
        replies: list[bytes] = []
        for session_name, requires_passphrase in (
            await self._session_manager.list_discoverable_sessions()
        ):
            reply = DiscoverReplyMessage(
                session=session_name,
                requires_passphrase=requires_passphrase,
                host=host,
                port=self.port,
            )
            replies.append(encode_message(reply))
        return replies

    def _advertise_host_for(self, requester_host: str) -> str:
        if self.host not in {"0.0.0.0", "::"}:
            return self.host
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        try:
            sock.connect((requester_host, 9))
            local_ip = sock.getsockname()[0]
            return local_ip
        except OSError:
            return "127.0.0.1"
        finally:
            sock.close()


class _DiscoveryServerProtocol(asyncio.DatagramProtocol):
    def __init__(self, server: ClipRoomServer) -> None:
        self.server = server
        self.transport: asyncio.DatagramTransport | None = None

    def connection_made(self, transport: asyncio.BaseTransport) -> None:
        self.transport = transport  # type: ignore[assignment]

    def datagram_received(self, data: bytes, addr: tuple[str, int]) -> None:
        if self.transport is None:
            return
        try:
            message = decode_message_line(data)
            if not isinstance(message, DiscoverMessage):
                return
        except Exception:
            return
        asyncio.create_task(self._reply_to_discover(addr))

    async def _reply_to_discover(self, addr: tuple[str, int]) -> None:
        if self.transport is None:
            return
        for payload in await self.server._build_discovery_replies(addr):
            self.transport.sendto(payload, addr)


async def run_server(
    host: str,
    port: int,
    *,
    fixed_session: str | None = None,
    fixed_auth_hash: str | None = None,
    discovery_port: int = DEFAULT_DISCOVERY_PORT,
    enable_discovery: bool = True,
) -> None:
    """Create and run a server process until interrupted."""
    server = ClipRoomServer(
        host=host,
        port=port,
        fixed_session=fixed_session,
        fixed_auth_hash=fixed_auth_hash,
        discovery_port=discovery_port,
        enable_discovery=enable_discovery,
    )
    await server.serve_forever()
