"""Caching system for Logos ADK."""
from __future__ import annotations

import time
from abc import ABC, abstractmethod
from collections import OrderedDict
from dataclasses import dataclass
from datetime import datetime, UTC
from typing import Any, Awaitable, Callable

from logos_adk.types import Response, ToolCall, Usage, normalize_content_parts


@dataclass
class CacheConfig:
    """Cache configuration."""

    ttl: int | None = 3600
    semantic: bool = False
    embed_fn: Callable[[str], Awaitable[list[float]]] | None = None
    similarity_threshold: float = 0.95


@dataclass
class CacheEntry:
    """A cached LLM response."""

    response_data: dict
    created_at: str
    hit_count: int = 0
    vector: list[float] | None = None
    scope: str | None = None

    @classmethod
    def from_response(cls, response: Response) -> CacheEntry:
        """Create a CacheEntry from a Response."""
        return cls(
            response_data={
                "content": response.content,
                "content_parts": response.content_parts,
                "tool_calls": [
                    {"id": tc.id, "name": tc.name, "arguments": tc.arguments}
                    for tc in response.tool_calls
                ],
                "usage": {
                    "input_tokens": response.usage.input_tokens,
                    "output_tokens": response.usage.output_tokens,
                },
            },
            created_at=datetime.now(UTC).isoformat(),
        )

    def to_response(self) -> Response:
        """Reconstruct a Response from cached data."""
        data = self.response_data
        return Response(
            content=data["content"],
            content_parts=normalize_content_parts(data.get("content_parts")),
            tool_calls=[
                ToolCall(id=tc["id"], name=tc["name"], arguments=tc["arguments"])
                for tc in data.get("tool_calls", [])
            ],
            usage=Usage(
                input_tokens=data["usage"]["input_tokens"],
                output_tokens=data["usage"]["output_tokens"],
            ),
            raw={},
        )


class CacheBackend(ABC):
    """Abstract base class for cache backends."""

    @abstractmethod
    async def get(self, key: str) -> CacheEntry | None:
        ...

    @abstractmethod
    async def put(self, key: str, entry: CacheEntry, ttl: int | None) -> None:
        ...

    @abstractmethod
    async def invalidate(self, key: str) -> None:
        ...

    @abstractmethod
    async def clear(self) -> None:
        ...

    @abstractmethod
    async def get_stats(self) -> dict[str, Any]:
        """Return cache performance statistics."""
        ...

    async def search_similar(
        self,
        vector: list[float],
        threshold: float,
        limit: int = 1,
        scope: str | None = None,
    ) -> list[CacheEntry]:
        """Search for semantically similar entries. Override in subclasses."""
        return []


from logos_adk.math_utils import cosine_similarity as _cosine_similarity  # noqa: F401 — re-export for backwards compat


class InMemoryCacheBackend(CacheBackend):
    """In-memory cache with LRU eviction."""

    def __init__(self, maxsize: int = 1000) -> None:
        self._maxsize = maxsize
        self._store: OrderedDict[str, tuple[CacheEntry, float | None]] = OrderedDict()
        self._vectors: dict[str, list[float]] = {}

    async def get(self, key: str) -> CacheEntry | None:
        if key not in self._store:
            return None
        entry, expires_at = self._store[key]
        if expires_at is not None and time.time() > expires_at:
            del self._store[key]
            self._vectors.pop(key, None)
            return None
        entry.hit_count += 1
        self._store.move_to_end(key)
        return entry

    async def put(self, key: str, entry: CacheEntry, ttl: int | None) -> None:
        expires_at = (time.time() + ttl) if ttl is not None else None
        if key in self._store:
            self._store.move_to_end(key)
        self._store[key] = (entry, expires_at)
        if entry.vector is not None:
            self._vectors[key] = entry.vector
        while len(self._store) > self._maxsize:
            evicted_key, _ = self._store.popitem(last=False)
            self._vectors.pop(evicted_key, None)

    async def invalidate(self, key: str) -> None:
        self._store.pop(key, None)
        self._vectors.pop(key, None)

    async def clear(self) -> None:
        self._store.clear()
        self._vectors.clear()

    async def get_stats(self) -> dict[str, Any]:
        hits = sum(e.hit_count for e, _ in self._store.values())
        return {
            "backend": "InMemory",
            "total_entries": len(self._store),
            "total_hits": hits,
        }

    async def search_similar(
        self,
        vector: list[float],
        threshold: float,
        limit: int = 1,
        scope: str | None = None,
    ) -> list[CacheEntry]:
        results: list[tuple[float, CacheEntry]] = []
        for key, stored_vec in self._vectors.items():
            entry, expires_at = self._store.get(key, (None, None))
            if entry is None:
                continue
            if expires_at is not None and time.time() > expires_at:
                continue
            if scope is not None and entry.scope != scope:
                continue
            sim = _cosine_similarity(vector, stored_vec)
            if sim >= threshold:
                results.append((sim, entry))
        results.sort(key=lambda x: x[0], reverse=True)
        return [entry for _, entry in results[:limit]]
