import chromadb
from chromadb.config import Settings
from datetime import datetime
import uuid

class MemoryManager:
    def __init__(self, persist_directory="memory/chroma_db"):
        self.client = chromadb.PersistentClient(path=persist_directory)
        
        # Vytvoříme collections pro různé typy paměti
        self.collections = {
            "rules": self.client.get_or_create_collection("rules"),
            "projects": self.client.get_or_create_collection("projects"),
            "preferences": self.client.get_or_create_collection("preferences"),
            "conversations": self.client.get_or_create_collection("conversations"),
            "errors": self.client.get_or_create_collection("errors"),
        }

    def add_memory(self, memory_type: str, content: str, metadata: dict = None):
        """
        Uloží novou vzpomínku do daného typu paměti.
        """
        if memory_type not in self.collections:
            raise ValueError(f"Neznámý typ paměti: {memory_type}")

        if metadata is None:
            metadata = {}

        metadata["timestamp"] = datetime.now().isoformat()
        metadata["type"] = memory_type

        self.collections[memory_type].add(
            documents=[content],
            metadatas=[metadata],
            ids=[str(uuid.uuid4())]
        )
        print(f"✅ Uloženo do paměti [{memory_type}]")

    def retrieve_memories(self, query: str, memory_type: str = None, limit: int = 5):
        """
        Vyhledá relevantní vzpomínky podle dotazu.
        """
        results = []

        collections_to_search = [memory_type] if memory_type else self.collections.keys()

        for col_name in collections_to_search:
            if col_name not in self.collections:
                continue

            collection = self.collections[col_name]
            res = collection.query(
                query_texts=[query],
                n_results=limit
            )

            for doc, meta, dist in zip(res["documents"][0], res["metadatas"][0], res["distances"][0]):
                results.append({
                    "type": col_name,
                    "content": doc,
                    "metadata": meta,
                    "relevance": 1 - dist  # převedeme vzdálenost na relevanci
                })

        # Seřadíme podle relevance (nejrelevantnější první)
        results.sort(key=lambda x: x["relevance"], reverse=True)
        return results[:limit]

    def get_all_memories(self, memory_type: str = None):
        """Vrátí všechny vzpomínky daného typu (pro debug)."""
        if memory_type:
            return self.collections[memory_type].get()
        else:
            all_data = {}
            for name, col in self.collections.items():
                all_data[name] = col.get()
            return all_data

def save_memory(memory_type: str, content: str, metadata: dict = None):
    """
    Uloží novou vzpomínku. Tuto funkci může volat model.
    """
    mm = MemoryManager()
    mm.add_memory(memory_type, content, metadata)
    return f"✅ Uloženo do paměti typu '{memory_type}'."