# Automatic function calling
# https://ai.google.dev/gemini-api/docs/function-calling

from crewai.tools import tool
#from crewai_tools import StagehandTool
from google import genai
from google.genai import types

WEB_SEARCH_MODEL        = "gemini-2.5-flash-preview-04-17"
IMAGE_ANALYSIS_MODEL    = "gemini-2.5-flash-preview-04-17"
AUDIO_ANALYSIS_MODEL    = "gemini-2.5-flash-preview-04-17"
VIDEO_ANALYSIS_MODEL    = "gemini-2.5-flash-preview-04-17"
YOUTUBE_ANALYSIS_MODEL  = "gemini-2.5-flash-preview-04-17"
DOCUMENT_ANALYSIS_MODEL = "gemini-2.5-flash-preview-04-17"
ARITHMETIC_MODEL        = "gemini-2.5-flash-preview-04-17"
CODE_GENERATION_MODEL   = "gemini-2.5-flash-preview-04-17"
CODE_EXECUTION_MODEL    = "gemini-2.5-flash-preview-04-17"

@tool("Web Search Tool")
def web_search_tool(question: str) -> str:
    """Given a question only, search the web to answer the question.

       Args:
           question (str): Question to answer
            
       Returns:
           str: Answer to the question
            
       Raises:
           RuntimeError: If processing fails"""
    try:
        client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
        
        response = client.models.generate_content(
            model=WEB_SEARCH_MODEL,
            contents=question,
            config=types.GenerateContentConfig(
                tools=[types.Tool(google_search=types.GoogleSearchRetrieval())]
            )
        )

        return response.text
    except Exception as e:
        raise RuntimeError(f"Processing failed: {str(e)}")

@tool("Image Analysis Tool")
def image_analysis_tool(question: str, file_path: str) -> str:
    """Given a question and image file, analyze the image to answer the question.

       Args:
           question (str): Question about an image file
           file_path (str): The image file path
            
       Returns:
           str: Answer to the question about the image file
            
       Raises:
           RuntimeError: If processing fails"""
    try:
        client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
        
        file = client.files.upload(file=file_path)

        response = client.models.generate_content(
            model=IMAGE_ANALYSIS_MODEL,
            contents=[file, question]
        )
      
        return response.text
    except Exception as e:
        raise RuntimeError(f"Processing failed: {str(e)}")

@tool("Audio Analysis Tool")
def audio_analysis_tool(question: str, file_path: str) -> str:
    """Given a question and audio file, analyze the audio to answer the question.

       Args:
           question (str): Question about an audio file
           file_path (str): The audio file path
            
       Returns:
           str: Answer to the question about the audio file
            
       Raises:
           RuntimeError: If processing fails"""
    try:
        client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])

        file = client.files.upload(file=file_path)
        
        response = client.models.generate_content(
            model=AUDIO_ANALYSIS_MODEL, 
            contents=[file, question]
        )
      
        return response.text
    except Exception as e:
        raise RuntimeError(f"Processing failed: {str(e)}")

@tool("Video Analysis Tool")
def video_analysis_tool(question: str, file_path: str) -> str:
    """Given a question and video file, analyze the video to answer the question.

       Args:
           question (str): Question about a video file
           file_path (str): The video file path
            
       Returns:
           str: Answer to the question about the video file
            
       Raises:
           RuntimeError: If processing fails"""
    try:
        client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])

        file = client.files.upload(file=file_path)
        
        response = client.models.generate_content(
            model=VIDEO_ANALYSIS_MODEL, 
            contents=[file, question]
        )

        return response.text
    except Exception as e:
        raise RuntimeError(f"Processing failed: {str(e)}")
        
@tool("YouTube Analysis Tool")
def youtube_analysis_tool(question: str, url: str) -> str:
    """Given a question and YouTube URL, analyze the video to answer the question.

       Args:
           question (str): Question about a YouTube video
           url (str): The YouTube URL
            
       Returns:
           str: Answer to the question about the YouTube video
            
       Raises:
           RuntimeError: If processing fails"""
    try:
        client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])

        return client.models.generate_content(
            model=YOUTUBE_ANALYSIS_MODEL,
            contents=types.Content(
                parts=[types.Part(file_data=types.FileData(file_uri=url)),
                       types.Part(text=question)]
            )
        )
    except Exception as e:
        raise RuntimeError(f"Processing failed: {str(e)}")

