"""
Memory Backends

Abstracts the underlying memory storage to allow pluggable backends.
Default implementation uses Mem0 for fact storage.

MCAL adds reasoning/intent layers ON TOP of these backends.
"""

from __future__ import annotations

from abc import ABC, abstractmethod
from typing import Any, Optional
from datetime import datetime, timezone

from pydantic import BaseModel, Field


def _utc_now() -> datetime:
    """Return current UTC time (timezone-aware)."""
    return datetime.now(timezone.utc)


# =============================================================================
# Backend Protocol / Abstract Base
# =============================================================================

class MemoryEntry(BaseModel):
    """Standardized memory entry across backends."""
    id: str
    content: str
    metadata: dict = Field(default_factory=dict)
    score: Optional[float] = None  # Relevance score from search
    created_at: Optional[datetime] = None
    updated_at: Optional[datetime] = None


class MemoryBackend(ABC):
    """
    Abstract base for memory backends.
    
    MCAL can work with different backends:
    - Mem0Backend (default, recommended)
    - StandaloneBackend (for testing or simple use cases)
    - Custom backends (implement this interface)
    """
    
    @abstractmethod
    def add(
        self,
        messages: list[dict],
        user_id: str,
        metadata: Optional[dict] = None
    ) -> list[MemoryEntry]:
        """
        Add messages to memory, extracting facts.
        
        Args:
            messages: List of message dicts [{"role": "user", "content": "..."}]
            user_id: User identifier
            metadata: Optional metadata to attach
            
        Returns:
            List of extracted/stored memory entries
        """
        pass
    
    @abstractmethod
    def search(
        self,
        query: str,
        user_id: str,
        limit: int = 10
    ) -> list[MemoryEntry]:
        """
        Search memories by semantic similarity.
        
        Args:
            query: Search query
            user_id: User identifier
            limit: Maximum results to return
            
        Returns:
            List of relevant memory entries with scores
        """
        pass
    
    @abstractmethod
    def get_all(self, user_id: str) -> list[MemoryEntry]:
        """Get all memories for a user."""
        pass
    
    @abstractmethod
    def delete(self, memory_id: str) -> bool:
        """Delete a specific memory."""
        pass
    
    @abstractmethod
    def delete_all(self, user_id: str) -> bool:
        """Delete all memories for a user."""
        pass


# =============================================================================
# Mem0 Backend Implementation
# =============================================================================

