import anyio
from anyio import move_on_after, create_memory_object_stream, create_task_group
from anyio.abc import SocketStream

from contextlib import AsyncExitStack

import socket as std_socket

from socket import SOCK_RAW

import struct

import anyio.abc
import anyio.streams
import anyio.streams.memory
import anyio.to_thread
from toolz.curried import pipe, map, filter, reduceby, concat, groupby, reduce
from typing import AsyncIterator, Iterator, Optional


from abc import abstractmethod

from oxitrait.trait import Trait
from oxitrait.impl import Impl
from oxitrait.struct import Struct
from oxitrait.enum import Enum, auto
from oxitrait.runtime import requires_traits

from result import Result, Ok, Err

from ..log import logger
from .consts import (
    AF_NETLINK,
    NLMSG_ALIGN,
    NLMSG_DONE,
    NLMSG_HDR_FORMAT,
    NLMSG_HDR_SIZE
)
from .rtnl.consts import (
    NDUSEROPT_MAX
)
from .message import NetlinkHeaderTrait
from .parsing import parse_peek, parse_full_message


class NetlinkSyncContextTrait(metaclass=Trait):
    @abstractmethod
    def __enter__(self):
        pass

    @abstractmethod
    def __exit__(self, exc_type, exc_value, traceback):
        pass

class NetlinkAsyncContextTrait(metaclass=Trait):
    @abstractmethod
    async def __aenter__(self):
        pass

    @abstractmethod
    async def __aexit__(self, exc_type, exc_value, traceback):
        pass

class NetlinkSocketTrait(metaclass=Trait):
    @abstractmethod
    def __post_init__(self):
        pass

    @abstractmethod
    async def initialize(self, task_group) -> Result[None, str]:
        """
        Open/bind the socket, spawn any needed background tasks.
        """
        pass

    @abstractmethod
    async def send_message(self, message: NetlinkHeaderTrait) -> Result[None, str]:
        """
        Serialize and send a netlink message.
        """
        pass

    @abstractmethod
    async def receive_message(self, timeout: float = 5.0) -> Result[NetlinkHeaderTrait, str]:
        """
        Retrieve one netlink message from the buffer (populated by a background task).
        """
        pass

    @abstractmethod
    async def close(self) -> Result[None, str]:
        """
        Cancel any background tasks and close the underlying socket.
        """
        pass

