from __future__ import annotations

import atexit
import asyncio
import os
from multiprocessing import shared_memory
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 ShmPacketTransport:
    """Shared-memory data plane with UDS control handoff."""

    _ACK_OK = b"\x01"
    _ACK_FAIL = b"\x00"

    def __init__(
        self,
        shm_path: str,
        handler: callable = None,
        *,
        max_payload_bytes: int = 67108864,
        max_inflight_handlers: int = 256,
        on_handler_error: callable | None = None,
    ):
        self.shm_path = shm_path
        self.handler = handler
        self.max_payload_bytes = max(1, int(max_payload_bytes))
        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"[SHM] Handler task failed: {exc}")
                return
            try:
                self._on_handler_error(exc)
            except Exception as cb_exc:  # pragma: no cover - defensive callback guard
                print(f"[SHM] Handler error callback failed: {cb_exc}")

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

        payload = cloudpickle.dumps(packet)
        payload_size = len(payload)
        self._validate_payload_size(payload_size, max_payload_bytes=self.max_payload_bytes)
        shm = shared_memory.SharedMemory(create=True, size=payload_size)
        try:
            shm.buf[:payload_size] = payload
            control = cloudpickle.dumps({"name": shm.name, "size": payload_size})
            reader, writer = await asyncio.open_unix_connection(peer_path)
            try:
                writer.write(len(control).to_bytes(4, "big") + control)
                await writer.drain()
                ack = await reader.readexactly(1)
                if ack != self._ACK_OK:
                    raise RuntimeError(f"[SHM] Peer {peer_path} rejected shared-memory payload.")
            finally:
                writer.close()
                await writer.wait_closed()
        except Exception as exc:
            raise RuntimeError(f"[SHM] Send failed to {peer_path}: {exc}") from exc
        finally:
            try:
                shm.close()
            finally:
                try:
                    shm.unlink()
                except FileNotFoundError:
                    pass

    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")
                control_data = await reader.readexactly(msg_len)
                control = cloudpickle.loads(control_data)
                shm_name = control["name"]
                payload_size = int(control["size"])
                self._validate_payload_size(
                    payload_size,
                    max_payload_bytes=self.max_payload_bytes,
                )
                shm = shared_memory.SharedMemory(name=shm_name)
                try:
                    packet_data = bytes(shm.buf[:payload_size])
                finally:
                    shm.close()
                request: FlownetPacket = cloudpickle.loads(packet_data)
                await self._dispatch_handler(request)
                writer.write(self._ACK_OK)
                await writer.drain()
            except Exception as exc:
                writer.write(self._ACK_FAIL)
                await writer.drain()
                if _transport_debug_logs_enabled():
                    print(f"[SHM] Connection lost: {exc}")
                break
        writer.close()
        await writer.wait_closed()

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

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

    @staticmethod
    def _validate_payload_size(payload_size: int, *, max_payload_bytes: int) -> None:
        if payload_size <= 0:
            raise ValueError(f"[SHM] Invalid payload size: {payload_size}.")
        if payload_size > max_payload_bytes:
            raise ValueError(
                f"[SHM] Payload size {payload_size} exceeds max {max_payload_bytes}."
            )
