"""
LangChain Integration for Engram

Provides EngramMemory as a LangChain-compatible memory backend.

Usage:
    pip install engram[langchain]
    
    from engram.langchain import EngramMemory
    from langchain.chains import ConversationChain
    from langchain.llms import OpenAI
    
    memory = EngramMemory(brain_path="memories.engram")
    chain = ConversationChain(llm=OpenAI(), memory=memory)
"""

from __future__ import annotations

from pathlib import Path
from typing import Any, Dict, List, Optional, Union

try:
    from langchain_core.memory import BaseMemory
    from langchain_core.messages import AIMessage, BaseMessage, HumanMessage
except ImportError:
    raise ImportError(
        "LangChain is required for this module. "
        "Install with: pip install engram[langchain]"
    )

from .core import MemoryTree
from .types import ContentType, FileMetadata


class EngramMemory(BaseMemory):
    """
    LangChain memory backend using Engram for persistent storage.
    
    Stores conversation history as engram nodes with embeddings for semantic search.
    
    Args:
        brain_path: Path to the .engram file (created if doesn't exist)
        memory_key: Key used to pass memories to the chain (default: "history")
        input_key: Key for human input (default: "input")
        output_key: Key for AI output (default: "output")
        return_messages: Return as Message objects vs string (default: False)
        max_history: Maximum messages to return (default: 10)
        
    Example:
        >>> memory = EngramMemory(brain_path="chat.engram")
        >>> memory.save_context({"input": "Hello"}, {"output": "Hi there!"})
        >>> memory.load_memory_variables({})
        {'history': 'Human: Hello\\nAI: Hi there!'}
    """
    
    brain_path: Union[str, Path]
    memory_key: str = "history"
    input_key: str = "input"
    output_key: str = "output"
    return_messages: bool = False
    max_history: int = 10
    
    _tree: Optional[MemoryTree] = None
    
    class Config:
        arbitrary_types_allowed = True
        underscore_attrs_are_private = True
    
    def __init__(self, **kwargs: Any):
        super().__init__(**kwargs)
        self._load_or_create_tree()
    
    def _load_or_create_tree(self) -> None:
        """Load existing brain or create new one."""
        path = Path(self.brain_path)
        if path.exists():
            self._tree = MemoryTree.from_file(path)
        else:
            self._tree = MemoryTree()
    
    @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 conversation history.
        
        Returns recent messages formatted as string or Message objects.
        """
        if self._tree is None:
            return {self.memory_key: "" if not self.return_messages else []}
        
        # Get conversation nodes sorted by creation time
        conv_nodes = self._tree.find_by_tag("conversation")
        conv_nodes.sort(key=lambda n: n.temporal.created if n.temporal else 0)
        
        # Take last max_history messages
        recent = conv_nodes[-self.max_history:] if len(conv_nodes) > self.max_history else conv_nodes
        
        if self.return_messages:
            messages: List[BaseMessage] = []
            for node in recent:
                if node.content and isinstance(node.content.data, str):
                    role = "human" if node.metadata and "human" in node.metadata.tags else "ai"
                    if role == "human":
                        messages.append(HumanMessage(content=node.content.data))
                    else:
                        messages.append(AIMessage(content=node.content.data))
            return {self.memory_key: messages}
        else:
            lines = []
            for node in recent:
                if node.content and isinstance(node.content.data, str):
                    role = "Human" if node.metadata and "human" in node.metadata.tags else "AI"
                    lines.append(f"{role}: {node.content.data}")
            return {self.memory_key: "\n".join(lines)}
    
    def save_context(self, inputs: Dict[str, Any], outputs: Dict[str, str]) -> None:
        """
        Save a conversation turn to memory.
        
        Args:
            inputs: Dict containing the human input
            outputs: Dict containing the AI output
        """
        if self._tree is None:
            self._tree = MemoryTree()
        
        # Save human input
        human_input = inputs.get(self.input_key, "")
        if human_input:
            self._tree.add(
                human_input,
                content_type=ContentType.TEXT,
                tags=["conversation", "human"],
            )
        
        # Save AI output
        ai_output = outputs.get(self.output_key, "")
        if ai_output:
            self._tree.add(
                ai_output,
                content_type=ContentType.TEXT,
                tags=["conversation", "ai"],
            )
        
        # Persist to disk
        self._save()
    
    def _save(self) -> None:
        """Save the tree to disk."""
        if self._tree:
            self._tree.save(
                self.brain_path,
                metadata=FileMetadata(
                    source="langchain",
                    description="LangChain conversation memory",
                    tags=["langchain", "conversation"],
                ),
            )
    
    def clear(self) -> None:
        """Clear all conversation history."""
        self._tree = MemoryTree()
        self._save()
    
    # ============== Extended Features ==============
    
    def search(self, query: str, top_k: int = 5) -> List[str]:
        """
        Search conversation history by text.
        
        Args:
            query: Search query
            top_k: Maximum results to return
            
        Returns:
            List of matching message contents
        """
        if self._tree is None:
            return []
        
        results = self._tree.search_text(query)
        contents = []
        for node in results[:top_k]:
            if node.content and isinstance(node.content.data, str):
                contents.append(node.content.data)
        return contents
    
    def add_knowledge(self, content: str, tags: Optional[List[str]] = None) -> str:
        """
        Add knowledge to the memory (not conversation).
        
        Useful for RAG-style knowledge injection.
        
        Args:
            content: The knowledge content
            tags: Optional tags for categorization
            
        Returns:
            The node ID
        """
        if self._tree is None:
            self._tree = MemoryTree()
        
        node = self._tree.add(
            content,
            content_type=ContentType.TEXT,
            tags=["knowledge"] + (tags or []),
        )
        self._save()
        return node.id
    
    def get_knowledge(self, tag: Optional[str] = None) -> List[str]:
        """
        Retrieve knowledge from memory.
        
        Args:
            tag: Optional tag to filter by
            
        Returns:
            List of knowledge contents
        """
        if self._tree is None:
            return []
        
        if tag:
            nodes = [n for n in self._tree.find_by_tag("knowledge") 
                     if n.metadata and tag in n.metadata.tags]
        else:
            nodes = self._tree.find_by_tag("knowledge")
        
        return [
            n.content.data for n in nodes 
            if n.content and isinstance(n.content.data, str)
        ]
    
    @property
    def tree(self) -> Optional[MemoryTree]:
        """Access the underlying MemoryTree for advanced operations."""
        return self._tree
