"""
MCALMemory: Memory nodes for LangGraph workflows.

Provides goal-aware memory that preserves reasoning context.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Callable, Dict, List, Optional, Sequence, TYPE_CHECKING

from mcal_langgraph._compat import check_langgraph, BaseMessage
from mcal_langgraph.store import MCALStore

if TYPE_CHECKING:
    from mcal import MCAL


@dataclass
class MCALMemoryConfig:
    """Configuration for MCAL memory in LangGraph."""
    llm_provider: str = "anthropic"
    embedding_provider: str = "openai"
    storage_path: Optional[str] = None
    user_id: str = "default"
    include_goals: bool = True
    include_decisions: bool = True
    max_context_tokens: int = 4000


class MCALMemory:
    """
    MCAL Memory component for LangGraph workflows.
    
    Provides goal-aware memory that preserves reasoning context
    across agent interactions.
    
    Usage:
        from mcal_langgraph import MCALMemory
        from langgraph.graph import StateGraph
        
        memory = MCALMemory(llm_provider="anthropic")
        
        # Add as a node in your graph
        graph = StateGraph(...)
        graph.add_node("update_memory", memory.update_node())
        graph.add_node("get_context", memory.context_node())
    """
    
    def __init__(
        self,
        mcal: Optional["MCAL"] = None,
        llm_provider: str = "anthropic",
        embedding_provider: str = "openai",
        storage_path: Optional[str] = None,
        user_id: str = "default",
        **mcal_kwargs
    ):
        check_langgraph()
        
        self.user_id = user_id
        
        if mcal is not None:
            # Use provided MCAL instance (Issue #106)
            self._mcal = mcal
        else:
            # Use singleton factory (Issue #106)
            from mcal import get_mcal
            self._mcal = get_mcal(
                storage_path=storage_path,
                llm_provider=llm_provider,
                **mcal_kwargs
            )
    
    @property
    def mcal(self) -> "MCAL":
        """Get the underlying MCAL instance."""
        return self._mcal
    
    def as_store(self) -> MCALStore:
        """Get a LangGraph BaseStore backed by this memory."""
        return MCALStore(self._mcal)
    
    def _convert_messages(self, messages: Sequence[BaseMessage]) -> List[Dict[str, str]]:
        """Convert LangChain messages to MCAL format."""
        result = []
        for msg in messages:
            if hasattr(msg, 'type'):
                if msg.type == "human":
                    result.append({"role": "user", "content": str(msg.content)})
                elif msg.type == "ai":
                    result.append({"role": "assistant", "content": str(msg.content)})
                else:
                    result.append({"role": "user", "content": str(msg.content)})
            else:
                result.append({"role": "user", "content": str(msg.content)})
        return result
    
    async def add(self, messages: Sequence[BaseMessage]) -> Dict[str, Any]:
        """
        Add messages to MCAL memory.
        
        Args:
            messages: LangChain message sequence
            
        Returns:
            Extraction result with goals and decisions
        """
        mcal_messages = self._convert_messages(messages)
        result = await self._mcal.add(mcal_messages, user_id=self.user_id)
        return {
            "goals": result.unified_graph.get_active_goals() if result.unified_graph else [],
            "decisions": result.unified_graph.get_all_decisions_with_detail() if result.unified_graph else [],
            "node_count": result.unified_graph.node_count if result.unified_graph else 0,
        }
    
    async def get_context(self, query: str, max_tokens: int = 4000) -> str:
        """
        Get goal-aware context for a query.
        
        Args:
            query: The current query/task
            max_tokens: Maximum context tokens
            
        Returns:
            Formatted context string
        """
        return await self._mcal.get_context(
            query=query,
            user_id=self.user_id,
            max_tokens=max_tokens
        )
    
    async def search(self, query: str, top_k: int = 10) -> List[Dict[str, Any]]:
        """
        Search memory for relevant information.
        
        Args:
            query: Search query
            top_k: Number of results
            
        Returns:
            List of search results
        """
        results = await self._mcal.search(query, user_id=self.user_id, limit=top_k)
        if not results:
            return []
        return [
            {
                "id": mem.id,
                "content": mem.content,
                "score": mem.score,
                "metadata": mem.metadata or {},
            }
            for mem in results.memories
        ]
    
    def update_node(self) -> Callable:
        """
        Create a LangGraph node that updates memory.
        
        Usage:
            graph.add_node("update_memory", memory.update_node())
        
        Returns:
            Async function suitable for add_node()
        """
        async def _update_memory(state: Dict[str, Any]) -> Dict[str, Any]:
            messages = state.get("messages", [])
            if messages:
                result = await self.add(messages)
                return {
                    "memory_updated": True,
                    "active_goals": result.get("goals", []),
                    "decisions": result.get("decisions", []),
                }
            return {"memory_updated": False}
        
        return _update_memory
    
    def context_node(self, query_key: str = "query") -> Callable:
        """
        Create a LangGraph node that retrieves context.
        
        Usage:
            graph.add_node("get_context", memory.context_node())
        
        Args:
            query_key: State key containing the query
            
        Returns:
            Async function suitable for add_node()
        """
        async def _get_context(state: Dict[str, Any]) -> Dict[str, Any]:
            query = state.get(query_key, "")
            if query:
                context = await self.get_context(query)
                return {"memory_context": context}
            return {"memory_context": ""}
        
        return _get_context


__all__ = ["MCALMemory", "MCALMemoryConfig"]
