import os
from typing import Any, Dict, List
from collections import defaultdict
from mem0 import Memory, MemoryClient
from crewai.utilities.chromadb import sanitize_collection_name

from crewai.memory.storage.interface import Storage

MAX_AGENT_ID_LENGTH_MEM0 = 255


class Mem0Storage(Storage):
    """
    Extends Storage to handle embedding and searching across entities using Mem0.
    """
    def __init__(self, type, crew=None, config=None):
        super().__init__()

        self._validate_type(type)
        self.memory_type = type
        self.crew = crew
        self.config = config or {}

        self._extract_config_values()
        self._initialize_memory()

    def _validate_type(self, type):
        supported_types = {"short_term", "long_term", "entities", "external"}
        if type not in supported_types:
            raise ValueError(
                f"Invalid type '{type}' for Mem0Storage. Must be one of: {', '.join(supported_types)}"
            )

    def _extract_config_values(self):
        self.mem0_run_id = self.config.get("run_id")
        self.includes = self.config.get("includes")
        self.excludes = self.config.get("excludes")
        self.custom_categories = self.config.get("custom_categories")
        self.infer = self.config.get("infer", True)

    def _initialize_memory(self):
        api_key = self.config.get("api_key") or os.getenv("MEM0_API_KEY")
        org_id = self.config.get("org_id")
        project_id = self.config.get("project_id")
        local_config = self.config.get("local_mem0_config")

        if api_key:
            self.memory = (
                MemoryClient(api_key=api_key, org_id=org_id, project_id=project_id)
                if org_id and project_id
                else MemoryClient(api_key=api_key)
            )
            if self.custom_categories:
                self.memory.update_project(custom_categories=self.custom_categories)
        else:
            self.memory = (
                Memory.from_config(local_config)
                if local_config and len(local_config)
                else Memory()
            )

    def _create_filter_for_search(self):
        """
        Returns:
            dict: A filter dictionary containing AND conditions for querying data.
                - Includes user_id and agent_id if both are present.
                - Includes user_id if only user_id is present.
                - Includes agent_id if only agent_id is present.
                - Includes run_id if memory_type is 'short_term' and mem0_run_id is present.
        """
        filter = defaultdict(list)

        if self.memory_type == "short_term" and self.mem0_run_id:
            filter["AND"].append({"run_id": self.mem0_run_id})
        else:
            user_id = self.config.get("user_id", "")
            agent_id = self.config.get("agent_id", "")

            if user_id and agent_id:
                filter["OR"].append({"user_id": user_id})
                filter["OR"].append({"agent_id": agent_id})
            elif user_id:
                filter["AND"].append({"user_id": user_id})
            elif agent_id:
                filter["AND"].append({"agent_id": agent_id})

        return filter

    def save(self, value: Any, metadata: Dict[str, Any]) -> None:
        user_id = self.config.get("user_id", "")
        assistant_message = [{"role" : "assistant","content" : value}]

        base_metadata = {
            "short_term": "short_term",
            "long_term": "long_term",
            "entities": "entity",
            "external": "external"
        }

        # Shared base params
        params: dict[str, Any] = {
            "metadata": {"type": base_metadata[self.memory_type], **metadata},
            "infer": self.infer
        }

        # MemoryClient-specific overrides
        if isinstance(self.memory, MemoryClient):
            params["includes"] = self.includes
            params["excludes"] = self.excludes
            params["output_format"] = "v1.1"
            params["version"] = "v2"

        if self.memory_type == "short_term" and self.mem0_run_id:
            params["run_id"] = self.mem0_run_id

        if user_id:
            params["user_id"] = user_id

        if agent_id := self.config.get("agent_id", self._get_agent_name()):
            params["agent_id"] = agent_id

        self.memory.add(assistant_message, **params)

    def search(self,query: str,limit: int = 3,score_threshold: float = 0.35) -> List[Any]:
        params = {
            "query": query,
            "limit": limit,
            "version": "v2",
            "output_format": "v1.1"
            }

        if user_id := self.config.get("user_id", ""):
            params["user_id"] = user_id

        memory_type_map = {
            "short_term": {"type": "short_term"},
            "long_term": {"type": "long_term"},
            "entities": {"type": "entity"},
            "external": {"type": "external"},
        }

        if self.memory_type in memory_type_map:
            params["metadata"] = memory_type_map[self.memory_type]
            if self.memory_type == "short_term":
                params["run_id"] = self.mem0_run_id

        # Discard the filters for now since we create the filters
        # automatically when the crew is created.

        params["filters"] = self._create_filter_for_search()
        params['threshold'] = score_threshold

        if isinstance(self.memory, Memory):
            del params["metadata"], params["version"], params['output_format']
            if params.get("run_id"):
                del params["run_id"]

        results = self.memory.search(**params)

        # This makes it compatible for Contextual Memory to retrieve
        for result in results["results"]:
            result["context"] = result["memory"]
        
        return [r for r in results["results"]]

    def reset(self):
        if self.memory:
            self.memory.reset()

    def _sanitize_role(self, role: str) -> str:
        """
        Sanitizes agent roles to ensure valid directory names.
        """
        return role.replace("\n", "").replace(" ", "_").replace("/", "_")

    def _get_agent_name(self) -> str:
        if not self.crew:
            return ""

        agents = self.crew.agents
        agents = [self._sanitize_role(agent.role) for agent in agents]
        agents = "_".join(agents)
        return sanitize_collection_name(name=agents, max_collection_length=MAX_AGENT_ID_LENGTH_MEM0)