import math
import threading
import time
from typing import Any, Dict, List, Optional, Tuple

from src.embeddings.embedder import embed_texts


def normalize_query(query: str) -> str:
    return " ".join((query or "").strip().lower().split())


class ExactMatchCache:
    def __init__(self, ttl_seconds: int = 3600):
        self.ttl_seconds = max(1, int(ttl_seconds))
        self._store: Dict[str, Dict[str, Any]] = {}
        self._lock = threading.Lock()
        self._hits = 0
        self._misses = 0

    def _is_expired(self, created_at: float) -> bool:
        return (time.time() - created_at) > self.ttl_seconds

    def get(self, query: str) -> Optional[Dict[str, Any]]:
        key = normalize_query(query)
        with self._lock:
            value = self._store.get(key)
            if not value:
                self._misses += 1
                return None
            if self._is_expired(value["timestamp"]):
                self._store.pop(key, None)
                self._misses += 1
                return None
            self._hits += 1
            return dict(value)

    def set(self, query: str, answer: str, docs: Optional[List[Dict[str, Any]]] = None):
        key = normalize_query(query)
        payload = {
            "answer": answer,
            "docs": list(docs or []),
            "timestamp": time.time(),
        }
        with self._lock:
            self._store[key] = payload

    def invalidate(self, query: Optional[str] = None):
        with self._lock:
            if query is None:
                self._store.clear()
                return
            self._store.pop(normalize_query(query), None)

    def stats(self) -> Dict[str, Any]:
        with self._lock:
            total = self._hits + self._misses
            ratio = (self._hits / total) if total else 0.0
            return {
                "entries": len(self._store),
                "hits": self._hits,
                "misses": self._misses,
                "hit_ratio": round(ratio, 4),
                "ttl_seconds": self.ttl_seconds,
            }


class SemanticCache:
    def __init__(self, ttl_seconds: int = 3600, similarity_threshold: float = 0.9):
        self.ttl_seconds = max(1, int(ttl_seconds))
        self.similarity_threshold = float(similarity_threshold)
        self._items: List[Dict[str, Any]] = []
        self._lock = threading.Lock()
        self._hits = 0
        self._misses = 0

    def _is_expired(self, created_at: float) -> bool:
        return (time.time() - created_at) > self.ttl_seconds

    def _cosine(self, a: List[float], b: List[float]) -> float:
        if not a or not b or len(a) != len(b):
            return 0.0
        dot = sum(x * y for x, y in zip(a, b))
        norm_a = math.sqrt(sum(x * x for x in a))
        norm_b = math.sqrt(sum(y * y for y in b))
        if norm_a == 0.0 or norm_b == 0.0:
            return 0.0
        return dot / (norm_a * norm_b)

    def _prune_expired(self):
        self._items = [
            item for item in self._items if not self._is_expired(item["timestamp"])
        ]

    def get(self, query: str) -> Optional[Dict[str, Any]]:
        normalized = normalize_query(query)
        if not normalized:
            self._misses += 1
            return None
        query_vector = embed_texts([normalized], batch_size=1)[0]

        with self._lock:
            self._prune_expired()
            best: Optional[Tuple[float, Dict[str, Any]]] = None
            for item in self._items:
                score = self._cosine(query_vector, item["vector"])
                if best is None or score > best[0]:
                    best = (score, item)
            if best is None or best[0] < self.similarity_threshold:
                self._misses += 1
                return None
            self._hits += 1
            result = dict(best[1]["value"])
            result["semantic_similarity"] = round(best[0], 4)
            return result

    def set(self, query: str, answer: str, docs: Optional[List[Dict[str, Any]]] = None):
        normalized = normalize_query(query)
        if not normalized:
            return
        vector = embed_texts([normalized], batch_size=1)[0]
        item = {
            "query": normalized,
            "vector": vector,
            "value": {
                "answer": answer,
                "docs": list(docs or []),
            },
            "timestamp": time.time(),
        }
        with self._lock:
            self._prune_expired()
            self._items.append(item)

    def invalidate(self):
        with self._lock:
            self._items = []

    def stats(self) -> Dict[str, Any]:
        with self._lock:
            total = self._hits + self._misses
            ratio = (self._hits / total) if total else 0.0
            return {
                "entries": len(self._items),
                "hits": self._hits,
                "misses": self._misses,
                "hit_ratio": round(ratio, 4),
                "ttl_seconds": self.ttl_seconds,
                "similarity_threshold": self.similarity_threshold,
            }