"""zg_storage.kv.cached_client module.

Write-behind KV client with async upload queue and LRU read cache.

The upload queue and read cache are decoupled:
  - Queue: bounded FIFO for uploads.  Items removed on success.
  - Cache: LRU dict for O(1) reads.  Eviction requires KV node commit.
"""

from __future__ import annotations

import base64
import logging
import queue
import threading
from collections import OrderedDict
from dataclasses import dataclass
from typing import Callable, 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 .client import KvClient
from .types import _ensure_bytes, _stream_id_bytes, create_tags

_SENTINEL = object()

_logger = logging.getLogger("zg_storage")


@dataclass
class CacheEntry:
    """A cached value with metadata for safe eviction."""

    value: bytes
    batch_id: int
    tx_seq: Optional[int] = None  # None until batch uploads successfully


@dataclass
class _PendingBatch:
    """Internal batch queued for upload."""

    batch_id: int
    builder: StreamDataBuilder
    keys: set[tuple[bytes, bytes]]  # set of (sid, key) in this batch


class CachedKvClient:
    """KV client with async upload queue and LRU read cache.

    Writes accumulate in a local buffer, are committed to a bounded upload
    queue, and uploaded in FIFO order by a background worker.  Reads are
    served from: current writes -> LRU cache -> KV node.

    The upload queue and read cache are decoupled:
      - Queue items are removed as soon as upload succeeds.
      - Cache entries persist until LRU eviction, which requires the
        entry's data to be committed on the KV node.

    Multi-stream: stream_id is passed per set()/get() call, like KVBatcher.
    """

    def __init__(
        self,
        kv_url: str,
        indexer_url: str,
        evm_client,
        flow_address: str,
        max_queue_size: int = 10,
        max_cache_entries: int = 10000,
        upload_option: Optional[UploadOption] = None,
        version: Optional[int] = None,
        on_commit_error: Optional[Callable[[Exception], None]] = None,
        encryption_key: Optional[bytes] = None,
    ) -> None:
        kv_client = KvClient(kv_rpc_url=kv_url)
        indexer = IndexerClient(
            url=indexer_url,
            evm_client=evm_client,
            flow_address=flow_address,
        )
        self._init_common(
            kv_client=kv_client,
            indexer=indexer,
            max_queue_size=max_queue_size,
            max_cache_entries=max_cache_entries,
            upload_option=upload_option,
            version=version,
            on_commit_error=on_commit_error,
            encryption_key=encryption_key,
        )

    def _init_common(
        self,
        kv_client: KvClient,
        indexer: IndexerClient,
        max_queue_size: int = 10,
        max_cache_entries: int = 10000,
        upload_option: Optional[UploadOption] = None,
        version: Optional[int] = None,
        on_commit_error: Optional[Callable[[Exception], None]] = None,
        encryption_key: Optional[bytes] = None,
    ) -> None:
        self._kv_client = kv_client
        self._indexer = indexer
        self._max_cache_entries = max_cache_entries
        self._upload_option = upload_option
        self._version = version if version is not None else MAX_VERSION
        self._on_commit_error = on_commit_error
        self._encryption_key = encryption_key

        # Current uncommitted batch
        self._current_builder = StreamDataBuilder(version=self._version)
        self._current_writes: dict[tuple[bytes, bytes], bytes] = {}

        # LRU read cache (committed data only): (stream_id, key) -> CacheEntry
        self._cache: OrderedDict[tuple[bytes, bytes], CacheEntry] = OrderedDict()
        self._batch_counter = 0

        # Finalization tracking
        self._pending_tx_seqs: set[int] = set()
        self._finalized_tx_seqs: set[int] = set()

        # Threading
        self._lock = threading.Lock()
        self._upload_queue: queue.Queue = queue.Queue(maxsize=max_queue_size)
        self._flush_condition = threading.Condition()
        self._closed = False
        self._failed: Optional[Exception] = None

        self._worker = threading.Thread(
            target=self._upload_worker,
            name="CachedKvClient-worker",
            daemon=True,
        )
        self._worker.start()

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def set(
        self,
        stream_id: bytes | bytearray | memoryview | str,
        key: bytes | bytearray | memoryview | str,
        value: bytes | bytearray | memoryview,
    ) -> CachedKvClient:
        """Add a key-value write to the current uncommitted batch."""
        sid = _stream_id_bytes(stream_id)
        k = _ensure_bytes(key)
        v = bytes(value) if not isinstance(value, bytes) else value
        cache_key = (sid, k)

        with self._lock:
            if self._closed:
                raise RuntimeError("CachedKvClient is closed")
            self._current_builder.set(sid, k, v)
            self._current_writes[cache_key] = v
        return self

    def get(
        self,
        stream_id: bytes | bytearray | memoryview | str,
        key: bytes | bytearray | memoryview | str,
    ) -> Value:
        """Read a value: current writes -> LRU cache -> KV node."""
        sid = _stream_id_bytes(stream_id)
        k = _ensure_bytes(key)
        cache_key = (sid, k)

        with self._lock:
            # Layer 1: current uncommitted writes
            if cache_key in self._current_writes:
                data = self._current_writes[cache_key]
                return Value(
                    version=self._version,
                    data=encode_bytes_base64(data),
                    size=len(data),
                )

            # Layer 2: LRU cache (committed data)
            if cache_key in self._cache:
                self._cache.move_to_end(cache_key)
                entry = self._cache[cache_key]
                return Value(
                    version=self._version,
                    data=encode_bytes_base64(entry.value),
                    size=len(entry.value),
                )

        # Layer 3: remote KV node
        stream_id_hex = "0x" + sid.hex()
        return self._kv_client.get_value(stream_id_hex, k, version=self._version)

    def get_bytes(
        self,
        stream_id: bytes | bytearray | memoryview | str,
        key: bytes | bytearray | memoryview | str,
    ) -> bytes:
        """Convenience: get value as raw bytes."""
        val = self.get(stream_id, key)
        return base64.b64decode(val.data)

    def commit(self) -> None:
        """Snapshot current writes and queue for async upload.

        Moves current writes into the LRU cache and queues the batch
        for background upload.  Blocks if the upload queue is full.
        """
        with self._lock:
            if self._closed:
                raise RuntimeError("CachedKvClient is closed")
            if self._failed is not None:
                raise RuntimeError(f"CachedKvClient failed: {self._failed}") from self._failed
            if not self._current_writes:
                return

            self._batch_counter += 1
            batch_id = self._batch_counter

            batch = _PendingBatch(
                batch_id=batch_id,
                builder=self._current_builder,
                keys=set(self._current_writes.keys()),
            )

            # Move current writes into the LRU cache
            for cache_key, value in self._current_writes.items():
                self._cache[cache_key] = CacheEntry(value=value, batch_id=batch_id)
                self._cache.move_to_end(cache_key)

            self._try_evict_cache()

            # Reset current batch
            self._current_builder = StreamDataBuilder(version=self._version)
            self._current_writes = {}

        # Queue for background worker (blocks if queue is full)
        self._upload_queue.put(batch)

    def flush(self, timeout: Optional[float] = None) -> None:
        """Block until all queued uploads are processed.

        Args:
            timeout: Maximum seconds to wait. None means wait indefinitely.

        Raises:
            TimeoutError: If timeout is exceeded.
        """
        if timeout is None:
            self._upload_queue.join()
            return

        import time

        deadline = time.monotonic() + timeout
        with self._flush_condition:
            while self._upload_queue.unfinished_tasks:
                remaining = deadline - time.monotonic()
                if remaining <= 0:
                    raise TimeoutError("flush timed out")
                self._flush_condition.wait(timeout=remaining)

    def close(self) -> None:
        """Commit any pending writes, flush, and shut down the worker."""
        with self._lock:
            if self._closed:
                return
            has_pending = bool(self._current_writes) and self._failed is None

        if has_pending:
            self.commit()

        with self._lock:
            self._closed = True

        self._upload_queue.join()
        self._upload_queue.put(_SENTINEL)
        self._worker.join(timeout=10.0)

    def reset(self) -> None:
        """Clear failed state so the client can accept new commits.

        After an upload failure the client enters a failed state where
        ``commit()`` raises.  Call ``reset()`` to clear the error, then
        re-``set()`` any keys that were in failed or drained batches and
        ``commit()`` again.

        The LRU cache is preserved so ``get()`` continues to return the
        latest local values.  Any writes accumulated via ``set()`` before
        or after this call remain in the uncommitted buffer.
        """
        with self._lock:
            if self._closed:
                raise RuntimeError("CachedKvClient is closed")
            self._failed = None

    @property
    def failed(self) -> bool:
        """Whether the client is in a failed state."""
        with self._lock:
            return self._failed is not None

    @property
    def cache_size(self) -> int:
        """Current number of entries in the read cache."""
        with self._lock:
            return len(self._cache)

    def __enter__(self) -> CachedKvClient:
        return self

    def __exit__(self, exc_type, exc, tb) -> None:
        self.close()

    # ------------------------------------------------------------------
    # Internal
    # ------------------------------------------------------------------

    def _try_evict_cache(self) -> None:
        """Best-effort LRU eviction.  Called with self._lock held.

        Only evicts entries whose tx_seq is known committed on the KV node
        (present in self._finalized_tx_seqs).  No RPC calls here.
        """
        while len(self._cache) > self._max_cache_entries:
            evicted = False
            # Iterate from oldest (front of OrderedDict)
            for cache_key in list(self._cache):
                entry = self._cache[cache_key]
                if entry.tx_seq is not None and entry.tx_seq in self._finalized_tx_seqs:
                    del self._cache[cache_key]
                    evicted = True
                    break
            if not evicted:
                break  # nothing evictable, allow cache to grow

    def _probe_finality(self) -> None:
        """Check committed status for recently uploaded tx_seqs.

        Called by the worker thread after each upload.  RPC calls are made
        outside ``self._lock``.
        """
        with self._lock:
            to_check = set(self._pending_tx_seqs)

        newly_finalized: set[int] = set()
        for tx_seq in to_check:
            try:
                result = self._kv_client._client.get_transaction_result(tx_seq)
                if result == "Commit":
                    newly_finalized.add(tx_seq)
            except Exception:
                _logger.debug("Failed to check tx_seq=%s", tx_seq, exc_info=True)

        if newly_finalized:
            with self._lock:
                self._finalized_tx_seqs.update(newly_finalized)
                self._pending_tx_seqs -= newly_finalized
                self._try_evict_cache()

    def _drain_queue(self) -> None:
        """Discard all remaining batches in the queue."""
        while True:
            try:
                self._upload_queue.get_nowait()
                self._upload_queue.task_done()
            except queue.Empty:
                break

    def _upload_worker(self) -> None:
        """Background thread: upload committed batches sequentially."""
        while True:
            item = self._upload_queue.get()
            if item is _SENTINEL:
                self._upload_queue.task_done()
                break

            batch: _PendingBatch = item
            try:
                self._do_upload(batch)
            except Exception as exc:
                _logger.error("CachedKvClient upload failed: %s", exc)
                with self._lock:
                    self._failed = exc
                self._drain_queue()
                if self._on_commit_error:
                    try:
                        self._on_commit_error(exc)
                    except Exception:
                        pass
            finally:
                self._upload_queue.task_done()
                self._probe_finality()
                with self._flush_condition:
                    self._flush_condition.notify_all()

    def _do_upload(self, batch: _PendingBatch) -> None:
        """Execute the upload pipeline for a committed batch."""
        data = batch.builder.build(sorted_items=True)
        encoded = data.encode()
        tags = create_tags(batch.builder.stream_ids(), sorted_ids=True)

        opt = self._upload_option.model_copy() if self._upload_option else UploadOption()
        opt.tags = tags
        source = BytesDataSource(encoded)
        if self._encryption_key is not None:
            source = EncryptedDataSource(source, self._encryption_key)

        tx_hash, root_hash, tx_seq = self._indexer.upload(file_path=source, tags=tags, option=opt)

        # Stamp tx_seq on cache entries that still belong to this batch
        with self._lock:
            for cache_key in batch.keys:
                entry = self._cache.get(cache_key)
                if entry is not None and entry.batch_id == batch.batch_id:
                    entry.tx_seq = tx_seq
            self._pending_tx_seqs.add(tx_seq)

        _logger.info(
            "CachedKvClient upload completed tx=%s root=%s tx_seq=%s",
            tx_hash,
            root_hash,
            tx_seq,
        )
