"""
In memory cache of files sent to the server to avoid sending the same file
multiple times, using an LRU (Least Recently Used) eviction policy.
"""

import hashlib
import threading
from collections import OrderedDict
from concurrent.futures import Future
from typing import Generic, Optional, TypeVar

from weave.trace_server.trace_server_interface import FileCreateReq, FileCreateRes

# Define generic type variables
K = TypeVar("K")  # Key type
V = TypeVar("V")  # Value type


class ThreadSafeLRUCache(Generic[K, V]):
    """Thread-safe LRU cache implementation using OrderedDict."""

    def __init__(self, max_size: int = 1000) -> None:
        """Initialize an LRU cache with a maximum size.

        Args:
            max_size: Maximum number of items to store in the cache.
                     If set to 0 or negative, cache size is unlimited.
        """
        self._max_size = max(0, max_size)  # Ensure max_size is non-negative
        self._cache: OrderedDict[K, V] = OrderedDict()
        self._lock = threading.RLock()

    def put(self, key: K, value: V) -> None:
        """Thread-safe insertion/update with LRU tracking.

        If the key already exists, it will be moved to the end (most recently used position).
        If max_size is reached, the least recently used item will be evicted.
        """
        with self._lock:
            # If key exists, move it to end (most recently used position)
            if key in self._cache:
                self._cache.move_to_end(key)
                self._cache[key] = value
                return

            # If we've reached max size and max_size is not 0 (unlimited), remove oldest item
            if 0 < self._max_size <= len(self._cache):
                self._cache.popitem(
                    last=False
                )  # last=False removes the first item (least recently used)

            # Add new key-value pair (will be at the end - most recently used)
            self._cache[key] = value

    def get(self, key: K) -> Optional[V]:
        """Thread-safe retrieval with LRU tracking.

        If the key exists, it will be moved to the end (most recently used position).
        """
        with self._lock:
            if key not in self._cache:
                return None

            # Move the key to the end (most recently used position)
            self._cache.move_to_end(key)
            return self._cache[key]

    def delete(self, key: K) -> None:
        """Thread-safe deletion."""
        with self._lock:
            if key in self._cache:
                del self._cache[key]

    def contains(self, key: K) -> bool:
        """Check if key exists in a thread-safe way."""
        with self._lock:
            return key in self._cache

    def clear(self) -> None:
        """Clear all items from the cache."""
        with self._lock:
            self._cache.clear()

    def size(self) -> int:
        """Return the current number of items in the cache."""
        with self._lock:
            return len(self._cache)

    @property
    def max_size(self) -> int:
        """Return the maximum size of the cache."""
        return self._max_size

    @max_size.setter
    def max_size(self, value: int) -> None:
        """Set a new maximum size for the cache.

        If the new size is smaller than the current cache size,
        the least recently used items will be evicted.
        """
        with self._lock:
            self._max_size = max(0, value)  # Ensure max_size is non-negative

            # If the new max_size is 0, it means unlimited size, so no eviction needed
            if self._max_size == 0:
                return

            # If we have more items than the new max_size, evict the oldest ones
            # Since OrderedDict maintains insertion order, we can just pop from the front
            while len(self._cache) > self._max_size:
                self._cache.popitem(last=False)  # Remove oldest item (first inserted)


class WeaveClientSendFileCache:
    """Cache for file create requests and responses with LRU eviction policy."""

    def __init__(self, max_size: int = 1000) -> None:
        """Initialize the file cache with a maximum size.

        Args:
            max_size: Maximum number of items to store in the cache.
                     If set to 0, cache size is unlimited.
        """
        self.cache: ThreadSafeLRUCache[str, Future[FileCreateRes]] = ThreadSafeLRUCache(
            max_size=max_size
        )

    def _key(self, req: FileCreateReq) -> str:
        """Generate a unique key for a file create request.

        The key is generated by hashing the project_id, file name, and file content.
        """
        cache_key_bytes = (req.project_id + "/" + req.name + "/").encode(
            "utf-8"
        ) + req.content
        return hashlib.sha256(cache_key_bytes).hexdigest()

    def get(self, req: FileCreateReq) -> Optional[Future[FileCreateRes]]:
        """Get a cached response for a file create request.

        Returns None if the request is not in the cache.
        """
        key = self._key(req)
        return self.cache.get(key)

    def put(self, req: FileCreateReq, res: Future[FileCreateRes]) -> None:
        """Cache a response for a file create request."""
        key = self._key(req)
        self.cache.put(key, res)

    def clear(self) -> None:
        """Clear all items from the cache."""
        self.cache.clear()

    def size(self) -> int:
        """Return the current number of items in the cache."""
        return self.cache.size()

    @property
    def max_size(self) -> int:
        """Return the maximum size of the cache."""
        return self.cache.max_size

    @max_size.setter
    def max_size(self, value: int) -> None:
        """Set a new maximum size for the cache."""
        self.cache.max_size = value
