"""
Stream link-up support for MoaT commands
"""

from __future__ import annotations

from moat.lib.micro import ACM, AC_exit, L, Lock, TaskGroup
from moat.lib.rpc import BaseCmd
from moat.lib.stream import BaseBlk, BaseBuf, BaseMsg

__all__ = ["BaseCmdBBM"]

# Typing
from typing import TYPE_CHECKING, cast  # isort:skip

if TYPE_CHECKING:
    from moat.lib.stream.base import Buffer, MutBuffer

    from collections.abc import Awaitable
    from typing import Protocol

    class _ConsoleProto(Protocol):
        async def crd(self, buf: MutBuffer) -> int: ...

        async def cwr(self, buf: Buffer) -> None: ...


class BaseCmdBBM(BaseCmd):
    """
    This is a command handler that connects MoaT's Cmd tree to a
    `~moat.lib.stream.BaseBuf`, `~moat.lib.stream.BaseBlk` or
    `~moat.lib.stream.BaseMsg` instance.

    Override :meth:`stream` to return that instance. Use
    :func:`moat.lib.micro.AC_use` if it's a context manager.

    This is a single class that adapts `~moat.lib.stream.BaseMsg`,
    `~moat.lib.stream.BaseBlk`, and `~moat.lib.stream.BaseBuf` streams.

    The difference between this and a
    :moat.lib.rpc.cmd.msg:`BaseCmdMsg`-derived class is that this
    class exposes commands that directly access the underlying stream
    (of whatever type).

    In contrast, :moat.lib.rpc.cmd.msg:`BaseCmdMsg` encapsulates
    arbitrary commands and requires a
    :moat.lib.rpc.cmd.msg:`BaseCmdMsg` handler on the other side to
    talk to.

    This class cannot wrap a pre-existing stream, by design.
    """

    s = None

    def __init__(self, cfg):
        super().__init__(cfg)
        self.w_lock = Lock()

    async def stream(self) -> BaseMsg | BaseBlk | BaseBuf:
        "create the actual data stream. Override this!"
        raise NotImplementedError("setup", self.path)

    async def setup(self):  # noqa:D102
        await super().setup()
        self.s = await self.stream()

    async def run(self):  # noqa:D102
        ACM(self)
        try:
            await super().run()
        finally:
            self.s = None
            await AC_exit(self)

    # Buf: rd/wr = .rd/.wr

    doc_rd = dict(_d="read bytestream", _0="int:len (64)")

    async def cmd_rd(self, n=64) -> Buffer:
        """read some data"""
        if L:
            await self.wait_ready()
        b = bytearray(n)
        if self.s is None:
            raise EOFError
        r = await cast(BaseBuf, self.s).rd(b)
        if r == n:
            return b
        elif r <= n >> 2:
            return bytes(b[:r])
        else:
            b = memoryview(b)
            return b[:r]

    doc_wr = dict(_d="write bytestream", _0="bytes:data")

    async def cmd_wr(self, b: Buffer):
        """write some data"""
        if L:
            await self.wait_ready()
        async with self.w_lock:
            if self.s is None:
                raise EOFError
            await cast(BaseBuf, self.s).wr(b)

    doc_rw = dict(
        _d="r/w byte stream", _0="int:rdbuflen (64)", _i="bytes:to send", _o="bytes:received"
    )

    async def stream_rw(self, msg):
        "read/write byte stream"
        n = msg.get(0, 64)
        async with msg.stream() as st, TaskGroup() as tg:

            @tg.start_soon
            async def rd():
                while True:
                    await st.send(await self.cmd_rd(n))

            async for m in st:
                await self.cmd_wr(m[0])
            tg.cancel()

    # Blk/Msg: Console crd/cwr = .crd/cwr

    doc_crd = dict(_d="read console", _0="int:len (64)", _r="bytes:data")

    async def cmd_crd(self, n=64) -> Buffer:
        """read some console data"""
        b = bytearray(n)
        if self.s is None:
            raise EOFError
        r = await cast("_ConsoleProto", self.s).crd(b)
        if r == n:
            return b
        elif r <= n >> 2:
            return bytes(b[:r])
        else:
            b = memoryview(b)
            return b[:r]

    doc_cwr = dict(_d="write console", _0="bytes:data")

    async def cmd_cwr(self, b: Buffer):
        """write some console data"""
        async with self.w_lock:
            if self.s is None:
                raise EOFError
            await cast("_ConsoleProto", self.s).cwr(b)

    doc_crw = dict(
        _d="r/w console stream", _0="int:rdbuflen (64)", _i="bytes:to send", _o="bytes:received"
    )

    async def stream_crw(self, msg):
        "read/write console stream"
        n = msg.get(0, 64)
        async with msg.stream() as st, TaskGroup() as tg:

            @tg.start_soon
            async def crd():
                while True:
                    await st.send(await self.cmd_crd(n))

            async for m in st:
                await self.cmd_cwr(m[0])
            tg.cancel()

    # Msg: s/r = .send/.recv

    doc_s = dict(_d="write message", _0="any:message")

    def cmd_s(self, m) -> Awaitable:  # pylint:disable=invalid-overridden-method
        """send a message"""
        if self.s is None:
            raise EOFError
        return cast(BaseMsg, self.s).send(m)

    doc_r = dict(_d="read message", _r="any:message")

    def cmd_r(self) -> Awaitable:  # pylint:disable=invalid-overridden-method
        """receive a message"""
        if self.s is None:
            raise EOFError
        return cast(BaseMsg, self.s).recv()

    doc_mrw = dict(_d="r/w msg stream", _i="any:message", _o="any:message")

    async def stream_mrw(self, msg):
        "read/write message stream"
        async with msg.stream() as st, TaskGroup() as tg:

            @tg.start_soon
            async def mrd():
                while True:
                    await st.send(await self.cmd_r())

            async for m in st:
                await self.cmd_s(m[0])
            tg.cancel()

    # Blk: sb/rb = .snd/.rcv

    doc_sb = dict(_d="write block", _0="bytes:encoded message")

    def cmd_sb(self, m) -> Awaitable:  # pylint:disable=invalid-overridden-method
        """send a binary message"""
        if self.s is None:
            raise EOFError
        return cast(BaseBlk, self.s).snd(m)

    doc_rb = dict(_d="read block", _r="any:encoded message")

    def cmd_rb(self) -> Awaitable:  # pylint:disable=invalid-overridden-method
        """receive a binary message"""
        if self.s is None:
            raise EOFError
        return cast(BaseBlk, self.s).rcv()

    doc_brw = dict(_d="r/w block stream", _i="bytes:message", _o="bytes:message")

    async def stream_brw(self, msg):
        "read/write block stream"
        async with msg.stream() as st, TaskGroup() as tg:

            @tg.start_soon
            async def brd():
                while True:
                    await st.send(await self.cmd_rb())

            async for m in st:
                await self.cmd_sb(m[0])
            tg.cancel()
