import base64

import pytest

from zg_storage.kv import KVBatcher

STREAM_ID = b"\x01" * 32
STREAM_ID_HEX = "0x" + "01" * 32


def test_get_returns_local_write():
    """get() returns locally set value without any RPC call."""
    batcher = KVBatcher.new()
    batcher.set(STREAM_ID, b"key1", b"value1")

    val = batcher.get(STREAM_ID, b"key1")

    assert base64.b64decode(val.data) == b"value1"
    assert val.size == 6
    assert val.version == batcher.builder.version


def test_get_local_with_hex_stream_id():
    """get() works when stream_id is passed as hex string."""
    batcher = KVBatcher.new()
    batcher.set(STREAM_ID_HEX, b"mykey", b"myval")

    val = batcher.get(STREAM_ID_HEX, b"mykey")

    assert base64.b64decode(val.data) == b"myval"
    assert val.size == 5


def test_get_local_with_hex_key():
    """get() works when key is passed as hex string."""
    batcher = KVBatcher.new()
    key_bytes = b"\xaa\xbb\xcc"
    key_hex = "0x" + key_bytes.hex()
    batcher.set(STREAM_ID, key_hex, b"hexval")

    val = batcher.get(STREAM_ID, key_hex)

    assert base64.b64decode(val.data) == b"hexval"


def test_get_missing_key_no_client_raises():
    """get() raises KeyError when key is not local and no KV client is set."""
    batcher = KVBatcher.new()

    with pytest.raises(KeyError, match="key not found locally"):
        batcher.get(STREAM_ID, b"missing")


def test_get_after_overwrite_returns_latest():
    """get() returns the latest value after overwriting a key."""
    batcher = KVBatcher.new()
    batcher.set(STREAM_ID, b"key1", b"first")
    batcher.set(STREAM_ID, b"key1", b"second")

    val = batcher.get(STREAM_ID, b"key1")

    assert base64.b64decode(val.data) == b"second"
    assert val.size == 6


def test_get_multiple_streams():
    """get() correctly isolates values across different stream ids."""
    stream_a = b"\x01" * 32
    stream_b = b"\x02" * 32

    batcher = KVBatcher.new()
    batcher.set(stream_a, b"key", b"val_a")
    batcher.set(stream_b, b"key", b"val_b")

    val_a = batcher.get(stream_a, b"key")
    val_b = batcher.get(stream_b, b"key")

    assert base64.b64decode(val_a.data) == b"val_a"
    assert base64.b64decode(val_b.data) == b"val_b"
