"""zg_storage.kv.types module."""

from __future__ import annotations

import hashlib
from dataclasses import dataclass
from typing import Optional

MAX_SET_SIZE = 1 << 16
MAX_KEY_SIZE = 1 << 24


class KvError(ValueError):
    """Pydantic model for RPC data.

    Purpose:
        Provide typed access to fields returned by RPC calls."""


def _ensure_bytes(data: bytes | bytearray | memoryview | str) -> bytes:
    """Convert bytes-like or hex string input to bytes."""
    if isinstance(data, (bytes, bytearray, memoryview)):
        return bytes(data)
    if isinstance(data, str):
        hex_str = data[2:] if data.startswith("0x") else data
        return bytes.fromhex(hex_str)
    raise TypeError("expected bytes or hex string")


def _stream_id_bytes(stream_id: bytes | bytearray | memoryview | str) -> bytes:
    """Validate and normalize a 32-byte stream id."""
    sid = _ensure_bytes(stream_id)
    if len(sid) != 32:
        raise KvError("stream_id must be 32 bytes")
    return sid


STREAM_DOMAIN = hashlib.sha256(b"STREAM").digest()


def create_tags(stream_ids: list[bytes], sorted_ids: bool = False) -> bytes:
    """Create flow tags for the provided stream ids."""
    ids = list(stream_ids)
    if sorted_ids:
        ids.sort(key=lambda x: x.hex())
    out = bytearray()
    out += STREAM_DOMAIN
    for sid in ids:
        out += sid
    return bytes(out)


def _encode_u32(value: int) -> bytes:
    return value.to_bytes(4, "big", signed=False)


def _encode_u64(value: int) -> bytes:
    return value.to_bytes(8, "big", signed=False)


def _encode_size24(size: int) -> bytes:
    """Encode a size into a 3-byte big-endian field."""
    if size == 0:
        raise KvError("key is empty")
    if size >= 1 << 24:
        raise KvError("key too large")
    return size.to_bytes(4, "big", signed=False)[1:]


@dataclass
class StreamRead:
    """Pydantic model for RPC data.

    Purpose:
        Provide typed access to fields returned by RPC calls."""

    stream_id: bytes
    key: bytes


@dataclass
class StreamWrite:
    """Pydantic model for RPC data.

    Purpose:
        Provide typed access to fields returned by RPC calls."""

    stream_id: bytes
    key: bytes
    data: bytes


@dataclass
class AccessControl:
    """Pydantic model for RPC data.

    Purpose:
        Provide typed access to fields returned by RPC calls."""

    type: int
    stream_id: bytes
    account: Optional[bytes] = None
    key: Optional[bytes] = None


@dataclass
class StreamData:
    """Pydantic model for RPC data.

    Purpose:
        Provide typed access to fields returned by RPC calls."""

    version: int
    reads: list[StreamRead]
    writes: list[StreamWrite]
    controls: list[AccessControl]

    def encode(self) -> bytes:
        """Serialize stream reads, writes, and controls into bytes."""
        encoded = bytearray()
        encoded += _encode_u64(self.version)

        encoded += _encode_u32(len(self.reads))
        for r in self.reads:
            encoded += r.stream_id
            encoded += _encode_size24(len(r.key))
            encoded += r.key

        encoded += _encode_u32(len(self.writes))
        for w in self.writes:
            encoded += w.stream_id
            encoded += _encode_size24(len(w.key))
            encoded += w.key
            encoded += _encode_u64(len(w.data))

        for w in self.writes:
            encoded += w.data

        encoded += _encode_u32(len(self.controls))
        for c in self.controls:
            encoded.append(c.type & 0xFF)
            encoded += c.stream_id
            if c.key is not None:
                encoded += _encode_size24(len(c.key))
                encoded += c.key
            if c.account is not None:
                if len(c.account) != 20:
                    raise KvError("account must be 20 bytes")
                encoded += c.account

        return bytes(encoded)
