import json
from datetime import datetime, timezone
from typing import Iterable, List, Optional
from uuid import uuid4

import numpy as np
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity

from schemas import Memory, RetrievedMemory, MemoryWriteDecision


class MemoryStore:
    def __init__(
        self,
        memories: Iterable[Memory],
        model_name: str = "all-MiniLM-L6-v2",
        embedding_model: Optional[SentenceTransformer] = None,
    ):
        self.memories = list(memories)
        self.embedding_model = embedding_model or SentenceTransformer(model_name)
        self.refresh_embeddings()

    def refresh_embeddings(self) -> None:
        self.memory_texts = [memory.content for memory in self.memories]
        if not self.memory_texts:
            self.memory_embeddings = np.empty((0, 0))
            return
        self.memory_embeddings = self.embedding_model.encode(self.memory_texts)

    def retrieve(self, query: str, top_k: int = 4) -> List[RetrievedMemory]:
        if not self.memories or not query.strip():
            return []

        safe_top_k = max(1, min(int(top_k), len(self.memories)))
        query_embedding = self.embedding_model.encode([query])
        similarities = cosine_similarity(query_embedding, self.memory_embeddings)[0]
        ranked_indices = similarities.argsort()[::-1][:safe_top_k]

        return [
            RetrievedMemory(
                memory=self.memories[index],
                similarity_score=float(similarities[index]),
                rank=rank,
            )
            for rank, index in enumerate(ranked_indices, start=1)
        ]

    def save_memory_from_decision(self, write_decision: MemoryWriteDecision) -> Optional[Memory]:
        if write_decision.outcome != "save" or not write_decision.should_save:
            return None
        if not write_decision.proposed_memory:
            return None

        new_memory = Memory(
            memory_id=f"session_{uuid4().hex[:10]}",
            memory_type=write_decision.memory_type or "episodic",
            content=write_decision.proposed_memory,
            tags=["session_memory", "policy_approved"],
            importance=write_decision.importance or 3,
            created_at=datetime.now(timezone.utc).date().isoformat(),
            source="session_write_policy",
        )
        self.memories.append(new_memory)
        self.refresh_embeddings()
        return new_memory

    def to_dicts(self) -> List[dict]:
        return [memory.model_dump() for memory in self.memories]

    def export_to_json(self, filepath: str) -> None:
        with open(filepath, "w", encoding="utf-8") as file:
            json.dump(self.to_dicts(), file, indent=2)


def memories_from_dicts(memory_data: Iterable[dict]) -> List[Memory]:
    return [Memory(**item) for item in memory_data]


def load_memories_from_json(filepath: str) -> List[Memory]:
    with open(filepath, "r", encoding="utf-8") as file:
        memory_data = json.load(file)
    return memories_from_dicts(memory_data)