"""
This package implements some basic infrastructure for handling data
streams in a structured manner. It contains three distinct classes:

- `BaseMsg` transports Python objects, using `send` and `recv`.
- `BaseBlk` transports delimited bytestrings, using `snd` and `rcv`.
- `BaseBuf` transports undelimited bytestrings, using `wr` and `rd`.

Additionally, message- and block-based classes understand `cwr` and `crd`
which transport out-of-band data. Typically these contain raw console bytes
that are interleaved with structured data; they are used on channels which
needs to multiplex both.

The common theme for using these classes is

- You assemble a communication stack bottom-up: start with a serial link;
  add packetizing, retransmission, and an object-codec.
- Using the top level as an async context managers establishes the complete
  stack; leaving the context tears it down.
- Most likely, connect the result to a [Base]MsgHandler.

Everything is fully asynchronous. There is no "new incoming data" callback:
it's the upper layer's job to repeatedly call the appropriate method to
read or receive new messages.

A `wrap` method provides a secondary context that can be used for
a persistent outer context, e.g. to keep a listening socket open.
"""

from __future__ import annotations

from moat.util import attrdict
from moat.lib.micro import ACM, AC_exit, AC_use

# Typing

from typing import TYPE_CHECKING  # isort:skip

if TYPE_CHECKING:
    from contextlib import AbstractAsyncContextManager
    from types import TracebackType

    from typing import Any, Self

    Buffer = bytes | bytearray | memoryview
    MutBuffer = bytearray | memoryview
    _AACMBase = AbstractAsyncContextManager
else:
    Buffer = bytes
    MutBuffer = bytearray
    _AACMBase = object


class _NullCtx:
    async def __aenter__(self) -> Self:
        return self

    async def __aexit__(self, *tb):
        pass


_nullctx = _NullCtx()


class Base(_AACMBase):
    """
    The MoaT stream base class for "something connected".

    This class *must* be used as an async context manager.

    Usage:

    Use the :py:func:`moat.lib.micro.AC_use` helper if you need
    to call an async context manager or to register a destructor.

    Augment `setup` or `teardown` to add non-stream related features.

    Override `wrap` to contain an async context manager that holds resources
    which must survive reconnection, e.g. a MQTT link's persistent state or
    a listening socket.
    """

    s = None
    cfg: attrdict

    def __init__(self, cfg: attrdict | None = None):
        if cfg is None:
            cfg = attrdict()
        self.cfg = cfg

    @property
    def cfg_name(self):
        "Some name to disambiguate this part"
        try:
            return self.cfg.name
        except AttributeError:
            return self.__class__.__name__

    def wrap(self) -> AbstractAsyncContextManager:
        """
        Async context manager for holding a cross-connection context.

        By default does nothing.
        """
        return _nullctx

    async def __aenter__(self) -> Self:
        res = self
        AC = ACM(self)
        await AC(self.teardown)
        try:
            await self.setup()
            if (ctx := getattr(self, "_ctx", None)) is not None:
                res = await AC(ctx())
            return res
        except BaseException as exc:
            await AC_exit(self, type(exc), exc, getattr(exc, "__traceback__", None))
            raise

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc: BaseException | None,
        tb: TracebackType | None,
    ) -> bool | None:
        return await AC_exit(self, exc_type, exc, tb)

    async def setup(self):
        """
        Basic async setup method. You may use AC_exit.

        Call the superclass first, when overriding.
        """

    async def teardown(self):
        """
        Object destructor.

        Should not fail when called with a partially-created object.
        """


class BaseConn(Base):
    """
    This is the MoaT stream base class for "something connected that talks".

    This class *must* be used as an async context manager.

    Usage:

    Override :py:meth:`BaseConn.stream` to create the data link.
    Use :py:func:`moat.lib.micro.AC_use` to
    call an async context manager or to register a destructor.

    Augment `setup` or `teardown` to add non-stream related features.
    """

    s: Any = None

    async def setup(self):
        """
        Object construction.

        By default, assigns the result of calling :py:meth:`BaseConn.stream` to the attribute
        ``s``.
        """
        if self.s is not None:
            raise RuntimeError("Busy!")

        self.s = await self.stream()

    async def teardown(self):
        """
        Object destructor.

        Should not fail when called with a partially-created object.
        """
        self.s = None

    async def stream(self):
        """
        Data stream setup.

        You need to use `moat.lib.micro.AC_use` for setting up an async context
        or to register a cleanup handler.
        """
        raise NotImplementedError(f"'stream' in {self!r}")