#
# 2) Provide the Implementation for that Trait
#
class ImplNetlinkSocket(NetlinkSocketTrait, metaclass=Impl, target="NetlinkSocket"):
    """
    Implementation for the NetlinkSocket struct.
    """

    def __post_init__(self):
        self.sock: Optional[std_socket.socket] = None
        self._message_buffer: dict[int, list[bytes]] = {}
    
        send_chan, recv_chan = create_memory_object_stream(100)
        self.inbound_send = send_chan
        self.inbound_recv = recv_chan

    async def initialize(self, task_group: anyio.abc.TaskGroup) -> Result[None, str]:
        """
        Initialize the netlink socket and spawn the _read_loop inside the provided task group.
        The caller is responsible for managing the lifecycle of the task group.
        """
        logger.trace("Entering initialize()")

        if self.sock is not None:
            return Ok(None)  # Already initialized

        try:
            # Create and bind raw Netlink socket
            proto: int = self.protocol
            assert isinstance(proto, int)
            assert isinstance(AF_NETLINK, int)
            raw_sock = std_socket.socket(AF_NETLINK, SOCK_RAW, proto)
            raw_sock.bind((0, 0))

            # Assign directly — no need to wrap in anyio stream
            self.sock = raw_sock

            self._running = True
            logger.debug(f"Is self running in initialize? Answer: {self._running}")

            # Store and use provided task group to spawn background reader
            self.tg = task_group
            self.tg.start_soon(self._read_loop)

            logger.trace("Exiting initialize()")
            return Ok(None)

        except Exception as e:
            return Err(f"Failed to initialize socket: {e}")

    @staticmethod
    def _concat(fragments: list[bytes]) -> bytes:
        return b"".join(fragments)

    @staticmethod
    def _log_and_discard(reason: str):
        logger.warning(f"_dispatch() skipped: {reason}")

    @staticmethod
    def _is_valid_frame(f: bytes) -> bool:
        return len(f) >= NLMSG_HDR_SIZE

    async def _dispatch(self, fragments: list[bytes]):
        logger.debug(f"_dispatch called with {len(fragments)} fragment(s)")

        def log_and_parse(raw: bytes):
            logger.debug(f"Raw message length: {len(raw)} bytes")
            logger.debug(f"Raw bytes (first 32): {raw[:32].hex()}")
            result = parse_full_message(raw)
            logger.debug("Called parse_full_message()")
            if result.is_ok():
                logger.debug("parse_full_message returned Ok")
            else:
                logger.error(f"parse_full_message returned Err: {result.unwrap_err()}")
            return result

        async def send_if_ok(result):
            if result.is_ok():
                logger.debug("Sending message to inbound channel")
                await self.inbound_send.send(result.unwrap())
            else:
                logger.error(f"Not sending, result was Err: {result.unwrap_err()}")

        return await pipe(
            fragments,
            b"".join,         # → bytes
            log_and_parse,    # → Result[dict, str]
            send_if_ok        # → awaitable effect
        )

    async def send_message(self, message: NetlinkHeaderTrait) -> Result[None, str]:
        """
        Send a serialized NetlinkMessage using a blocking system socket.
        """
        if not self.sock:
            return Err("Socket not initialized. Call initialize() first.")

        try:
            # Convert the message to bytes, if needed
            data = bytes(message)  # assuming message implements __bytes__()

            await anyio.to_thread.run_sync(self.sock.send, data)
            return Ok(None)
        except Exception as e:
            return Err(f"Send error: {e}")

    async def receive_message(self, timeout: float = 1.0) -> Result[dict, str]:
        """
        Retrieve one Netlink message from the inbound queue, subject to timeout.
        """
        if not self.inbound_recv:
            logger.error("Socket not initialized. inbound_recv channel is missing.")
            return Err("Socket not initialized. inbound_recv channel is missing.")

        logger.debug("⏳ Entering receive_message(), waiting for a message...")

        with move_on_after(timeout) as cancel_scope:
            try:
                msg = await self.inbound_recv.receive()

                if cancel_scope.cancelled_caught:
                    logger.warning("receive_message() timed out waiting for a message")
                    return Err("Timeout occurred waiting for Netlink message.")

                logger.debug("Message received in receive_message()")
                return Ok(msg)

            except Exception as e:
                logger.exception("Exception while receiving from inbound channel")
                return Err(e)

    async def listen(self, groups: list[int]) -> AsyncIterator[tuple[dict, bytes]]:
        """
        Subscribes to given multicast groups and yields unsolicited Netlink messages.

        This generator runs indefinitely until cancelled.
        """
        if not self.sock:
            raise RuntimeError("Socket not initialized. Call initialize() first.")

        if not getattr(self, "_multicast_bound", False):
            try:
                group_mask = sum(1 << (g - 1) for g in groups)
                self.sock.bind((0, group_mask))
                self._multicast_bound = True
                logger.info(f"📡 Subscribed to groups: {groups} (mask={group_mask})")
            except Exception as e:
                logger.error(f"Failed to bind for multicast: {e}")
                return  # silent fail — stops generator

        while self._running:
            try:
                msg = await self.inbound_recv.receive()
                if isinstance(msg, dict) and "header" in msg and "payload" in msg:
                    yield msg["header"], msg["payload"]
                else:
                    logger.warning("listen(): malformed msg")
            except anyio.get_cancelled_exc_class():
                logger.debug("listen(): cancelled")
                break
            except Exception as e:
                logger.warning(f"listen() error: {e}")

    async def close(self) -> Result[None, str]:
        """
        Stop the background loop, close the socket, and clean up channels.
        """
        self._running = False
        if self.sock:
            try:
                await self.sock.close()
            except Exception as e:
                return Err(f"Close error: {e}")
            finally:
                self.sock = None

        # If you want to also terminate the read loop more definitively:
        if self.tg:
            self.tg.cancel_scope.cancel()

        # Close the memory channel
        if self.inbound_send:
            await self.inbound_send.close()

        return Ok(None)

    async def _rx_chunks(self, sock):
        while True:
            try:
                logger.trace("⛏ Calling sock.recv()")
                data = await anyio.to_thread.run_sync(sock.recv, 65536)
                logger.trace(f"Got {len(data)} bytes from kernel")
                yield data
            except OSError as e:
                if self._running:
                    logger.warning(f"Netlink recv OSError: {e}")
                break
            except anyio.get_cancelled_exc_class():
                break

    # 2) pure helpers
    @staticmethod
    def _is_valid_header(data: bytes, offset: int) -> bool:
        return len(data) - offset >= NLMSG_HDR_SIZE

    @staticmethod
    def _unpack_len(data: bytes, offset: int) -> int:
        return struct.unpack_from(NLMSG_HDR_FORMAT, data, offset)[0]

    @staticmethod
    def _can_fit_frame(data: bytes, offset: int, aligned_len: int) -> bool:
        return offset + aligned_len <= len(data)

    @staticmethod
    def split_frames(buf: bytes) -> Iterator[bytes]:
        """
        Yield Netlink-aligned message frames from a raw buffer.

        Ensures:
        - Header exists (min 16 bytes)
        - msg_len is >= header size
        - msg_len doesn't run past buffer end
        """
        def _next_offset(offset: int, msg_len: int) -> int:
            return NLMSG_ALIGN(offset + msg_len)

        offset = 0
        buflen = len(buf)

        while offset + NLMSG_HDR_SIZE <= buflen:
            try:
                (msg_len,) = struct.unpack_from("=I", buf, offset)
            except struct.error:
                logger.error(f"struct.unpack_from failed at offset={offset}")
                break

            if msg_len < NLMSG_HDR_SIZE:
                logger.error(f"Invalid Netlink length: {msg_len} < header size ({NLMSG_HDR_SIZE})")
                break

            end = offset + msg_len
            if end > buflen:
                logger.error(f"Truncated Netlink message: {end} > buffer length {buflen}")
                break

            yield buf[offset:end]
            offset = _next_offset(offset, msg_len)

    @staticmethod
    def parse_header(frame: bytes) -> Result[tuple[dict, bytes], str]:
        try:
            ln, ty, fl, seq, pid = struct.unpack_from(NLMSG_HDR_FORMAT, frame)
            hdr = {"len": ln, "type": ty, "flags": fl, "seq": seq, "pid": pid}
            return Ok((hdr, frame))             #  ← keep full bytes here
        except Exception as e:
            return Err(str(e))

    @staticmethod
    def _is_valid_chunk(raw: bytes) -> bool:
        return len(raw) >= NLMSG_HDR_SIZE

    async def _read_loop(self):
        logger.trace("Starting Netlink read loop")

        async for chunk in self._rx_chunks(self.sock):
            logger.debug(f"📥 Received {len(chunk)} bytes from socket")

            try:
                frames = pipe(
                    chunk,
                    self.split_frames,
                    filter(self._is_valid_chunk),
                    map(self.parse_header),
                    filter(self._is_ok),
                    map(self._unwrap),
                )
            except Exception as e:
                logger.exception("Exception during frame processing")
                continue

            frame_list = list(frames)  # materialize for logging & reuse
            logger.debug(f"Extracted {len(frame_list)} valid frame(s) from chunk")

            try:
                by_seq = pipe(
                    frame_list,
                    groupby(self._by_seq_key),
                )
            except Exception as e:
                logger.exception("Exception during sequence grouping")
                continue

            logger.debug(f"Grouped into {len(by_seq)} sequence(s): {list(by_seq.keys())}")

            for seq, group in by_seq.items():
                logger.debug(f"Processing seq {seq} with {len(group)} message(s)")

                if seq not in self._message_buffer:
                    logger.debug(f"Creating new buffer for seq {seq}")
                    self._message_buffer[seq] = []

                seen_done = False

                for hdr, full_frame in group:
                    logger.debug(f"Msg type={hdr['type']} len={hdr['len']} pid={hdr['pid']}")
                    self._message_buffer[seq].append(full_frame)

                    if hdr["type"] == NLMSG_DONE:
                        seen_done = True
                        logger.debug(f"Detected NLMSG_DONE in seq {seq}")

                if seen_done:
                    fragments = self._message_buffer.pop(seq)
                    logger.debug(f"Dispatching {len(fragments)} buffered frame(s) for seq {seq}")
                    try:
                        await self._dispatch(fragments)
                    except Exception as e:
                        logger.exception(f"Exception during dispatch of seq {seq}")

            logger.trace("Waiting for next Netlink message...")

    @staticmethod
    def _is_ok(result):
        return result.is_ok()

    @staticmethod
    def _unwrap(result):
        return result.unwrap()

    @staticmethod
    def _by_seq_key(tpl):
        header, _ = tpl
        return header["seq"]

