import os

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

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 get_crew():
    search_tool = SerperDevTool()
    web_rag_tool = WebsiteSearchTool()
    
    research_agent = Agent(
        role="Web Research Agent",
        goal="Search the web for question \"{topic}\" and scrape the most relevant web page.",
        backstory="As an expert web research assistant, you are searching the web for question \"{topic}\" and scraping the most relevant web page.",
        allow_delegation=False,
        tools=[search_tool, web_rag_tool],
        verbose=True
    )

    answer_agent = Agent(
        role="Final Answer Agent",
        goal="Provide the final answer to question \"{topic}\".",
        backstory="As an expert question/answer assistant, you are providing the final answer to question \"{topic}\".",
        allow_delegation=False,
        tools=[final_answer_tool],
        verbose=True
    )

    manager_agent = Agent(
        role="Project Manager",
        goal="Efficiently manage the crew and ensure high-quality task completion.",
        backstory="You're an experienced project manager, skilled in overseeing complex projects and guiding teams to success.",
        allow_delegation=True
    )
    
    #research = Task(
    #    agent=research_agent,
    #    description="Search the web for topic \"{topic}\" and scrape the most relevant web page.",
    #    expected_output="Content on topic \"{topic}\"."
    #)

    #answer = Task(
    #    agent=answer_agent,
    #    description="Given topic \"{topic}\" and an initial answer, get the final answer.",
    #    expected_output="The final answer to topic \"{topic}\"."
    #)

    task = Task(
        description="Find the final answer to question \"{topic}\".",
        expected_output="The final answer to question \"{topic}\"."
    )
    
    return Crew(
        agents=[research_agent, answer_agent],
        manager_agent=manager_agent,
        planning=True,
        process=Process.hierarchical,
        #tasks=[research, answer],
        tasks=[task],
        verbose=True
    )

@tool("Final answer tool.")
def final_answer_tool(question: str, initial_answer: str) -> str:
    """Given a question and an initial answer, provide the final answer."""
    prompt_template = """
        You are given a question and an initial answer.
        Your final answer must a number OR as few words as possible OR a comma separated list of numbers and/or strings.
        If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise.
        If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise.
        If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
        **Question:** """ + question + """
        **Initial answer:** """ + initial_answer + """
        **Example:** What is the opposite of white? Black
        **Final answer:**:
        """

    client = OpenAI()
    completion = client.chat.completions.create(
        messages=[{"role": "user", "content": [{"type": "text", "text": prompt_template}]}],
        model="gpt-4o"
    )
        
    return completion.choices[0].message.content