class BaseMsg(BaseConn):
    """
    A stream base module for messages. May not be useful.

    Implement send/recv.
    """

    async def send(self, m: Any) -> Any:
        """
        Send a message.
        """
        raise NotImplementedError(f"'send' in {self!r}")

    async def recv(self) -> Any:
        """
        Receive a message.
        """
        raise NotImplementedError(f"'recv' in {self!r}")


class BaseBlk(BaseConn):
    """
    A stream base module for bytestrings. May not be useful.

    Implement snd/rcv.
    """

    async def snd(self, m: Buffer | bytes) -> None:
        """
        Send a block of bytes.
        """
        raise NotImplementedError(f"'send' in {self!r}")

    async def rcv(self) -> Buffer | bytes:
        """
        Receive a block of bytes.
        """
        raise NotImplementedError(f"'recv' in {self!r}")


class BaseBuf(BaseConn):
    """
    A stream base module for bytestreams.

    Implement rd/wr.
    """

    async def rd(self, buf: MutBuffer) -> int:
        """
        Read some bytes.

        @buf is a bytearray to read data into. The return value is the
        number of bytes filled.

        This method never returns zero. End-of-file raises `EOFError`.
        """
        raise NotImplementedError(f"'rd' in {self!r}")

    async def wr(self, data: Buffer) -> int:
        """
        Write some bytes.
        """
        raise NotImplementedError(f"'wr' in {self!r}")


class StackedConn(BaseConn):
    """
    Base class for connection stacking.

    Connection stacks have a lower layer. Our `stream` methods uses it as an
    async context manager to create connection from it, using its

    Args:
        link(BaseConn): The lower layer to run on top of.
        cfg: Config data
    """

    link: BaseConn | None = None

    def __init__(self, link: BaseConn, cfg):
        super().__init__(cfg=cfg)
        self.link = link

    def wrap(self):  # noqa:D102
        if self.link is None:
            raise RuntimeError("No link")
        return self.link.wrap()

    async def stream(self) -> BaseConn:
        """
        Generate the low-level connection this module uses.

        By default, returns the linked stream's async context.
        """
        if self.link is None:
            raise RuntimeError("No link")
        return await AC_use(self, self.link)


class StackedMsg(StackedConn, BaseMsg):
    """
    A no-op stack module for messages. Override to implement interesting features.

    Use the attribute "s" to store the linked stream's context.

    Args:
        link(BaseMsg): The lower layer to run on top of.
        cfg: Config data
    """

    async def send(self, m: Any) -> Any:
        "Send. Transmits a structured message"
        return await self.s.send(m)

    async def recv(self) -> Any:
        "Receive. Returns a message."
        return await self.s.recv()

    async def cwr(self, buf: Buffer) -> None:
        "Console Send. Returns when the buffer is transmitted."
        await self.s.cwr(buf)

    async def crd(self, buf: MutBuffer) -> int:
        "Console Receive. Returns data by reading into a buffer."
        return await self.s.crd(buf)


class StackedBuf(StackedConn, BaseBuf):
    """
    A no-op stack module for byte steams. Override to implement interesting features.

    Use the attribute "s" to store the linked stream's context.

    Args:
        link(BaseBuf): The lower layer to run on top of.
        cfg: Config data
    """

    async def wr(self, data: Buffer) -> int:
        "Send. Returns when the buffer is transmitted."
        return await self.s.wr(data)

    async def rd(self, buf: MutBuffer) -> int:
        "Receive. Returns data by reading into a buffer."
        return await self.s.rd(buf)


class StackedBlk(StackedConn, BaseBlk):
    """
    A no-op stack module for bytestrings. Override to implement interesting features.

    Use the attribute "s" to store the linked stream's context.

    Args:
        link(BaseBlk): The lower layer to run on top of.
        cfg: Config data
    """

    cwr = StackedMsg.cwr
    crd = StackedMsg.crd

    async def snd(self, m: Buffer | bytes) -> None:
        "Send. Transmits a structured message"
        await self.s.snd(m)

    async def rcv(self) -> Buffer | bytes:
        "Receive. Returns a message."
        return await self.s.rcv()
