"""zg_storage.kv.client module."""

from __future__ import annotations

import base64
from typing import Iterable, Optional

from ..node import ZgsKvClient
from ..types import KeyValue, Value
from ..utils import encode_bytes_base64, hex_to_bytes
from .builder import MAX_VERSION

MAX_QUERY_SIZE = 1024 * 256


class KvIterator:
    """Iterator over key-value pairs in a KV stream.

    Mirrors the iterator pattern from the Go SDK.
    Use KvClient.new_iterator() to create instances."""

    def __init__(self, kv_client: KvClient, stream_id: str, version: Optional[int] = None) -> None:
        self._kv_client = kv_client
        self._rpc = kv_client._client
        self._stream_id = stream_id
        self._version = version if version is not None else MAX_VERSION
        self._current: Optional[KeyValue] = None

    def _move(self, kv: Optional[KeyValue]) -> None:
        """Fetch the full value for a navigated key and update current pair."""
        if kv is None:
            self._current = None
            return
        val = self._kv_client.get_value(
            self._stream_id, base64.b64decode(kv.key), version=self._version
        )
        self._current = KeyValue(
            version=val.version,
            key=kv.key,
            data=val.data,
            size=val.size,
        )

    def valid(self) -> bool:
        """Return True if the iterator points to a valid key-value pair."""
        return self._current is not None

    @property
    def key(self) -> bytes:
        """Return the current key as bytes. Raises if not valid."""
        if self._current is None:
            raise StopIteration("iterator is not valid")
        return base64.b64decode(self._current.key)

    @property
    def data(self) -> bytes:
        """Return the current value data as bytes. Raises if not valid."""
        if self._current is None:
            raise StopIteration("iterator is not valid")
        return base64.b64decode(self._current.data)

    @property
    def size(self) -> int:
        """Return the total size of the current value. Raises if not valid."""
        if self._current is None:
            raise StopIteration("iterator is not valid")
        return self._current.size

    @property
    def current_pair(self) -> Optional[KeyValue]:
        """Return the current KeyValue or None."""
        return self._current

    def seek_to_first(self) -> None:
        """Position the iterator at the first key in the stream."""
        kv = self._rpc.get_first(self._stream_id, 0, 0, version=self._version)
        self._move(kv)

    def seek_to_last(self) -> None:
        """Position the iterator at the last key in the stream."""
        kv = self._rpc.get_last(self._stream_id, 0, 0, version=self._version)
        self._move(kv)

    def seek_before(self, key: str | bytes | bytearray | memoryview) -> None:
        """Seek to the position before the given key (inclusive)."""
        key_param = KvClient._normalize_key(key)
        kv = self._rpc.get_prev(self._stream_id, key_param, 0, 0, True, version=self._version)
        self._move(kv)

    def seek_after(self, key: str | bytes | bytearray | memoryview) -> None:
        """Seek to the position after the given key (inclusive)."""
        key_param = KvClient._normalize_key(key)
        kv = self._rpc.get_next(self._stream_id, key_param, 0, 0, True, version=self._version)
        self._move(kv)

    def next(self) -> None:
        """Advance the iterator to the next key."""
        if self._current is None:
            return
        kv = self._rpc.get_next(
            self._stream_id,
            self._current.key,
            0,
            0,
            False,
            version=self._version,
        )
        self._move(kv)

    def prev(self) -> None:
        """Move the iterator to the previous key."""
        if self._current is None:
            return
        kv = self._rpc.get_prev(
            self._stream_id,
            self._current.key,
            0,
            0,
            False,
            version=self._version,
        )
        self._move(kv)


class KvClient:
    """High-level KV client with paging logic.

    Purpose:
        Fetch full values and decode data across multiple pages."""

    def __init__(self, kv_rpc_url: str, timeout: float = 30.0) -> None:
        self._client = ZgsKvClient(kv_rpc_url, timeout=timeout)

    @staticmethod
    def _normalize_key(key: str | bytes | bytearray | memoryview) -> str:
        """Normalize a key to base64 encoding for RPC calls."""
        if isinstance(key, (bytes, bytearray, memoryview)):
            return encode_bytes_base64(bytes(key))
        if key.startswith("0x"):
            return encode_bytes_base64(hex_to_bytes(key))
        return encode_bytes_base64(key.encode("utf-8"))

    def new_iterator(self, stream_id: str, version: Optional[int] = None) -> KvIterator:
        """Create a new iterator for the given stream."""
        return KvIterator(self, stream_id, version=version)

    def get_value(
        self,
        stream_id: str,
        key: str | bytes | bytearray | memoryview,
        version: Optional[int] = None,
    ) -> Value:
        """Fetch the full value for a key, paging until complete."""
        if version is None:
            version = MAX_VERSION

        key_param = self._normalize_key(key)
        data = bytearray()
        size = 0
        current_version = version

        while True:
            seg = self._client.get_value(
                stream_id,
                key_param,
                len(data),
                MAX_QUERY_SIZE,
                version=current_version,
            )
            seg_bytes = base64.b64decode(seg.data) if seg.data else b""

            if current_version == MAX_VERSION:
                current_version = seg.version
            elif current_version != seg.version:
                current_version = seg.version
                data = bytearray()

            size = seg.size
            data.extend(seg_bytes)

            if len(data) >= size:
                break

        return Value(version=current_version, data=encode_bytes_base64(bytes(data)), size=size)

    def get_value_bytes(
        self,
        stream_id: str,
        key: str | bytes | bytearray | memoryview,
        version: Optional[int] = None,
    ) -> bytes:
        """Fetch the value bytes for a key."""
        val = self.get_value(stream_id, key, version=version)
        return base64.b64decode(val.data)

    def get_values(
        self,
        stream_id: str,
        keys: Iterable[str | bytes | bytearray | memoryview],
        version: Optional[int] = None,
    ) -> dict[str, str]:
        """Fetch multiple values and return decoded strings by key."""
        result: dict[str, str] = {}
        for key in keys:
            data = self.get_value_bytes(stream_id, key, version=version)
            if isinstance(key, str):
                key_label = key
            else:
                key_label = encode_bytes_base64(bytes(key))
            result[key_label] = data.decode("utf-8", errors="replace")
        return result
