"""
Test script to demonstrate using the ManagerAgent's research function
"""
import os
from dotenv import load_dotenv
from llama_index.llms.anthropic import Anthropic
from manager_agent import ManagerAgent
from models import TaskPrompt

# Load environment variables
load_dotenv()

# ANSI color codes for prettier output
COLOR_RESET = "\033[0m"
COLOR_CYAN = "\033[96m"
COLOR_GREEN = "\033[92m"
COLOR_YELLOW = "\033[93m"
COLOR_RED = "\033[91m"

def color_text(text, color):
    return f"{color}{text}{COLOR_RESET}"

def main():
    # Check if API key is available
    api_key = os.environ.get("ANTHROPIC_API_KEY")
    if not api_key:
        print(color_text("Error: ANTHROPIC_API_KEY not found in environment variables.", COLOR_RED))
        print("Please set your Anthropic API key with:")
        print("  export ANTHROPIC_API_KEY='your-api-key'")
        print("  or create a .env file with ANTHROPIC_API_KEY=your-api-key")
        return
    
    # Initialize LLM
    print(color_text("Initializing Anthropic Claude...", COLOR_CYAN))
    llm = Anthropic(model="claude-3-5-sonnet-20241022", api_key=api_key)
    
    # Initialize ManagerAgent
    print(color_text("Creating ManagerAgent...", COLOR_CYAN))
    manager = ManagerAgent(llm=llm)
    
    print(color_text("\nTest 1: Using research function directly", COLOR_GREEN))
    query = "What are the latest developments in AI agents and autonomous systems?"
    print(color_text(f"Research Query: {query}", COLOR_YELLOW))
    
    # Call research function directly
    report = manager.research(query=query, verbose=True)
    print(color_text("\n=== Research Report ===", COLOR_GREEN))
    print(report)
    
    print(color_text("\nTest 2: Using research as a tool through the agent", COLOR_GREEN))
    prompt_text = "I need a comprehensive report on recent developments in quantum computing. Please research this topic thoroughly."
    print(color_text(f"User Prompt: {prompt_text}", COLOR_YELLOW))
    
    # Create task prompt
    task_prompt = TaskPrompt(text=prompt_text)
    
    # Run through agent
    response = manager.run_task(task_prompt)
    print(color_text("\n=== Agent Response ===", COLOR_GREEN))
    print(response)

if __name__ == "__main__":
    main()