import os

from crewai import Agent, Crew, Process, Task
from crewai.tools import tool
from crewai_tools import (
    SerperDevTool,
    WebsiteSearchTool
)
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"
)

###
#from opentelemetry import trace
#from opentelemetry.sdk.trace import TracerProvider
#from opentelemetry.sdk.trace.export import BatchSpanProcessor

from openinference.instrumentation.crewai import CrewAIInstrumentor
#from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
#from opentelemetry.sdk.trace.export import ConsoleSpanExporter, SimpleSpanProcessor

#endpoint = "https://app.phoenix.arize.com"
#trace_provider = TracerProvider()
#trace_provider.add_span_processor(SimpleSpanProcessor(OTLPSpanExporter(endpoint)))

CrewAIInstrumentor().instrument(tracer_provider=tracer_provider)
###

#agentops.init(os.environ["AGENTOPS_API_KEY"])

def get_crew():
    search_tool = SerperDevTool()
    web_rag_tool = WebsiteSearchTool()
    
    research_agent = Agent(
        role="Web Research Agent",
        goal="Search the web for topic \"{topic}\" and scrape the most relevant web page.",
        backstory="As an expert web research assistant, you are searching the web for topic \"{topic}\" and scraping the most relevant web page.",
        tools=[search_tool, web_rag_tool],
        verbose=True
    )

    answer_agent = Agent(
        role="Final Answer Agent",
        goal="Provide the final answer on topic \"{topic}\".",
        backstory="As an expert answer assistant, you are providing the final answer on topic \"{topic}\".",
        tools=[final_answer_tool],
        verbose=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}\"."
    )
    
    return Crew(
        agents=[research_agent, answer_agent],
        tasks=[research, answer],
        planning=True,
        verbose=True
    )

@tool("Final answer tool.")
def final_answer_tool(question: str, 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. You must **precisely** answer the question based on the initial answer.
        If asked for a number, only provide the number. Do not include explanations, steps, reasoning, or additional text.
        **Question:** """ + question + """
        **Initial answer:** """ + answer + """
        **Example 1:** How many states are in the USA? 50
        **Example 2:** What is the superlative of good? Best
        **Example 3:** What is the opposite of left? Right
        **Final answer:**:
        """

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