#!/usr/bin/env python3
"""
LangChain tools for Long Term Memory integration with Ollama
"""

from langchain.tools import tool
from langchain_ollama import OllamaLLM
from langchain.agents import create_react_agent, AgentExecutor
from langchain import hub
from typing import Optional, Dict, Any, List
import requests
import json

# Import your existing LTM demo class  
from gradio_demo import LongTermMemoryDemo

# Initialize shared memory instance
ltm = LongTermMemoryDemo()

@tool
def save_memory(content: str, title: str, tags: str = "", context: str = "") -> str:
    """
    Save important insights, conclusions, or context to long-term memory.
    Use this to remember key information from conversations that might be useful later.
    
    Args:
        content: The insight or information to save
        title: A brief descriptive title
        tags: Optional comma-separated tags
        context: Optional additional context
    """
    try:
        return ltm.save_memory(content, title, tags, context)
    except Exception as e:
        return f"Error saving memory: {str(e)}"

@tool
def search_memory(query: str, limit: int = 5, threshold: float = 0.3) -> str:
    """
    Search through long-term memory for relevant information.
    Use this to find previously saved insights or context related to current discussion.
    
    Args:
        query: What to search for
        limit: Max number of results (default: 5)
        threshold: Similarity threshold 0-1 (default: 0.3)
    """
    try:
        return ltm.search_memory(query, limit, threshold)
    except Exception as e:
        return f"Error searching memory: {str(e)}"

@tool
def list_memories(limit: int = 10) -> str:
    """
    List all stored memories to see what information is available.
    Useful for getting an overview of stored knowledge.
    
    Args:
        limit: Maximum number of memories to show (default: 10)
    """
    try:
        return ltm.list_memories(limit)
    except Exception as e:
        return f"Error listing memories: {str(e)}"

@tool
def memory_stats() -> str:
    """
    Get statistics about stored memories.
    Shows total count, tags, and other metadata.
    """
    try:
        return ltm.get_memory_stats()
    except Exception as e:
        return f"Error getting stats: {str(e)}"

# Example usage with Ollama
def create_memory_enabled_agent(model_name: str = "llama3.2"):
    """Create a LangChain agent with memory capabilities"""
    
    # Initialize Ollama LLM
    llm = OllamaLLM(model=model_name)
    
    # Create tools list
    tools = [save_memory, search_memory, list_memories, memory_stats]
    
    # Get the react prompt from hub
    try:
        prompt = hub.pull("hwchase17/react")
    except:
        # Fallback prompt if hub is not available
        from langchain.prompts import PromptTemplate
        
        template = """Answer the following questions as best you can. You have access to the following tools:

{tools}

Use the following format:

Question: the input question you must answer
Thought: you should always think about what to do
Action: the action to take, should be one of [{tool_names}]
Action Input: the input to the action
Observation: the result of the action
... (this Thought/Action/Action Input/Observation can repeat N times)
Thought: I now know the final answer
Final Answer: the final answer to the original input question

Begin!

Question: {input}
Thought:{agent_scratchpad}"""
        
        prompt = PromptTemplate.from_template(template)
    
    # Create agent
    agent = create_react_agent(llm, tools, prompt)
    
    # Create agent executor
    agent_executor = AgentExecutor(
        agent=agent,
        tools=tools,
        verbose=True,
        handle_parsing_errors=True,
        max_iterations=10
    )
    
    return agent_executor

# Example conversation loop
def main():
    """Example usage"""
    print("🧠 Initializing Memory-Enabled Agent with Ollama...")
    
    try:
        agent = create_memory_enabled_agent("llama3.2")  # или любая другая модель в Ollama
        
        print("✅ Agent ready! Type 'quit' to exit.")
        print("💡 Try commands like:")
        print("   - 'Save this insight: quantum computers might revolutionize AI with title Quantum AI and tags quantum,ai,future' ")
        print("   - 'Search my memories for information about quantum computing'")
        print("   - 'What memories do I have stored?'")
        print("   - 'Show me memory statistics'")
        print()
        
        while True:
            try:
                user_input = input("You: ").strip()
                
                if user_input.lower() in ['quit', 'exit', 'bye']:
                    print("Goodbye!")
                    break
                
                if not user_input:
                    continue
                
                # Run the agent
                response = agent.invoke({"input": user_input})
                print(f"Agent: {response['output']}")
                print()
                
            except KeyboardInterrupt:
                print("\nGoodbye!")
                break
            except Exception as e:
                print(f"Error: {e}")
                continue
                
    except Exception as e:
        print(f"Failed to initialize agent: {e}")
        print("Make sure Ollama is running and the model is available.")
        print("Try: ollama pull llama3.2")

if __name__ == "__main__":
    main()