from __future__ import annotations

import math
from datetime import datetime, timezone
from uuid import uuid4

from pydantic import BaseModel, Field

from nexus.core.state.embedding_fallback import deterministic_embedding


class WarmEntry(BaseModel):
    """Compressed episodic memory unit with embedding and recall stats."""

    entry_id: str = Field(default_factory=lambda: str(uuid4()))
    content: str
    original_keys: list[str]
    embedding: list[float]
    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
    last_recalled: datetime | None = None
    recall_count: int = 0
    compression_ratio: float = 1.0


class WarmLayer:
    """Session memory store with embedding-based retrieval."""

    def __init__(self, agent_id: str, model_name: str = "all-MiniLM-L6-v2", max_entries: int = 500) -> None:
        self.agent_id = agent_id
        self.model_name = model_name
        self.max_entries = max_entries
        self._entries: list[WarmEntry] = []
        self._model: object | None = None
        self._use_fallback = False

    def _get_model(self) -> object:
        if self._model is None and not self._use_fallback:
            try:
                from sentence_transformers import SentenceTransformer  # lazy import

                self._model = SentenceTransformer(self.model_name)
            except Exception:
                self._use_fallback = True
        return self._model

    def _embed(self, text: str) -> list[float]:
        model = self._get_model()
        if model is None or self._use_fallback:
            return deterministic_embedding(text)
        encoded = model.encode(text).tolist()  # type: ignore[attr-defined]
        return [float(x) for x in encoded]

    def add(self, content: str, original_keys: list[str], compression_ratio: float = 1.0) -> WarmEntry:
        """Embed and add compressed content; evict LRU on overflow."""
        emb = self._embed(content)
        entry = WarmEntry(
            content=content,
            original_keys=list(original_keys),
            embedding=emb,
            compression_ratio=float(compression_ratio),
        )
        self._entries.append(entry)
        if len(self._entries) > self.max_entries:
            self.evict_lru(len(self._entries) - self.max_entries)
        return entry

    def recall(self, query: str, top_k: int = 5, min_similarity: float = 0.65) -> list[WarmEntry]:
        """Recall top-k semantically similar entries above threshold."""
        if not self._entries:
            return []
        q = self._embed(query)
        scored = [(self._cosine_similarity(q, e.embedding), e) for e in self._entries]
        scored.sort(key=lambda item: item[0], reverse=True)

        recalled: list[WarmEntry] = []
        now = datetime.now(timezone.utc)
        for sim, entry in scored:
            if sim < min_similarity:
                continue
            entry.last_recalled = now
            entry.recall_count += 1
            recalled.append(entry)
            if len(recalled) >= top_k:
                break
        return recalled

    def _cosine_similarity(self, a: list[float], b: list[float]) -> float:
        """Compute cosine similarity and guard against zero vectors."""
        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 evict_lru(self, n: int = 1) -> list[WarmEntry]:
        """Evict n least-recently-recalled entries, then oldest first."""
        if n <= 0 or not self._entries:
            return []
        ranked = sorted(
            self._entries,
            key=lambda e: (
                e.last_recalled if e.last_recalled is not None else datetime.fromtimestamp(0, tz=timezone.utc),
                e.created_at,
            ),
        )
        to_remove = set(e.entry_id for e in ranked[:n])
        evicted = [e for e in ranked[:n]]
        self._entries = [e for e in self._entries if e.entry_id not in to_remove]
        return evicted

    def to_dict(self) -> list[dict]:
        """Serialize warm entries for checkpointing."""
        return [entry.model_dump() for entry in self._entries]

    @classmethod
    def from_dict(cls, data: list[dict]) -> WarmLayer:
        """Rehydrate a warm layer from serialized list; defaults agent id to restored."""
        layer = cls(agent_id="restored")
        layer._entries = [WarmEntry.model_validate(x) for x in data]
        return layer
