import os
from langchain.chat_models import init_chat_model
from langchain_community.tools.shell.tool import ShellTool
from langgraph.prebuilt import create_react_agent
from langchain_core.messages import HumanMessage
from custom_tools.poem_tool import print_poem


def create_agent():
    """Create and return a ReAct agent with shell and poem capabilities."""
    
    # Initialize the chat model (using OpenAI GPT-4)
    # Make sure you have set your OPENAI_API_KEY environment variable
    llm = init_chat_model("openai:gpt-4o-mini")
    
    # Define the tools available to the agent
    shell_tool = ShellTool()
    tools = [shell_tool, print_poem]
    
    
    # Create a ReAct agent with simplified system prompt
    system_prompt = """You are a helpful assistant with access to shell commands and poem generation capabilities.
    
    You can:
    1. Execute shell commands using the shell tool to interact with the system
    2. Generate and print beautiful poems using the print_poem tool
    
    The poem tool can create different types of poems:
    - "nature": Poems about nature and the environment
    - "tech": Poems about technology and programming
    - "motivational": Inspirational and motivational poems
    - "friendship": Poems about friendship and relationships
    - "random": Randomly selects from any available poem type
    
    When helping users:
    - Be friendly and helpful
    - Use appropriate tools based on the user's request
    - Always be cautious with shell commands and explain what they do
    - For system monitoring commands like 'top', 'ps', etc., use appropriate flags to avoid hanging
    - When users ask for poems or want to brighten their day, use the poem tool
    """


    # Create the ReAct agent
    agent = create_react_agent(
        llm, 
        tools, 
        prompt=system_prompt
    )
    
    return agent


def run_agent_query(agent, user_input: str, conversation_history: list):
    """Run a simple agent query and display results cleanly."""
    # Create a human message
    message = HumanMessage(content=user_input)
    
    # Add the new message to conversation history
    conversation_history.append(message)
    
    # Use the agent's stream method for clean output like the LangChain tutorial
    for step in agent.stream({"messages": conversation_history}, stream_mode="values"):
        step["messages"][-1].pretty_print()
    
    # Add the agent's response to conversation history
    if step and "messages" in step:
        conversation_history.append(step["messages"][-1])



def main():
    # Check if required API keys are set
    if not os.getenv("OPENAI_API_KEY"):
        print("Please set your OPENAI_API_KEY environment variable.")
        print("You can set it by running: export OPENAI_API_KEY='your-api-key-here'")
        return
    
    print("🤖 LangGraph Simple Demo Agent")
    print("Type 'quit', 'exit', or 'q' to exit the chat.")
    print("Type 'help' or 'h' for help and examples.")
    print("Type 'clear' or 'reset' to clear conversation history.")
    print("⚠️  WARNING: This agent has shell access - use with caution!")
    print("🎭 Available capabilities:")
    print("   - Generate beautiful poems for any occasion")
    print("   - Execute shell commands for system tasks")
    print("   - Help with general assistance")
    print("-" * 60)
    
    # Create the agent
    try:
        agent = create_agent()
        print("✅ Simple Demo Agent initialized successfully!")
        print("💡 Try asking: 'Write me a nature poem'")
        print("💡 Or: 'Show me the current directory'")
        
    except Exception as e:
        print(f"❌ Error initializing agent: {e}")
        return
    
    # Start the chat loop
    conversation_history = []  # Initialize conversation history
    
    while True:
        try:
            user_input = input("\nUser: ")
            if user_input.lower() in ["quit", "exit", "q"]:
                print("👋 Goodbye!")
                break
            elif user_input.lower() in ["help", "h"]:
                print("\n🆘 Help:")
                print("Commands:")
                print("  - quit/exit/q: Exit the agent")
                print("  - help/h: Show this help")
                print("  - clear/reset: Clear conversation history")
                print("\nExample queries:")
                print("  - 'Write me a motivational poem'")
                print("  - 'Generate a tech poem'")
                print("  - 'Show me a random poem'")
                print("  - 'List files in current directory'")
                print("  - 'Check disk usage on the system'")
                continue
            elif user_input.lower() in ["clear", "reset"]:
                conversation_history = []
                print("🗑️ Conversation history cleared!")
                continue
            
            if user_input.strip():
                run_agent_query(agent, user_input, conversation_history)
            else:
                print("Please enter a message.")
                
        except KeyboardInterrupt:
            print("\n👋 Goodbye!")
            break
        except Exception as e:
            print(f"❌ Error: {e}")


if __name__ == "__main__":
    main()