from __future__ import annotations

import time
from collections import OrderedDict
from dataclasses import dataclass, field
from typing import TYPE_CHECKING

import numpy as np

from secrag.core.config import Settings, get_settings
from secrag.core.logging import get_logger
from secrag.core.types import QueryResponse
from secrag.observability.tracing import span
from secrag.retrieval.embedder import Embedder

if TYPE_CHECKING:
    from numpy.typing import NDArray

log = get_logger(__name__)


@dataclass(slots=True)
class CacheEntry:
    question: str
    vector: NDArray[np.float32]
    response: QueryResponse
    created_at: float = field(default_factory=time.time)

    def is_expired(self, ttl_s: int) -> bool:
        return ttl_s > 0 and (time.time() - self.created_at) > ttl_s


@dataclass(slots=True)
class CacheStats:
    hits: int = 0
    misses: int = 0
    evictions: int = 0
    expirations: int = 0

    @property
    def lookups(self) -> int:
        return self.hits + self.misses

    @property
    def hit_rate(self) -> float:
        return self.hits / self.lookups if self.lookups else 0.0

    def to_dict(self) -> dict[str, float | int]:
        return {
            "hits": self.hits,
            "misses": self.misses,
            "evictions": self.evictions,
            "expirations": self.expirations,
            "lookups": self.lookups,
            "hit_rate": round(self.hit_rate, 4),
        }


class SemanticCache:
    def __init__(self, settings: Settings | None = None, embedder: Embedder | None = None) -> None:
        self.settings = settings or get_settings()
        self.embedder = embedder or Embedder(self.settings)
        self.stats = CacheStats()
        self._partitions: dict[str, OrderedDict[str, CacheEntry]] = {}

    @property
    def enabled(self) -> bool:
        return self.settings.cache_enabled

    @property
    def size(self) -> int:
        return sum(len(p) for p in self._partitions.values())

    def clear(self) -> None:
        self._partitions.clear()
        self.stats = CacheStats()

    def get(self, question: str, partition: str = "") -> QueryResponse | None:
        if not self.enabled:
            return None

        bucket = self._partitions.get(partition)
        if not bucket:
            self.stats.misses += 1
            return None

        with span("cache_lookup", partition=partition or "default"):
            vector = self.embedder.embed_query(question)
            best_key, best_score = None, 0.0

            for key, entry in list(bucket.items()):
                if entry.is_expired(self.settings.cache_ttl_s):
                    del bucket[key]
                    self.stats.expirations += 1
                    continue
                score = float(np.dot(vector, entry.vector))
                if score > best_score:
                    best_key, best_score = key, score

            if best_key is not None and best_score >= self.settings.cache_similarity_threshold:
                entry = bucket[best_key]
                bucket.move_to_end(best_key)
                self.stats.hits += 1
                log.info(
                    "cache_hit",
                    similarity=round(best_score, 4),
                    original=entry.question[:60],
                )
                hit = entry.response.model_copy(deep=True)
                hit.cached = True
                return hit

        self.stats.misses += 1
        return None

    def put(self, question: str, response: QueryResponse, partition: str = "") -> None:
        if not self.enabled:
            return
        bucket = self._partitions.setdefault(partition, OrderedDict())
        vector = self.embedder.embed_query(question)

        stored = response.model_copy(deep=True)
        stored.cached = False
        bucket[question] = CacheEntry(question=question, vector=vector, response=stored)
        bucket.move_to_end(question)

        while len(bucket) > self.settings.cache_max_entries:
            bucket.popitem(last=False)
            self.stats.evictions += 1

    def snapshot(self) -> dict[str, object]:
        return {
            "enabled": self.enabled,
            "size": self.size,
            "partitions": len(self._partitions),
            "threshold": self.settings.cache_similarity_threshold,
            **self.stats.to_dict(),
        }


def partition_key(tickers: list[str], fiscal_years: list[int], sections: list[str]) -> str:
    return "|".join(
        (
            ",".join(sorted(t.upper() for t in tickers)),
            ",".join(str(y) for y in sorted(fiscal_years)),
            ",".join(sorted(sections)),
        )
    )