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 run_crew(question):
    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 research assistant, you search the web for question \"{topic}\" and scrape the most relevant web page. "
                  "Your output is the basis for the Final Answer Agent to provide a final answer to the question.",
        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 answering assistant, you provide the final answer to question \"{topic}\". "
                  "You base your final answer to the question on the output of the Web Research Agent.",
        allow_delegation=False,
        tools=[final_answer_tool],
        verbose=True
    )

    #manager_agent = Agent(
    #    role="Project Manager",
    #    goal="Use the Web Research Agent to help answer question \"{topic}\". Always delegate to the Final Answer Agent at the end.",
    #    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 question \"{topic}\" and scrape the most relevant web page.",
        expected_output="Content to help answer question \"{topic}\"."
    )

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

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

    answer = crew().kickoff(inputs={"topic": question})

    return final_answer(question, answer)
    
@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"
    )

    print("###")
    print(completion.choices[0].message.content)
    print("###")
    
    return completion.choices[0].message.content