"""Unit tests for KV encryption support in KVBatcher and CachedKvClient."""

import os
from unittest.mock import MagicMock

from zg_storage.core.data import BytesDataSource
from zg_storage.core.encrypted_data import EncryptedDataSource
from zg_storage.kv import KVBatcher
from zg_storage.kv.cached_client import CachedKvClient
from zg_storage.transfer.encryption import decrypt_file

STREAM_ID = b"\x01" * 32


# ---------------------------------------------------------------------------
# KVBatcher encryption
# ---------------------------------------------------------------------------


def test_batcher_stores_encryption_key():
    key = os.urandom(32)
    batcher = KVBatcher.new(encryption_key=key)
    assert batcher._encryption_key == key


def test_batcher_no_encryption_key_by_default():
    batcher = KVBatcher.new()
    assert batcher._encryption_key is None


def test_batcher_exec_encrypts_payload():
    """exec() wraps data with EncryptedDataSource when encryption_key is set."""
    key = os.urandom(32)
    batcher = KVBatcher.new(encryption_key=key)
    batcher.set(STREAM_ID, b"k1", b"v1")

    # Build expected unencrypted payload
    plain_batcher = KVBatcher.new()
    plain_batcher.set(STREAM_ID, b"k1", b"v1")
    expected_encoded = plain_batcher.builder.build().encode()

    mock_indexer = MagicMock()
    mock_indexer.upload.return_value = ("0xtxhash", "0xroothash", 42)

    batcher.exec(mock_indexer)

    mock_indexer.upload.assert_called_once()
    call_kwargs = mock_indexer.upload.call_args
    source = call_kwargs.kwargs.get("file_path") or call_kwargs[1].get("file_path")
    if source is None:
        source = call_kwargs[0][0] if call_kwargs[0] else None

    assert isinstance(source, EncryptedDataSource)

    encrypted_data = source.read_at(0, source.size())
    decrypted = decrypt_file(key, encrypted_data)
    assert decrypted == expected_encoded


def test_batcher_exec_no_encryption_without_key():
    """exec() uses plain BytesDataSource when no encryption_key is set."""
    batcher = KVBatcher.new()
    batcher.set(STREAM_ID, b"k1", b"v1")

    mock_indexer = MagicMock()
    mock_indexer.upload.return_value = ("0xtxhash", "0xroothash", 42)

    batcher.exec(mock_indexer)

    mock_indexer.upload.assert_called_once()
    call_kwargs = mock_indexer.upload.call_args
    source = call_kwargs.kwargs.get("file_path") or call_kwargs[1].get("file_path")
    if source is None:
        source = call_kwargs[0][0] if call_kwargs[0] else None

    assert isinstance(source, BytesDataSource)


# ---------------------------------------------------------------------------
# CachedKvClient encryption
# ---------------------------------------------------------------------------


def _make_cached_client(mock_kv=None, mock_indexer=None, **kwargs):
    """Create a CachedKvClient with mocked internals (bypasses __init__ URL params)."""
    mock_kv = mock_kv or MagicMock()
    mock_indexer = mock_indexer or MagicMock()
    client = CachedKvClient.__new__(CachedKvClient)
    client._init_common(kv_client=mock_kv, indexer=mock_indexer, **kwargs)
    return client


def test_cached_client_stores_encryption_key():
    key = os.urandom(32)
    client = _make_cached_client(encryption_key=key)
    assert client._encryption_key == key
    client.close()


def test_cached_client_no_encryption_key_by_default():
    client = _make_cached_client()
    assert client._encryption_key is None
    client.close()


def test_cached_client_do_upload_encrypts_payload():
    """_do_upload wraps data with EncryptedDataSource when encryption_key is set."""
    key = os.urandom(32)
    mock_indexer = MagicMock()
    mock_indexer.upload.return_value = ("0xtxhash", "0xroothash", 42)

    client = _make_cached_client(mock_indexer=mock_indexer, encryption_key=key)

    # Build expected unencrypted payload
    plain_batcher = KVBatcher.new()
    plain_batcher.set(STREAM_ID, b"ck1", b"cv1")
    expected_encoded = plain_batcher.builder.build(sorted_items=True).encode()

    client.set(STREAM_ID, b"ck1", b"cv1")
    client.commit()
    client.flush(timeout=10)

    mock_indexer.upload.assert_called_once()
    call_kwargs = mock_indexer.upload.call_args
    source = call_kwargs.kwargs.get("file_path") or call_kwargs[1].get("file_path")
    if source is None:
        source = call_kwargs[0][0] if call_kwargs[0] else None

    assert isinstance(source, EncryptedDataSource)

    encrypted_data = source.read_at(0, source.size())
    decrypted = decrypt_file(key, encrypted_data)
    assert decrypted == expected_encoded

    client.close()


def test_cached_client_do_upload_no_encryption_without_key():
    """_do_upload uses plain BytesDataSource when no encryption_key is set."""
    mock_indexer = MagicMock()
    mock_indexer.upload.return_value = ("0xtxhash", "0xroothash", 42)

    client = _make_cached_client(mock_indexer=mock_indexer)

    client.set(STREAM_ID, b"ck1", b"cv1")
    client.commit()
    client.flush(timeout=10)

    mock_indexer.upload.assert_called_once()
    call_kwargs = mock_indexer.upload.call_args
    source = call_kwargs.kwargs.get("file_path") or call_kwargs[1].get("file_path")
    if source is None:
        source = call_kwargs[0][0] if call_kwargs[0] else None

    assert isinstance(source, BytesDataSource)

    client.close()
