"""
agent_memory.py
---------------
Agent framework integrations for HIEROMEM.

Provides compatible memory interfaces for LangChain, CrewAI, and other
agent frameworks, enabling HIEROMEM to serve as the persistent memory
layer for AI agents.

Author: Antigravity
"""

from __future__ import annotations

from dataclasses import dataclass, field
from datetime import datetime
from typing import Any, Dict, List, Optional
import networkx as nx

# Lazy imports for optional dependencies
_langchain_available = None
_crewai_available = None


def _check_langchain():
    global _langchain_available
    if _langchain_available is None:
        try:
            from langchain.memory.chat_memory import BaseChatMemory  # noqa: F401
            _langchain_available = True
        except ImportError:
            _langchain_available = False
    return _langchain_available


def _check_crewai():
    global _crewai_available
    if _crewai_available is None:
        try:
            import crewai  # noqa: F401
            _crewai_available = True
        except ImportError:
            _crewai_available = False
    return _crewai_available


@dataclass
class AgentAction:
    """Recorded agent action/tool call."""
    action_type: str  # "tool_call", "reasoning", "response"
    name: str
    input: Any
    output: Any
    timestamp: datetime = field(default_factory=datetime.now)
    metadata: Dict[str, Any] = field(default_factory=dict)


class AgentMemory:
    """
    Agent-compatible memory wrapper for HIEROMEM.

    Provides integration with popular agent frameworks:
    - LangChain (BaseChatMemory interface)
    - CrewAI (memory protocol)
    - Direct API for custom agents

    Example:
        from hieromem import AgentMemory

        # Direct usage
        memory = AgentMemory(persist=True)
        memory.add_message("user", "What's the weather?")
        memory.add_message("assistant", "It's sunny today.")

        # LangChain integration
        chain = ConversationChain(memory=memory.as_langchain())

        # CrewAI integration
        agent = Agent(memory=memory.as_crewai())
    """

    def __init__(
        self,
        persist: bool = True,
        user_id: Optional[str] = None,
        session_id: Optional[str] = None,
        track_tool_calls: bool = True,
        track_reasoning: bool = True,
        memory_kwargs: Optional[Dict[str, Any]] = None,
    ):
        self.persist = persist
        self.user_id = user_id
        self.session_id = session_id or datetime.now().strftime("%Y%m%d_%H%M%S")
        self.track_tool_calls = track_tool_calls
        self.track_reasoning = track_reasoning

        # Initialize HIEROMEM
        from hieromem_core.memory import Memory
        memory_kwargs = memory_kwargs or {}
        self._memory = Memory(**memory_kwargs)

        # Message buffer for conversation
        self._messages: List[Dict[str, Any]] = []
        self._actions: List[AgentAction] = []

    @property
    def memory(self):
        """Access underlying HIEROMEM Memory instance."""
        return self._memory

    def add_message(
        self,
        role: str,
        content: str,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        """
        Add a conversation message.

        Args:
            role: 'user', 'assistant', 'system', or 'tool'
            content: Message content
            metadata: Optional metadata dict
        """
        message = {
            "role": role,
            "content": content,
            "timestamp": datetime.now().isoformat(),
            "metadata": metadata or {}
        }
        self._messages.append(message)

        # Persist to HIEROMEM
        if self.persist:
            G = nx.DiGraph()
            G.add_node(role)
            if self.user_id:
                G.add_edge(self.user_id, role, relation="sent")

            self._memory.add_text_episode(
                content=content,
                G=G,
                meta={
                    "type": "message",
                    "role": role,
                    "user_id": self.user_id,
                    "session_id": self.session_id,
                    **(metadata or {})
                }
            )

    def add_tool_call(
        self,
        tool_name: str,
        tool_input: Any,
        tool_output: Any,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        """
        Record a tool call.

        Args:
            tool_name: Name of the tool
            tool_input: Input to the tool
            tool_output: Output from the tool
            metadata: Optional metadata
        """
        if not self.track_tool_calls:
            return

        action = AgentAction(
            action_type="tool_call",
            name=tool_name,
            input=tool_input,
            output=tool_output,
            metadata=metadata or {}
        )
        self._actions.append(action)

        # Persist
        if self.persist:
            G = nx.DiGraph()
            G.add_node(tool_name, type="tool")

            self._memory.add_text_episode(
                content=f"Tool: {tool_name}\nInput: {tool_input}\nOutput: {tool_output}",
                G=G,
                meta={
                    "type": "tool_call",
                    "tool_name": tool_name,
                    "user_id": self.user_id,
                    "session_id": self.session_id,
                    **(metadata or {})
                }
            )

    def add_reasoning(
        self,
        thought: str,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        """
        Record agent reasoning/thought.

        Args:
            thought: The reasoning text
            metadata: Optional metadata
        """
        if not self.track_reasoning:
            return

        action = AgentAction(
            action_type="reasoning",
            name="thought",
            input=thought,
            output=None,
            metadata=metadata or {}
        )
        self._actions.append(action)

        if self.persist:
            self._memory.add_text_episode(
                content=f"Thought: {thought}",
                meta={
                    "type": "reasoning",
                    "user_id": self.user_id,
                    "session_id": self.session_id,
                    **(metadata or {})
                }
            )

    def get_history(self, k: int = 10) -> List[Dict[str, Any]]:
        """Get recent message history."""
        return self._messages[-k:]

    def get_context(self, query: str, k: int = 5) -> str:
        """
        Get relevant context for a query.

        Args:
            query: The query to find context for
            k: Number of relevant memories to retrieve

        Returns:
            Formatted context string
        """
        result = self._memory.recall(query, k=k)

        context_parts = []
        for ep in result.episodes:
            content = ep.meta.get("content_preview", "")
            role = ep.meta.get("role", "memory")
            context_parts.append(f"[{role}]: {content}")

        return "\n".join(context_parts)

    def clear(self) -> None:
        """Clear message buffer (does not clear persistent memory)."""
        self._messages.clear()
        self._actions.clear()

    # ----------------------------------------------------------------
    # LangChain Integration
    # ----------------------------------------------------------------

    def as_langchain(self):
        """
        Return LangChain-compatible memory interface.

        Returns:
            BaseChatMemory-compatible object

        Raises:
            ImportError: If langchain is not installed
        """
        if not _check_langchain():
            raise ImportError(
                "langchain not installed. Run: pip install langchain"
            )

        return HieromemLangChainMemory(self)

    # ----------------------------------------------------------------
    # CrewAI Integration
    # ----------------------------------------------------------------

    def as_crewai(self):
        """
        Return CrewAI-compatible memory interface.

        Returns:
            CrewAI memory-compatible object

        Raises:
            ImportError: If crewai is not installed
        """
        if not _check_crewai():
            raise ImportError(
                "crewai not installed. Run: pip install crewai"
            )

        return HieromemCrewAIMemory(self)


class HieromemLangChainMemory:
    """
    LangChain BaseChatMemory implementation using HIEROMEM.

    Implements the interface expected by LangChain chains.
    """

    def __init__(self, agent_memory: AgentMemory):
        self._agent_memory = agent_memory
        self.memory_key = "history"
        self.input_key = "input"
        self.output_key = "output"
        self.return_messages = False

    @property
    def memory_variables(self) -> List[str]:
        """Return memory variables."""
        return [self.memory_key]

    def load_memory_variables(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
        """Load memory variables for chain."""
        # Get relevant context if there's input
        query = inputs.get(self.input_key, "")

        if query:
            context = self._agent_memory.get_context(query, k=5)
        else:
            # Return recent history
            history = self._agent_memory.get_history(k=10)
            context = "\n".join(
                f"{m['role']}: {m['content']}" for m in history
            )

        if self.return_messages:
            # Return as message objects (for chat models)
            try:
                from langchain.schema import HumanMessage, AIMessage
                messages = []
                for m in self._agent_memory.get_history(k=10):
                    if m["role"] == "user":
                        messages.append(HumanMessage(content=m["content"]))
                    else:
                        messages.append(AIMessage(content=m["content"]))
                return {self.memory_key: messages}
            except ImportError:
                pass

        return {self.memory_key: context}

    def save_context(self, inputs: Dict[str, Any], outputs: Dict[str, Any]) -> None:
        """Save context from chain interaction."""
        user_input = inputs.get(self.input_key, "")
        ai_output = outputs.get(self.output_key, "")

        if user_input:
            self._agent_memory.add_message("user", user_input)
        if ai_output:
            self._agent_memory.add_message("assistant", ai_output)

    def clear(self) -> None:
        """Clear memory."""
        self._agent_memory.clear()


class HieromemCrewAIMemory:
    """
    CrewAI memory implementation using HIEROMEM.

    Implements the interface expected by CrewAI agents.
    """

    def __init__(self, agent_memory: AgentMemory):
        self._agent_memory = agent_memory

    def save(self, key: str, value: str) -> None:
        """Save a memory."""
        self._agent_memory.add_message("memory", f"{key}: {value}")

    def search(self, query: str, k: int = 5) -> List[str]:
        """Search memories."""
        context = self._agent_memory.get_context(query, k=k)
        return context.split("\n") if context else []

    def reset(self) -> None:
        """Reset memory."""
        self._agent_memory.clear()