@tool("Document Analysis Tool")
def document_analysis_tool(question: str, file_path: str) -> str:
    """Given a question and document file, analyze the document to answer the question.

       Args:
           question (str): Question about a document file
           file_path (str): The document file path
            
       Returns:
           str: Answer to the question about the document file
            
       Raises:
           RuntimeError: If processing fails"""
    try:
        client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])

        contents = []
        
        if is_ext(file_path, ".docx"):
            text_data = read_docx_text(file_path)
            contents = [f"{question}\n{text_data}"]
            print(f"=> Text data:\n{text_data}")
        elif is_ext(file_path, ".pptx"):
            text_data = read_pptx_text(file_path)
            contents = [f"{question}\n{text_data}"]
            print(f"=> Text data:\n{text_data}")
        else:
            file = client.files.upload(file=file_path)
            contents = [file, question]
        
        response = client.models.generate_content(
            model=DOCUMENT_ANALYSIS_MODEL,
            contents=contents
        )
      
        return response.text
    except Exception as e:
        raise RuntimeError(f"Processing failed: {str(e)}")

@tool("Arithmetic Tool")
def arithmetic_tool(question: str, a: float, b: float) -> float:
    """Given a question and two numbers, perform the calculation to answer the question.

       Args:
           question (str): Question to answer
           a (float): First number
           b (float): Second number
            
       Returns:
           float: Result number
            
       Raises:
           RuntimeError: If processing fails"""
    try:
        client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
        
        response = client.models.generate_content(
            model=ARITHMETIC_MODEL,
            contents=question,
            config=types.GenerateContentConfig(
                tools=[add, subtract, multiply, divide, modulus]
            )
        )

        return response.text
    except Exception as e:
        raise RuntimeError(f"Processing failed: {str(e)}")

@tool("Code Execution Tool")
def code_execution_tool(question: str, file_path: str) -> str:
    """Given a question and Python file, execute the file to answer the question.

       Args:
           question (str): Question to answer
           file_path (str): The Python file path
            
       Returns:
           str: Answer to the question
            
       Raises:
           RuntimeError: If processing fails"""
    try:
        client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])

        file = client.files.upload(file=file_path)

        response = client.models.generate_content(
            model=CODE_EXECUTION_MODEL,
            contents=[file, question],
            config=types.GenerateContentConfig(
                tools=[types.Tool(code_execution=types.ToolCodeExecution)]
            ),
        )
        
        for part in response.candidates[0].content.parts:
            if part.code_execution_result is not None:
                return part.code_execution_result.output
    except Exception as e:
        raise RuntimeError(f"Processing failed: {str(e)}")

@tool("Code Generation Tool")
def code_generation_tool(question: str, json_data: str) -> str:
    """Given a question and JSON data, generate and execute code to answer the question.
       Args:
           question (str): Question to answer
            file_path (str): The JSON data
       Returns:
           str: Answer to the question
           
       Raises:
           RuntimeError: If processing fails"""
    try:
        client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
                
        response = client.models.generate_content(
            model=CODE_GENERATION_MODEL,
            contents=[f"{question}\n{json_data}"],
            config=types.GenerateContentConfig(
                tools=[types.Tool(code_execution=types.ToolCodeExecution)]
            ),
        )
        
        for part in response.candidates[0].content.parts:
            if part.code_execution_result is not None:
                return part.code_execution_result.output
    except Exception as e:
        raise RuntimeError(f"Processing failed: {str(e)}")

# Arithmetic functions

@tool("Addition Tool")
def add(a: float, b: float) -> float:
    """Add two numbers.
    
    Args:
        a: First number
        b: Second number

    Returns:
        Result number
    """
    return a + b

@tool("Subtraction Tool")
def subtract(a: float, b: float) -> float:
    """Subtract two numbers.
    
    Args:
        a: First number
        b: Second number

    Returns:
        Result number
    """
    return a - b

@tool("Multiplication Tool")
def multiply(a: float, b: float) -> float:
    """Multiply two numbers.
    Args:
        a: First number
        b: Second number

    Returns:
        Result number
    """
    return a * b

@tool("Division Tool")
def divide(a: float, b: float) -> float:
    """Divide two numbers.
    
    Args:
        a: First number
        b: Second number

    Returns:
        Result number
    """
    if b == 0:
        raise ValueError("Cannot divide by zero.")
    return a / b

@tool("Modulus Tool")
def modulus(a: float, b: float) -> float:
    """Get the modulus of two numbers.
    
    Args:
        a: First number
        b: Second number

    Returns:
        Result number
    """
    return a % b

# TODO: Other functions