"""
Persistent per-user memory storage backed by ChromaDB and LlamaIndex.
"""

from __future__ import annotations
import logging
import os
from typing import Any, Optional
from uuid import uuid4
from datetime import datetime, timezone
import chromadb
import yaml
from llama_index.core import Document, Settings, StorageContext, VectorStoreIndex
from llama_index.core.vector_stores import MetadataFilter, MetadataFilters
from llama_index.vector_stores.chroma import ChromaVectorStore


# Configure logging for this module
logger = logging.getLogger(__name__)


class UserMemoryStore:
    """Stores and retrieves user-specific memories in a dedicated vector index."""

    def __init__(
        self,
        config: Optional[dict[str, Any]] = None,
        config_path: Optional[str] = None,
    ) -> None:
        """Initializes the user memory store from a config dict or YAML file."""
        # Guardrail: prevent both config and config_path from being provided simultaneously
        if config is not None and config_path is not None:
            raise ValueError("Provide either 'config' or 'config_path', not both.")

        # If a config dict is provided, validate and use it
        if config is not None:
            self.config = self._validate_config(config)

        # Otherwise, load it from the YAML file
        else:
            if config_path is None:
                config_path = os.path.join(os.path.dirname(__file__), "rag_config.yaml")
            self.config = self._load_config(config_path)

        self._setup_vector_store()
        self._index: Optional[VectorStoreIndex] = None

        logger.info("UserMemoryStore initialized successfully.")

    def _validate_config(self, config: dict[str, Any]) -> dict[str, Any]:
        """Validates the config structure required for the user memory store."""
        if not isinstance(config, dict):
            raise ValueError("RAG configuration must be a dictionary.")

        # Validate the keys needed for memory storage
        self._require_config_key(config, ("storage",))
        self._require_config_key(config, ("storage", "persist_dir"))
        self._require_config_key(config, ("storage", "user_memory_collection_name"))
        self._require_config_key(config, ("memory_retrieval",))
        self._require_config_key(config, ("memory_retrieval", "similarity_top_k"))

        return config

    def _load_config(self, config_path: str) -> dict[str, Any]:
        """Loads the memory store configuration."""
        logger.info("Loading user memory configuration from: %s", config_path)
        if not os.path.exists(config_path):
            raise FileNotFoundError(f"Configuration file not found at {config_path}")

        with open(config_path, "r", encoding="utf-8") as config_file:
            config = yaml.safe_load(config_file)

            return self._validate_config(config)

    def _require_config_key(
        self, config: dict[str, Any], key_path: tuple[str, ...]
    ) -> Any:
        """Returns a required config value or raises an error if the key is missing."""
        current: Any = config

        for key in key_path:
            if not isinstance(current, dict) or key not in current:
                dotted_path = ".".join(key_path)
                raise KeyError(f"Missing required config key: {dotted_path}")
            current = current[key]

        return current

    def _setup_vector_store(self) -> None:
        """Creates or opens the dedicated Chroma collection for user memory."""
        persist_dir = self.config["storage"]["persist_dir"]
        collection_name = self.config["storage"]["user_memory_collection_name"]

        # Create the persistence directory if it doesn't exist
        if not os.path.exists(persist_dir):
            logger.info("Creating user memory persistence directory: %s", persist_dir)
            os.makedirs(persist_dir, exist_ok=True)

        # Create the Chroma client and open/create the user memory collection
        db = chromadb.PersistentClient(path=persist_dir)
        chroma_collection = db.get_or_create_collection(collection_name)

        # Wrap the Chroma collection in a LlamaIndex vector store and storage context
        self.vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
        self.storage_context = StorageContext.from_defaults(
            vector_store=self.vector_store
        )

        logger.info(
            "User memory vector store ready at %s (collection=%s)",
            persist_dir,
            collection_name,
        )

    def _ensure_embed_model(self) -> None:
        """Fails fast if no global embedding model has been configured."""
        if Settings.embed_model is None:
            raise RuntimeError(
                "LlamaIndex Settings.embed_model is not configured. "
                "Configure it before using UserMemoryStore."
            )

    def _get_or_create_index(self) -> VectorStoreIndex:
        """Loads or creates the vector index for the user memory store."""
        if self._index is None:
            # Check that the embedding model is configured before loading the index
            self._ensure_embed_model()

            # Load the index from the existing Chroma vector store
            logger.info("Loading user memory index from vector store.")
            self._index = VectorStoreIndex.from_vector_store(
                self.vector_store,
                storage_context=self.storage_context,
            )
        return self._index

    def add_memory(
        self, user_id: str, memory_text: str, metadata: Optional[dict[str, Any]] = None
    ) -> None:
        """Inserts a single memory for a specific user into persistent storage with metadata."""
        # Strip whitespace and validate inputs
        user_id = user_id.strip()
        memory_text = memory_text.strip()

        if not user_id:
            raise ValueError("user_id must be a non-empty string.")
        if not memory_text:
            raise ValueError("memory_text must be a non-empty string.")

        # Combine provided metadata with the required user_id field for filtering
        memory_metadata: dict[str, Any] = dict(metadata or {})
        memory_metadata["user_id"] = user_id
        memory_metadata.setdefault("created_at", datetime.now(timezone.utc).isoformat())

        # Build a LlamaIndex Document with a unique ID for this memory entry
        document = Document(
            text=memory_text,
            metadata=memory_metadata,
            doc_id=f"user-memory-{user_id}-{uuid4()}",
        )

        logger.info("Adding memory for user_id=%s", user_id)
        index = self._get_or_create_index()

        # Insert the document into the vector index
        index.insert(document)

    def query_memories(
        self, user_id: str, query: str, top_k: Optional[int] = None
    ) -> Any:
        """Retrieves the most relevant memories for a specific user."""
        # Strip whitespace and validate inputs
        user_id = user_id.strip()
        query = query.strip()

        if not user_id:
            raise ValueError("user_id must be a non-empty string.")
        if not query:
            raise ValueError("query must be a non-empty string.")

        # Use the memory retrieval top_k if not explicitly provided
        if top_k is None:
            top_k = self.config["memory_retrieval"]["similarity_top_k"]
        if top_k <= 0:
            raise ValueError("top_k must be greater than 0.")

        # Only retrieve documents that match the user_id in their metadata
        filters = MetadataFilters(
            filters=[MetadataFilter(key="user_id", value=user_id)]
        )

        logger.info("Querying memories for user_id=%s with top_k=%s", user_id, top_k)
        index = self._get_or_create_index()

        # Build retriever that applies the user_id filter and retrieves most similar memories
        retriever = index.as_retriever(similarity_top_k=top_k, filters=filters)
        return retriever.retrieve(query)