from __future__ import annotations

import atexit
import asyncio
import os
from typing import TYPE_CHECKING

import cloudpickle

from flownet.runtime import loops

if TYPE_CHECKING:
    from flownet.core.request import FlownetPacket
    from flownet.runtime.transport.packet_transport import EndpointRef


def _transport_debug_logs_enabled() -> bool:
    value = (os.environ.get("FLOWNET_TRANSPORT_DEBUG_LOGS") or "").strip().lower()
    return value in {"1", "true", "yes", "on"}


class UdsPacketTransport:
    def __init__(
        self,
        uds_path: str,
        handler: callable = None,
        *,
        max_inflight_handlers: int = 256,
        on_handler_error: callable | None = None,
    ):
        self.uds_path = uds_path
        self.handler = handler
        self.connections: dict[str, asyncio.StreamWriter] = {}
        self._handler_max_inflight = max(1, int(max_inflight_handlers))
        self._handler_semaphore = asyncio.Semaphore(self._handler_max_inflight)
        self._handler_inflight = 0
        self._handler_error_count = 0
        self._on_handler_error = on_handler_error
        atexit.register(self._cleanup_socket_path)
        loops.submit_runtime_coro(self.start())

    def handler_metrics_snapshot(self) -> dict[str, int]:
        return {
            "handler_max_inflight": self._handler_max_inflight,
            "handler_inflight": self._handler_inflight,
            "handler_errors": self._handler_error_count,
        }

    async def _dispatch_handler(self, request: FlownetPacket) -> None:
        if self.handler is None:
            return
        await self._handler_semaphore.acquire()
        self._handler_inflight += 1
        task = asyncio.get_running_loop().create_task(self.handler(request))
        task.add_done_callback(self._on_handler_task_done)

    def _on_handler_task_done(self, task: asyncio.Task) -> None:
        self._handler_inflight = max(0, self._handler_inflight - 1)
        self._handler_semaphore.release()
        try:
            task.result()
        except asyncio.CancelledError:
            return
        except Exception as exc:
            self._handler_error_count += 1
            if self._on_handler_error is None:
                print(f"[UDS] Handler task failed: {exc}")
                return
            try:
                self._on_handler_error(exc)
            except Exception as cb_exc:  # pragma: no cover - defensive callback guard
                print(f"[UDS] Handler error callback failed: {cb_exc}")

    async def connect_to(self, path: str) -> asyncio.StreamWriter:
        writer = self.connections.get(path)
        if writer is not None:
            return writer
        _reader, writer = await asyncio.open_unix_connection(path)
        self.connections[path] = writer
        return writer

    async def send(
        self,
        address: str,
        packet: FlownetPacket,
        endpoint: EndpointRef | None = None,
    ) -> None:
        peer_path = endpoint.uds_path if endpoint is not None else None
        if not peer_path:
            raise RuntimeError(f"[UDS] Missing UDS endpoint for address {address}")

        writer = self.connections.get(peer_path)
        if writer is None:
            try:
                writer = await self.connect_to(peer_path)
            except Exception as exc:
                raise RuntimeError(f"[UDS] Failed to connect to {peer_path}: {exc}") from exc

        try:
            data = cloudpickle.dumps(packet)
            writer.write(len(data).to_bytes(4, "big") + data)
            await writer.drain()
        except Exception as exc:
            self.connections.pop(peer_path, None)
            raise RuntimeError(f"[UDS] Send failed to {peer_path}: {exc}") from exc

    async def handle_client(
        self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter
    ) -> None:
        while not reader.at_eof():
            try:
                raw_len = await reader.readexactly(4)
                msg_len = int.from_bytes(raw_len, "big")
                data = await reader.readexactly(msg_len)
                request: FlownetPacket = cloudpickle.loads(data)
                await self._dispatch_handler(request)
            except Exception as exc:
                if _transport_debug_logs_enabled():
                    print(f"[UDS] Connection lost: {exc}")
                break
        writer.close()
        await writer.wait_closed()

    async def start(self) -> None:
        try:
            os.makedirs(os.path.dirname(self.uds_path), exist_ok=True)
            self._cleanup_socket_path()
            server = await asyncio.start_unix_server(self.handle_client, path=self.uds_path)
            if _transport_debug_logs_enabled():
                print(f"[UDS] Listening on {self.uds_path}")
            async with server:
                await server.serve_forever()
        except Exception as exc:
            print(f"[UDS ERROR] Failed to bind on {self.uds_path} -> {exc}")
        finally:
            self._cleanup_socket_path()

    def _cleanup_socket_path(self) -> None:
        try:
            if os.path.exists(self.uds_path):
                os.unlink(self.uds_path)
        except FileNotFoundError:
            return
        except Exception:
            return
