from dataclasses import InitVar, dataclass, field
from io import BytesIO

from ordeq import IO


@dataclass(frozen=True, eq=False)
class BytesBuffer(IO[bytes]):
    """IO that uses an in-memory bytes buffer to load and save data. Useful
    for buffering data across nodes without writing to disk.

    Example:

        >>> from ordeq_common import StringBuffer
        >>> buffer = BytesBuffer()
        >>> buffer.load()
        b''

    The buffer is initially empty, unless provided with initial data:

        >>> buffer = BytesBuffer(b"Initial data")
        >>> buffer.load()
        b'Initial data'

    Saving to the buffer appends data to the existing content:

        >>> buffer.save(b"New data")
        >>> buffer.load()
        b'Initial dataNew data'

    Example in a node:

        >>> from ordeq_args import CommandLineArg
        >>> from ordeq_common import StringBuffer, Static
        >>> from ordeq import node, run
        >>> result = BytesBuffer(b"Greeting")
        >>> @node(
        ...     inputs=[BytesBuffer(b"Hello"), Static(b"you")],
        ...     outputs=result
        ... )
        ... def greet(greeting: bytes, name: bytes) -> bytes:
        ...     return greeting + b', ' + name + b'!'
        >>> run(greet).get(result)
        b'Hello, you!'

    """

    _buffer: BytesIO = field(default_factory=BytesIO, init=False)
    value: InitVar[bytes | None] = None

    def __post_init__(self, value: bytes | None):
        if value is not None:
            self._buffer.write(value)

    def load(self) -> bytes:
        return self._buffer.getvalue()

    def save(self, data: bytes) -> None:
        self._buffer.write(data)
