# References:

# https://docs.crewai.com/introduction
# https://ai.google.dev/gemini-api/docs

import os
from crewai import Agent, Crew, Task
from crewai.agents.agent_builder.base_agent import BaseAgent
from crewai.project import CrewBase, agent, crew, task
from google import genai
from openinference.instrumentation.crewai import CrewAIInstrumentor
from phoenix.otel import register
from tools.ai_tools import AITools
from tools.arithmetic_tools import ArithmeticTools
from typing import List
from utils import read_file_json, is_ext

## LLMs

MANAGER_MODEL      = "gpt-4.5-preview"
AGENT_MODEL        = "gpt-4.1-mini"

FINAL_ANSWER_MODEL = "gemini-2.5-pro-preview-03-25"

# LLM evaluation

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)

@CrewBase
class GAIACrew():
    agents: List[BaseAgent]
    tasks: List[Task]

    @agent
    def web_search_agent(self) -> Agent:
        return Agent(
            config=self.agents_config["web_search_agent"],
            allow_delegation=False,
            llm=AGENT_MODEL,
            max_iter=2,
            tools=[AITools.web_search_tool],
            verbose=True
        )

    @agent
    def web_browser_agent(self) -> Agent:
        return Agent(
            config=self.agents_config["web_browser_agent"],
            allow_delegation=False,
            llm=AGENT_MODEL,
            max_iter=3,
            tools=[AITools.web_browser_tool],
            verbose=True
        )
    
    @agent
    def image_analysis_agent(self) -> Agent:
        return Agent(
            config=self.agents_config["image_analysis_agent"],
            allow_delegation=False,
            llm=AGENT_MODEL,
            max_iter=2,
            tools=[AITools.image_analysis_tool],
            verbose=True
        )

    @agent
    def audio_analysis_agent(self) -> Agent:
        return Agent(
            config=self.agents_config["audio_analysis_agent"],
            allow_delegation=False,
            llm=AGENT_MODEL,
            max_iter=2,
            tools=[AITools.audio_analysis_tool],
            verbose=True
        )

    @agent
    def video_analysis_agent(self) -> Agent:
        return Agent(
            config=self.agents_config["video_analysis_agent"],
            allow_delegation=False,
            llm=AGENT_MODEL,
            max_iter=2,
            tools=[AITools.video_analysis_tool],
            verbose=True
        )

    @agent
    def youtube_analysis_agent(self) -> Agent:
        return Agent(
            config=self.agents_config["youtube_analysis_agent"],
            allow_delegation=False,
            llm=AGENT_MODEL,
            max_iter=2,
            tools=[AITools.youtube_analysis_tool],
            verbose=True
        )

    @agent
    def document_analysis_agent(self) -> Agent:
        return Agent(
            config=self.agents_config["document_analysis_agent"],
            allow_delegation=False,
            llm=AGENT_MODEL,
            max_iter=2,
            tools=[AITools.document_analysis_tool],
            verbose=True
        )

    @agent
    def arithmetic_agent(self) -> Agent:
        return Agent(
            config=self.agents_config["document_analysis_agent"],
            allow_delegation=False,
            llm=AGENT_MODEL,
            max_iter=2,
            tools=[ArithmeticTools.add, ArithmeticTools.subtract, ArithmeticTools.multiply, ArithmeticTools.divide, ArithmeticTools.modulus],
            verbose=True
        )

    @agent
    def code_generation_agent(self) -> Agent:
        return Agent(
            config=self.agents_config["code_generation_agent"],
            allow_delegation=False,
            llm=AGENT_MODEL,
            max_iter=3,
            tools=[AITools.code_generation_tool],
            verbose=True
        )

    @agent
    def code_execution_agent(self) -> Agent:
        return Agent(
            config=self.agents_config["code_execution_agent"],
            allow_delegation=False,
            llm=AGENT_MODEL,
            max_iter=3,
            tools=[AITools.code_execution_tool],
            verbose=True
        )

    @agent
    def manager_agent(self) -> Agent:
        return Agent(
            config=self.agents_config["manager_agent"],
            allow_delegation=True,
            llm=MANAGER_MODEL,
            max_iter=5,
            verbose=True
        )

    @task
    def manager_task(self) -> Task:
        return Task(
            config=self.tasks_config["manager_task"]
        )

    @crew
    def crew(self) -> Crew:
        return Crew(
            agents=self.agents,
            tasks=self.tasks,
            verbose=True
        )

def run_crew(question, file_path):
    final_question = question
    
    if file_path:
        if is_ext(file_path, ".csv") or is_ext(file_path, ".xls") or is_ext(file_path, ".xlsx") or is_ext(file_path, ".json") or is_ext(file_path, ".jsonl"):
            json_data = read_file_json(file_path)
            final_question = f"{question} JSON data:\n{json_data}."
        else:
            final_question = f"{question} File path: {file_path}."

    answer = GAIACrew().crew().kickoff(inputs={"question": final_question})
    final_answer = get_final_answer(FINAL_ANSWER_MODEL, question, str(answer))

    print(f"=> Initial question: {question}")
    print(f"=> Final question: {final_question}")
    print(f"=> Initial answer: {answer}")
    print(f"=> Final answer: {final_answer}")
    
    return final_answer

def get_final_answer(model, question, answer):
    prompt_template = """
        You are an expert question answering assistant. Given a question and an initial answer, your task is to provide the final answer.
        Your final answer must be a number and/or string 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 USD, $, percent, or % unless specified otherwise.
        If you are asked for a string, don't use articles, neither abbreviations (for example 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.
        If the final answer is a number, use a number not a word.
        If the final answer is a string, start with an uppercase character.
        If the final answer is a comma-separated list of numbers, use a space character after each comma.
        If the final answer is a comma-separated list of strings, use a space character after each comma and start with a lowercase character.
        Do not add any content to the final answer that is not in the initial answer.
        **Question:** """ + question + """
        
        **Initial answer:** """ + answer + """
        
        **Example 1:** What is the biggest city in California? Los Angeles
        **Example 2:** How many 'r's are in strawberry? 3
        **Example 3:** What is the opposite of black? White
        **Example 4:** What are the first 5 numbers in the Fibonacci sequence? 0, 1, 1, 2, 3
        **Example 5:** What is the opposite of bad, worse, worst? good, better, best
        
        **Final answer:** 
        """

    client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])

    response = client.models.generate_content(
        model=model, 
        contents=[prompt_template]
    )
    
    return response.text