"""
MCAL Memory Backend for Microsoft AutoGen

Implements AutoGen's Memory interface with goal-aware memory capabilities.
"""

from __future__ import annotations

import threading
import time
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional, TYPE_CHECKING

from mcal_autogen._compat import (
    AUTOGEN_CORE_AVAILABLE,
    check_autogen_core,
    Memory,
    MemoryContent,
    MemoryQueryResult,
    UpdateContextResult,
    MemoryMimeType,
    CancellationToken,
    ChatCompletionContext,
    SystemMessage,
)

if TYPE_CHECKING:
    from mcal import MCAL


class MCALMemory(Memory if AUTOGEN_CORE_AVAILABLE else object):
    """
    MCAL memory backend for Microsoft AutoGen.
    
    This class implements AutoGen's Memory interface, providing
    goal-aware memory with context preservation for multi-agent systems.
    
    Key Features:
    - Goal-aware search that prioritizes project-relevant memories
    - Decision tracking for architectural choices
    - User isolation for multi-tenant scenarios
    - TTL support with automatic expiration
    - Thread-safe operations
    
    Usage:
        from mcal import MCAL
        from mcal_autogen import MCALMemory
        from autogen_agentchat.agents import AssistantAgent
        
        mcal = MCAL(goal="Build fraud detection system")
        memory = MCALMemory(mcal, user_id="user_123")
        
        agent = AssistantAgent(
            name="analyst",
            model_client=model_client,
            memory=[memory],
        )
        
    Args:
        mcal: Initialized MCAL instance
        user_id: User identifier for memory isolation
        name: Memory instance name (default: "mcal_memory")
        max_results: Maximum results to return from queries (default: 10)
        score_threshold: Minimum relevance score 0-1 (default: 0.0)
        default_ttl_minutes: Default TTL in minutes (None = no expiration)
        enable_goal_tracking: Whether to extract goals from content
        include_decisions: Whether to include decisions in search results
    """
    
    def __init__(
        self,
        mcal: "MCAL",
        user_id: str = "default",
        name: str = "mcal_memory",
        max_results: int = 10,
        score_threshold: float = 0.0,
        default_ttl_minutes: Optional[float] = None,
        enable_goal_tracking: bool = True,
        include_decisions: bool = True,
    ):
        check_autogen_core()
        
        self._mcal = mcal
        self._user_id = user_id
        self._name = name
        self._max_results = max_results
        self._score_threshold = score_threshold
        self._default_ttl_minutes = default_ttl_minutes
        self._enable_goal_tracking = enable_goal_tracking
        self._include_decisions = include_decisions
        
        # Thread safety
        self._lock = threading.RLock()
        
        # In-memory storage for MemoryContent items
        # Maps content_id -> MemoryContent
        self._items: Dict[str, MemoryContent] = {}
        self._created_at: Dict[str, datetime] = {}
        self._expires_at: Dict[str, Optional[float]] = {}  # Unix timestamp
        
        # Counter for generating unique IDs
        self._id_counter = 0
    
    @property
    def name(self) -> str:
        """Get the memory instance identifier."""
        return self._name
    
    @property
    def user_id(self) -> str:
        """Get the user identifier."""
        return self._user_id
    
    @property
    def mcal(self) -> "MCAL":
        """Get the underlying MCAL instance."""
        return self._mcal
    
    def _generate_id(self) -> str:
        """Generate a unique content ID."""
        with self._lock:
            self._id_counter += 1
            return f"mem_{self._user_id}_{self._id_counter}_{int(time.time() * 1000)}"
    
    def _is_expired(self, content_id: str) -> bool:
        """Check if a memory entry has expired."""
        expires = self._expires_at.get(content_id)
        if expires is None:
            return False
        return time.time() > expires
    
    def _cleanup_expired(self) -> int:
        """Remove expired entries. Returns count of removed items."""
        removed = 0
        with self._lock:
            expired_ids = [
                cid for cid in list(self._items.keys())
                if self._is_expired(cid)
            ]
            for cid in expired_ids:
                del self._items[cid]
                self._created_at.pop(cid, None)
                self._expires_at.pop(cid, None)
                removed += 1
        return removed
    
    def _extract_text(self, content: MemoryContent | str) -> str:
        """Extract searchable text from content."""
        if isinstance(content, str):
            return content
        
        if hasattr(content, 'content'):
            val = content.content
            if isinstance(val, str):
                return val
            elif isinstance(val, dict):
                # Try common keys
                for key in ['text', 'content', 'message', 'data']:
                    if key in val:
                        return str(val[key])
                # Fallback to JSON-like string
                import json
                try:
                    return json.dumps(val)
                except (TypeError, ValueError):
                    return str(val)
            elif isinstance(val, bytes):
                try:
                    return val.decode('utf-8')
                except UnicodeDecodeError:
                    return str(val)
            else:
                return str(val)
        
        return str(content)
    
    async def add(
        self,
        content: MemoryContent,
        cancellation_token: Optional[CancellationToken] = None,
    ) -> None:
        """
        Add new content to memory.
        
        Args:
            content: Memory content to store
            cancellation_token: Optional token to cancel operation
        """
        # Check cancellation
        if cancellation_token and hasattr(cancellation_token, 'cancelled'):
            if getattr(cancellation_token, 'cancelled', False):
                return
        
        content_id = self._generate_id()
        now = datetime.now(timezone.utc)
        
        # Calculate expiration
        expires_at: Optional[float] = None
        ttl_minutes = self._default_ttl_minutes
        
        # Check for TTL in metadata
        if content.metadata and 'ttl_minutes' in content.metadata:
            ttl_minutes = content.metadata['ttl_minutes']
        
        if ttl_minutes is not None:
            expires_at = time.time() + (ttl_minutes * 60)
        
        # Store the content
        with self._lock:
            self._items[content_id] = content
            self._created_at[content_id] = now
            self._expires_at[content_id] = expires_at
        
        # Also store in MCAL for goal-aware search
        text = self._extract_text(content)
        if text:
            metadata = dict(content.metadata) if content.metadata else {}
            metadata['content_id'] = content_id
            metadata['user_id'] = self._user_id
            metadata['source'] = 'autogen_memory'
            
            # Add to MCAL for goal-aware retrieval
            try:
                await self._mcal.add(
                    messages=[{"role": "user", "content": text}],
                    user_id=self._user_id,
                    metadata=metadata,
                )
            except Exception:
                # Silently continue if MCAL storage fails
                # The in-memory storage still works
                pass
    
    async def query(
        self,
        query: str | MemoryContent,
        cancellation_token: Optional[CancellationToken] = None,
        **kwargs: Any,
    ) -> MemoryQueryResult:
        """
        Query memory for relevant content.
        
        Uses MCAL's goal-aware search to find the most relevant memories
        based on the current project goal and query.
        
        Args:
            query: Query string or MemoryContent
            cancellation_token: Optional token to cancel operation
            **kwargs: Additional query parameters
        
        Returns:
            MemoryQueryResult containing matching memories
        """
        # Check cancellation
        if cancellation_token and hasattr(cancellation_token, 'cancelled'):
            if getattr(cancellation_token, 'cancelled', False):
                return MemoryQueryResult(results=[])
        
        # Cleanup expired entries
        self._cleanup_expired()
        
        query_text = self._extract_text(query)
        if not query_text:
            return MemoryQueryResult(results=[])
        
        # Get search parameters
        max_results = kwargs.get('max_results', self._max_results)
        score_threshold = kwargs.get('score_threshold', self._score_threshold)
        
        results: List[MemoryContent] = []
        
        # Try MCAL search first for goal-aware results
        try:
            search_result = await self._mcal.search(
                query=query_text,
                user_id=self._user_id,
                limit=max_results,
                include_decisions=self._include_decisions,
            )
            
            # Convert MCAL MemoryEntry results to MemoryContent
            memories = search_result.memories if search_result else []
            for mem in memories[:max_results]:
                score = mem.score or 0.0
                text = mem.content
                metadata: Dict[str, Any] = dict(mem.metadata) if mem.metadata else {}
                
                # Apply score threshold
                if score >= score_threshold:
                    metadata['score'] = score
                    results.append(MemoryContent(
                        content=text,
                        mime_type=MemoryMimeType.TEXT,
                        metadata=metadata,
                    ))
        except Exception:
            # Fallback to in-memory search if MCAL fails
            pass
        
        # If no MCAL results, fall back to simple in-memory search
        if not results:
            query_lower = query_text.lower()
            with self._lock:
                for content_id, content in self._items.items():
                    if self._is_expired(content_id):
                        continue
                    
                    text = self._extract_text(content)
                    # Simple text matching
                    if query_lower in text.lower():
                        metadata = dict(content.metadata) if content.metadata else {}
                        metadata['score'] = 0.5  # Default score for text match
                        results.append(MemoryContent(
                            content=text,
                            mime_type=content.mime_type,
                            metadata=metadata,
                        ))
                        
                        if len(results) >= max_results:
                            break
        
        return MemoryQueryResult(results=results[:max_results])
    
    async def update_context(
        self,
        model_context: ChatCompletionContext,
    ) -> UpdateContextResult:
        """
        Update the model context with relevant memories.
        
        This method retrieves the conversation history, uses the last
        message as a query to find relevant memories, and adds them
        to the context as a system message.
        
        Args:
            model_context: The model context to update
        
        Returns:
            UpdateContextResult containing memories added to context
        """
        # Get messages from context
        messages = await model_context.get_messages()
        if not messages:
            return UpdateContextResult(memories=MemoryQueryResult(results=[]))
        
        # Use the last message as query
        last_message = messages[-1]
        query_text = (
            last_message.content
            if isinstance(last_message.content, str)
            else str(last_message)
        )
        
        # Query memory
        query_results = await self.query(query_text)
        
        # Add relevant memories to context
        if query_results.results:
            memory_strings = [
                f"{i}. {self._extract_text(memory)}"
                for i, memory in enumerate(query_results.results, 1)
            ]
            memory_context = (
                "\nRelevant context from memory:\n" + 
                "\n".join(memory_strings)
            )
            
            # Add as system message
            await model_context.add_message(
                SystemMessage(content=memory_context)
            )
        
        return UpdateContextResult(memories=query_results)
    
    async def clear(self) -> None:
        """Clear all memory entries for this user."""
        with self._lock:
            # Filter to only clear this user's entries
            user_prefix = f"mem_{self._user_id}_"
            to_remove = [
                cid for cid in self._items.keys()
                if cid.startswith(user_prefix) or self._user_id == "default"
            ]
            
            for cid in to_remove:
                self._items.pop(cid, None)
                self._created_at.pop(cid, None)
                self._expires_at.pop(cid, None)
        
        # Also clear from MCAL if possible
        try:
            if hasattr(self._mcal, 'clear_user'):
                self._mcal.clear_user(self._user_id)
            elif hasattr(self._mcal, 'reset'):
                # Only reset if default user (affects all)
                if self._user_id == "default":
                    self._mcal.reset()
        except Exception:
            pass
    
    async def close(self) -> None:
        """Clean up resources."""
        # Clear in-memory data
        with self._lock:
            self._items.clear()
            self._created_at.clear()
            self._expires_at.clear()
    
    # -------------------------------------------------------------------------
    # Additional helper methods
    # -------------------------------------------------------------------------
    
    def get_item_count(self) -> int:
        """Get the number of non-expired items in memory."""
        self._cleanup_expired()
        with self._lock:
            return len(self._items)
    
    def get_all_items(self) -> List[MemoryContent]:
        """Get all non-expired items (for debugging)."""
        self._cleanup_expired()
        with self._lock:
            return list(self._items.values())
    
    async def add_text(
        self,
        text: str,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        """
        Convenience method to add plain text to memory.
        
        Args:
            text: Text content to add
            metadata: Optional metadata
        """
        content = MemoryContent(
            content=text,
            mime_type=MemoryMimeType.TEXT,
            metadata=metadata,
        )
        await self.add(content)
    
    async def query_text(
        self,
        query: str,
        max_results: Optional[int] = None,
    ) -> List[str]:
        """
        Convenience method to query and return plain text results.
        
        Args:
            query: Search query
            max_results: Optional max results override
        
        Returns:
            List of text strings from matching memories
        """
        kwargs = {}
        if max_results is not None:
            kwargs['max_results'] = max_results
        
        results = await self.query(query, **kwargs)
        return [self._extract_text(r) for r in results.results]


__all__ = ["MCALMemory"]
