"""zg_storage.kv.builder module."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Dict, Set

from .types import (
    MAX_KEY_SIZE,
    MAX_SET_SIZE,
    KvError,
    StreamData,
    StreamRead,
    StreamWrite,
    _ensure_bytes,
    _stream_id_bytes,
)

MAX_VERSION = (1 << 64) - 1


@dataclass
class StreamDataBuilder:
    """Builder for KV stream payloads.

    Purpose:
        Collect read/write operations and produce serialized stream data."""

    version: int = MAX_VERSION
    _stream_ids: Set[bytes] = field(default_factory=set)
    _reads: Dict[bytes, Set[bytes]] = field(default_factory=dict)
    _writes: Dict[bytes, Dict[bytes, bytes]] = field(default_factory=dict)

    def set_version(self, version: int) -> StreamDataBuilder:
        """Override the batch version for subsequent writes."""
        self.version = version
        return self

    def watch(
        self,
        stream_id: bytes | bytearray | memoryview | str,
        key: bytes | bytearray | memoryview | str,
    ) -> StreamDataBuilder:
        """Add a read/watch entry for a stream key."""
        sid = _stream_id_bytes(stream_id)
        k = _ensure_bytes(key)
        self._reads.setdefault(sid, set()).add(k)
        return self

    def set(
        self,
        stream_id: bytes | bytearray | memoryview | str,
        key: bytes | bytearray | memoryview | str,
        data: bytes | bytearray | memoryview,
    ) -> StreamDataBuilder:
        """Add a write entry for a stream key."""
        sid = _stream_id_bytes(stream_id)
        k = _ensure_bytes(key)
        d = _ensure_bytes(data)
        self._stream_ids.add(sid)
        self._writes.setdefault(sid, {})[k] = d
        return self

    def build(self, sorted_items: bool = False) -> StreamData:
        """Construct the stream data payload."""
        reads: list[StreamRead] = []
        writes: list[StreamWrite] = []

        for sid, keys in self._reads.items():
            for k in keys:
                if len(k) == 0:
                    raise KvError("key is empty")
                if len(k) > MAX_KEY_SIZE:
                    raise KvError("key too large")
                reads.append(StreamRead(stream_id=sid, key=k))
                if len(reads) > MAX_SET_SIZE:
                    raise KvError("size too large")

        for sid, kvs in self._writes.items():
            for k, v in kvs.items():
                if len(k) == 0:
                    raise KvError("key is empty")
                if len(k) > MAX_KEY_SIZE:
                    raise KvError("key too large")
                writes.append(StreamWrite(stream_id=sid, key=k, data=v))
                if len(writes) > MAX_SET_SIZE:
                    raise KvError("size too large")

        if sorted_items:
            reads.sort(key=lambda r: (r.stream_id.hex(), r.key.hex()))
            writes.sort(key=lambda w: (w.stream_id.hex(), w.key.hex()))

        return StreamData(version=self.version, reads=reads, writes=writes, controls=[])

    def stream_ids(self) -> list[bytes]:
        return list(self._stream_ids)