class Mem0Backend(MemoryBackend):
    """
    Mem0-based memory backend.
    
    Leverages Mem0's production-ready infrastructure:
    - Intelligent fact extraction
    - Vector + Graph + Key-Value storage
    - Deduplication and conflict resolution
    - 90% token reduction
    
    MCAL adds reasoning/intent layers on top.
    
    Usage:
        backend = Mem0Backend(config={
            "llm": {"provider": "anthropic", "config": {"model": "claude-3-5-sonnet"}},
            "embedder": {"provider": "openai", "config": {"model": "text-embedding-3-small"}}
        })
        
        # Or with API key for hosted Mem0
        backend = Mem0Backend(api_key="your-mem0-api-key")
    """
    
    def __init__(
        self,
        config: Optional[dict] = None,
        api_key: Optional[str] = None
    ):
        """
        Initialize Mem0 backend.
        
        Args:
            config: Mem0 configuration dict (for self-hosted)
            api_key: Mem0 API key (for hosted service)
        """
        try:
            from mem0 import Memory, MemoryClient
        except ImportError:
            raise ImportError(
                "mem0ai is required for Mem0Backend. "
                "Install with: pip install mem0ai"
            )
        
        if api_key:
            # Use hosted Mem0 service
            self._client = MemoryClient(api_key=api_key)
            self._is_hosted = True
        else:
            # Use self-hosted Mem0
            self._client = Memory.from_config(config or self._default_config())
            self._is_hosted = False
    
    def _default_config(self) -> dict:
        """Default Mem0 configuration using HuggingFace for local embeddings."""
        return {
            "llm": {
                "provider": "aws_bedrock",
                "config": {
                    "model": "us.meta.llama3-3-70b-instruct-v1:0",
                    "temperature": 0.1,
                }
            },
            "embedder": {
                "provider": "huggingface",
                "config": {
                    "model": "sentence-transformers/all-MiniLM-L6-v2"
                }
            },
            "vector_store": {
                "provider": "faiss",
                "config": {
                    "embedding_model_dims": 384,  # all-MiniLM-L6-v2 dimension
                    "path": "/tmp/mcal_faiss"
                }
            }
        }
    
    def add(
        self,
        messages: list[dict],
        user_id: str,
        metadata: Optional[dict] = None
    ) -> list[MemoryEntry]:
        """Add messages to Mem0, extracting facts."""
        result = self._client.add(
            messages=messages,
            user_id=user_id,
            metadata=metadata or {}
        )
        
        # Convert Mem0 response to standardized format
        entries = []
        
        # Handle different Mem0 response formats
        if isinstance(result, dict):
            results_list = result.get("results", result.get("memories", []))
        elif isinstance(result, list):
            results_list = result
        else:
            results_list = []
        
        for item in results_list:
            if isinstance(item, dict):
                entries.append(MemoryEntry(
                    id=item.get("id", ""),
                    content=item.get("memory", item.get("content", "")),
                    metadata=item.get("metadata", {}),
                    created_at=item.get("created_at"),
                    updated_at=item.get("updated_at")
                ))
        
        return entries
    
    def search(
        self,
        query: str,
        user_id: str,
        limit: int = 10
    ) -> list[MemoryEntry]:
        """Search Mem0 memories by semantic similarity."""
        result = self._client.search(
            query=query,
            user_id=user_id,
            limit=limit
        )
        
        # Convert to standardized format
        entries = []
        
        if isinstance(result, dict):
            results_list = result.get("results", result.get("memories", []))
        elif isinstance(result, list):
            results_list = result
        else:
            results_list = []
        
        for item in results_list:
            if isinstance(item, dict):
                entries.append(MemoryEntry(
                    id=item.get("id", ""),
                    content=item.get("memory", item.get("content", "")),
                    metadata=item.get("metadata") or {},
                    score=item.get("score", item.get("similarity", 0.0))
                ))
        
        return entries
    
    def get_all(self, user_id: str) -> list[MemoryEntry]:
        """Get all memories for a user from Mem0."""
        result = self._client.get_all(user_id=user_id)
        
        entries = []
        if isinstance(result, dict):
            results_list = result.get("results", result.get("memories", []))
        elif isinstance(result, list):
            results_list = result
        else:
            results_list = []
        
        for item in results_list:
            if isinstance(item, dict):
                entries.append(MemoryEntry(
                    id=item.get("id", ""),
                    content=item.get("memory", item.get("content", "")),
                    metadata=item.get("metadata", {})
                ))
        
        return entries
    
    def delete(self, memory_id: str) -> bool:
        """Delete a specific memory from Mem0."""
        try:
            self._client.delete(memory_id=memory_id)
            return True
        except Exception:
            return False
    
    def delete_all(self, user_id: str) -> bool:
        """Delete all memories for a user from Mem0."""
        try:
            self._client.delete_all(user_id=user_id)
            return True
        except Exception:
            return False


# =============================================================================
# Standalone Backend (for testing / simple use)
# =============================================================================

class StandaloneBackend(MemoryBackend):
    """
    Simple in-memory backend for testing or lightweight use.
    
    Does NOT have Mem0's intelligent extraction - just stores raw content.
    Useful for:
    - Unit testing
    - Development without API keys
    - Simple use cases
    
    For production, use Mem0Backend.
    """
    
    def __init__(self):
        self._memories: dict[str, dict[str, MemoryEntry]] = {}  # user_id -> {memory_id -> entry}
        self._counter = 0
    
    def add(
        self,
        messages: list[dict],
        user_id: str,
        metadata: Optional[dict] = None
    ) -> list[MemoryEntry]:
        """Store messages as simple memories (no intelligent extraction)."""
        if user_id not in self._memories:
            self._memories[user_id] = {}
        
        entries = []
        for msg in messages:
            self._counter += 1
            memory_id = f"mem_{self._counter}"
            
            entry = MemoryEntry(
                id=memory_id,
                content=msg.get("content", ""),
                metadata={
                    "role": msg.get("role", "unknown"),
                    **(metadata or {})
                },
                created_at=_utc_now()
            )
            
            self._memories[user_id][memory_id] = entry
            entries.append(entry)
        
        return entries
    
    def search(
        self,
        query: str,
        user_id: str,
        limit: int = 10
    ) -> list[MemoryEntry]:
        """Simple keyword-based search (no semantic similarity)."""
        if user_id not in self._memories:
            return []
        
        query_lower = query.lower()
        results = []
        
        for entry in self._memories[user_id].values():
            if query_lower in entry.content.lower():
                entry.score = 1.0  # Simple match
                results.append(entry)
        
        return results[:limit]
    
    def get_all(self, user_id: str) -> list[MemoryEntry]:
        """Get all memories for a user."""
        if user_id not in self._memories:
            return []
        return list(self._memories[user_id].values())
    
    def delete(self, memory_id: str) -> bool:
        """Delete a specific memory."""
        for user_memories in self._memories.values():
            if memory_id in user_memories:
                del user_memories[memory_id]
                return True
        return False
    
    def delete_all(self, user_id: str) -> bool:
        """Delete all memories for a user."""
        if user_id in self._memories:
            self._memories[user_id] = {}
            return True
        return False
