import os

from crewai import Agent, Crew, Process, Task
from crewai.tools import tool
from crewai_tools import (
    CodeInterpreterTool,
    SerperDevTool,
    VisionTool,
    WebsiteSearchTool,
    YoutubeVideoSearchTool
)
from openinference.instrumentation.crewai import CrewAIInstrumentor
from phoenix.otel import register
from util import get_final_answer

MANAGER_MODEL = "gpt-4.1"
AGENT_MODEL = "gpt-4.1"

PHOENIX_API_KEY = os.environ["PHOENIX_API_KEY"]

os.environ["PHOENIX_CLIENT_HEADERS"] = f"api_key={PHOENIX_API_KEY}"
os.environ["PHOENIX_COLLECTOR_ENDPOINT"] = "https://app.phoenix.arize.com"

tracer_provider = register(
    auto_instrument=True,
    project_name="gaia"
)

CrewAIInstrumentor().instrument(tracer_provider=tracer_provider)

def run_crew(question, file_name):
    # Tools
    
    image_analysis_tool = VisionTool()
    python_coding_tool = CodeInterpreterTool()
    web_search_tool = SerperDevTool()
    web_rag_tool = WebsiteSearchTool()
    youtube_analysis_tool = YoutubeVideoSearchTool()

    # Agents

    image_analysis_agent = Agent(
        role="Image Analysis Agent",
        goal="Analyze image to help answer question \"{topic}\"",
        backstory="As an expert image analysis assistant, you analyze the image to help answer the question.",
        allow_delegation=False,
        llm=AGENT_MODEL,
        max_iter=3,
        tools=[image_analysis_tool],
        verbose=False
    )
    
    python_coding_agent = Agent(
        role="Python Coding Agent",
        goal="Write and/or execute Python code to help answer question \"{topic}\"",
        backstory="As an expert Python coding assistant, you write and/or execute Python code to help answer the question.",
        allow_delegation=False,
        llm=AGENT_MODEL,
        max_iter=5,
        tools=[python_coding_tool],
        verbose=False
    )
   
    web_search_agent = Agent(
        role="Web Search Agent",
        goal="Search the web to help answer question \"{topic}\", then scrape the most relevant web page.",
        backstory="As an expert web search assistant, you search the web to help answer the question.",
        allow_delegation=False,
        llm=AGENT_MODEL,
        max_iter=3,
        tools=[web_search_tool, web_rag_tool],
        verbose=False
    )

    youtube_analysis_agent = Agent(
        role="YouTube Analysis Agent",
        goal="Analyze YouTube video to help answer question \"{topic}\"",
        backstory="As an expert YouTube video analysis assistant, you analyze the video to help answer the question.",
        allow_delegation=False,
        llm=AGENT_MODEL,
        max_iter=3,
        tools=[youtube_analysis_tool],
        verbose=False
    )

    manager_agent = Agent(
        role="Manager Agent",
        goal="Try to answer the following question. If needed, delegate to **one** of your coworkers, Image Analysis Agent, Python Coding Agent, Web Search Agent, or YouTube Analysis Agent for help. "
             "If there is no good coworker, delegate to the Python Coding Agent to implement a tool for the task. "
             "Question: \"{topic}\"",
        backstory="As an expert manager assistant, you answer the question.",
        allow_delegation=True,
        llm=MANAGER_MODEL,
        max_iter=5,
        verbose=True
    )

    # Tasks
  
    manager_task = Task(
        agent=manager_agent,
        description="Try to answer the following question. If needed, delegate to **one** of your coworkers, Image Analysis Agent, Python Coding Agent, Web Search Agent, or YouTube Analysis Agent for help. Question: \"{topic}\"",
        expected_output="The answer to the question."
    )
    
    # Crew
    
    crew = Crew(
        agents=[image_analysis_agent, python_coding_agent, web_search_agent, youtube_analysis_agent],
        manager_agent=manager_agent,
        tasks=[manager_task],
        verbose=True
    )

    # Processing
    
    if file_name:
        question = f"{question} File name: data/{file_name}."
        if file_name.endswith(".py") or file_name.endswith(".xlsx"):
            with open(f"data/{file_name}", "r") as file:
                question = f"{question} File data:\n{file.read()}"
    
    print(f"Question: {question}")
    initial_answer = crew.kickoff(inputs={"topic": question})
    print(f"Initial answer: {initial_answer}")
    final_answer = get_final_answer(question, str(initial_answer))
    print(f"Final answer: {final_answer}")
    
    return final_answer