class ImplNetlinkSyncContext(NetlinkSyncContextTrait, metaclass=Impl, target="NetlinkSocket"):
    def __enter__(self):
        proto: int = self.protocol
        assert isinstance(proto, int)
        assert isinstance(AF_NETLINK, int)
        self.sock = std_socket.socket(AF_NETLINK, SOCK_RAW, proto)
        return self  # Allow `with NetlinkSocket() as pyr:`

    def __exit__(self, exc_type, exc_value, traceback):
        if self.sock:
            self.sock.close()

class ImplNetlinkAsyncContext(NetlinkAsyncContextTrait, metaclass=Impl, target="NetlinkSocket"):
    async def __aenter__(self):
        # Structured async context manager stack
        self._exit_stack = AsyncExitStack()
        await self._exit_stack.__aenter__()

        # Create and bind the raw Netlink socket
        proto: int = self.protocol
        assert isinstance(proto, int)
        assert isinstance(AF_NETLINK, int)
        self.sock = std_socket.socket(AF_NETLINK, SOCK_RAW, proto)
        self.sock.bind((0, 0))
        self._running = True

        # Launch background read loop under managed task group
        self.tg = await self._exit_stack.enter_async_context(create_task_group())
        self.tg.start_soon(self._read_loop)

        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        self._running = False

        # Safely close the socket, if open
        if self.sock:
            self.sock.close()
            self.sock = None

        # Exit all context managers in LIFO order
        return await self._exit_stack.__aexit__(exc_type, exc_val, exc_tb)

class NetlinkSocket(metaclass=Struct):
    """
    A concrete struct (dataclass) that includes:
      - protocol: the netlink protocol to use (e.g., NETLINK_ROUTE)
      - sock: the actual Trio socket object
      - a background task group
      - channels for buffering inbound messages
      - a running flag
    """

    protocol: int
    sock: Optional[std_socket.socket] = None
    tg: Optional[anyio.abc.TaskGroup] = None
    _running: bool = False
    _message_buffer: Optional[dict[int, list[bytes]]] = None

    inbound_send: Optional[anyio.streams.memory.MemoryObjectSendStream[NetlinkHeaderTrait]] = None
    inbound_recv: Optional[anyio.streams.memory.MemoryObjectReceiveStream[NetlinkHeaderTrait]] = None
