from smolagents import CodeAgent, LiteLLMModel, GoogleSearchTool, WikipediaSearchTool, Tool, VisitWebpageTool, HfApiModel
import os
import json
import requests
import pandas as pd
from huggingface_hub import InferenceClient

class GetFileTool(Tool):
    name = "get_file_tool"
    description = "This tool allows to download the file attached to the question"
    output_type = "string"
    inputs = {
        "task_id": {
            "type": "string",
            "description": "The task id of the question",
        },
        "file_name": {
            "type": "string",
            "description": "The name of the file to download"
        }
    }

    def forward(self, task_id: str, file_name: str) -> str:
        # Download the file from the task id
        file = f"https://agents-course-unit4-scoring.hf.space/files/{task_id}"
        # Save the file with the file name
        with open(file_name, "wb") as f:
            f.write(requests.get(file).content)

        # Return the file name  
        return os.path.abspath(file_name)
    
class LoadXlsxFileTool(Tool):
    name = "load_xlsx_file_tool"
    description = """This tool loads xlsx file into pandas and returns it"""
    inputs = {
        "file_path": {"type": "string", "description": "File path"}
    }
    output_type = "object"

    def forward(self, file_path: str) -> object:
        return pd.read_excel(file_path)

class LoadTextFileTool(Tool):
    name = "load_text_file_tool"
    description = """This tool loads any text file"""
    inputs = {
        "file_path": {"type": "string", "description": "File path"}
    }
    output_type = "string"

    def forward(self, file_path: str) -> object:
        with open(file_path, 'r', encoding='utf-8') as file:
            return file.read() 

class AudioToTextTool(Tool):
    name = "audio_to_text_tool"
    description = """This tool transcribes audio files into text"""
    inputs = {
        "file_path": {"type": "string", "description": "File path"}
    }
    output_type = "string"

    def forward(self, file_path: str) -> str:
        try:
            # Check if file exists
            if not os.path.exists(file_path):
                return f"Error: File {file_path} does not exist"

            # Read the audio file as raw bytes
            with open(file_path, "rb") as f:
                audio_data = f.read()

            # Set up the API URL and headers
            api_url = "https://router.huggingface.co/hf-inference/models/openai/whisper-large-v3-turbo"
            headers = {
                "Authorization": f"Bearer {os.getenv('HF_TOKEN')}",
                "Content-Type": "audio/mpeg"  # Assuming MP3 format
            }

            # Make the API request
            response = requests.post(api_url, headers=headers, data=audio_data)
            response.raise_for_status()  # Raise an exception for bad status codes
            
            # Parse and return the response
            output = response.json()
            return output["text"]

        except Exception as e:
            return f"Error transcribing audio: {str(e)}"
    
class ImageAnalysisTool(Tool):
    name = "image_analysis_tool"
    description = """This tool analyzes images and returns the text and the information in the image"""
    inputs = {
        "task_id": {"type": "string", "description": "The task id of the question"},
    }
    output_type = "string"

    def forward(self, task_id: str) -> str:
        client = InferenceClient(
            provider="nebius",
            api_key=os.getenv("HF_TOKEN"),
        )
            
        completion = client.chat.completions.create(
            model="google/gemma-3-27b-it",
            messages=[
                {
                    "role": "user",
                    "content": [
                    {
                        "type": "text",
                        "text": "Describe this image in markdown format"
                    },
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"https://agents-course-unit4-scoring.hf.space/files/{task_id}"
                        }
                    }
                ]
            }
        ],
        )

        return completion.choices[0].message.content


def final_answer_formatting(answer, question):
    model = LiteLLMModel(
        model_id="gemini/gemini-2.0-flash",
        api_key=os.getenv("GOOGLE_API_KEY"),
    )

    prompt = f"""
    You are an AI assistant specialized in the GAIA benchmark. For the question provided, generate the answer in the exact format requested by the question. Do not include any other text or creative additions.
    Question: {question}
    Answer: {answer}
    """

    messages = [
    {"role": "user", "content": [{"type": "text", "text": prompt}]}
    ]

    output = model(messages).content
    return output


web_agent = CodeAgent(
    model=LiteLLMModel(
        model_id="gemini/gemini-2.0-flash",
        api_key=os.getenv("GOOGLE_API_KEY"),
    ),
    tools=[
        WikipediaSearchTool(),
        GoogleSearchTool(provider="serper"), 
        VisitWebpageTool()
    ],
    add_base_tools=False,
    additional_authorized_imports=[
        "os", "requests", "inspect", "pandas", 
        "datetime", "re", "bs4", "markdownify"
    ],
    max_steps=10,
    name="web_agent",
    description="This agent is used to search the web for information"
)

audio_agent = CodeAgent(
    model=LiteLLMModel(
        model_id="gemini/gemini-2.0-flash",
        api_key=os.getenv("GOOGLE_API_KEY"),
    ),
    tools=[AudioToTextTool()],
    add_base_tools=False,
    max_steps=10,
    name="audio_agent",
    description="This agent is used to analyze and transcribe audio files"
)


manager_agent = CodeAgent(
    name="manager_agent",
    model=LiteLLMModel(
        model_id="gemini/gemini-2.5-flash-preview-04-17",
        api_key=os.getenv("GOOGLE_API_KEY"),
    ),
    tools=[GetFileTool(), LoadXlsxFileTool(), LoadTextFileTool(), ImageAnalysisTool()],
    managed_agents=[web_agent, audio_agent],
    additional_authorized_imports=[
        "pandas"
    ],
    planning_interval=5,
    verbosity_level=1,
    max_steps=10,
)