"""zg_storage.kv.batcher module."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Optional

from ..core.data import BytesDataSource
from ..core.encrypted_data import EncryptedDataSource
from ..indexer import IndexerClient
from ..transfer import UploadOption
from ..types import Value
from ..utils import encode_bytes_base64
from .builder import MAX_VERSION, StreamDataBuilder
from .types import _ensure_bytes, _stream_id_bytes, create_tags

if TYPE_CHECKING:
    from .client import KvClient


@dataclass
class KVBatcher:
    """KV write helper that serializes and uploads in-memory data.

    Purpose:
        Build stream data and upload via the standard pipeline."""

    builder: StreamDataBuilder
    _kv_client: Optional[KvClient] = field(default=None, repr=False)
    _encryption_key: Optional[bytes] = field(default=None, repr=False)

    @classmethod
    def new(
        cls,
        version: Optional[int] = None,
        encryption_key: Optional[bytes] = None,
    ) -> KVBatcher:
        """Create a new KV batcher.

        Args:
            version: Key version for writes. Defaults to MAX_VERSION if None.
            encryption_key: Optional 32-byte AES-256 key. When provided,
                the KV payload is encrypted before upload.
        """
        if version is None:
            version = MAX_VERSION
        return cls(
            builder=StreamDataBuilder(version=version),
            _encryption_key=encryption_key,
        )

    def with_kv_client(self, kv_client: KvClient) -> KVBatcher:
        """Set a KV client to enable Get (read-your-own-writes with remote fallback)."""
        self._kv_client = kv_client
        return self

    def set(
        self,
        stream_id: bytes | bytearray | memoryview | str,
        key: bytes | bytearray | memoryview | str,
        value: bytes | bytearray | memoryview,
    ) -> KVBatcher:
        """Add a key-value write to the current batch."""
        self.builder.set(stream_id, key, value)
        return self

    def get(
        self,
        stream_id: bytes | bytearray | memoryview | str,
        key: bytes | bytearray | memoryview | str,
    ) -> Value:
        """Get a value by key. Checks local write cache first, falls back to KV node.

        Args:
            stream_id: The stream to read from.
            key: The key to look up.

        Returns:
            Value with version, base64-encoded data, and size.

        Raises:
            KeyError: If the key is not found locally and no KV client is configured.
        """
        sid = _stream_id_bytes(stream_id)
        k = _ensure_bytes(key)

        # Check local writes first
        local_writes = self.builder._writes.get(sid)
        if local_writes is not None and k in local_writes:
            data = local_writes[k]
            return Value(
                version=self.builder.version,
                data=encode_bytes_base64(data),
                size=len(data),
            )

        # Fall back to KV node
        if self._kv_client is None:
            raise KeyError("key not found locally and no KV client configured")
        sid_hex = "0x" + sid.hex()
        return self._kv_client.get_value(sid_hex, k, version=self.builder.version)

    def exec(
        self,
        indexer: IndexerClient,
        option: Optional[UploadOption] = None,
    ) -> tuple[str, str, int]:
        """Serialize the batch and upload via the indexer."""
        data = self.builder.build()
        encoded = data.encode()
        tags = create_tags(self.builder.stream_ids(), sorted_ids=True)

        opt = option or UploadOption()
        opt.tags = tags
        source = BytesDataSource(encoded)
        if self._encryption_key is not None:
            source = EncryptedDataSource(source, self._encryption_key)
        return indexer.upload(file_path=source, tags=tags, option=opt)
