import importlib
import os

import requests
import yaml
import pandas as pd


DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
groq_api_key = os.getenv("GROQ_API_KEY")

from smolagents import CodeAgent, DuckDuckGoSearchTool, VisitWebpageTool, WikipediaSearchTool, Tool, SpeechToTextTool
from langchain_groq import ChatGroq


# === TOOLS ===
class GetTaskFileTool(Tool):
    name = "get_task_file_tool"
    description = """This tool downloads the file content associated with the given task_id if exists. Returns absolute file path"""
    inputs = {
        "task_id": {"type": "string", "description": "Task id"},
        "file_name": {"type": "string", "description": "File name"},
    }
    output_type = "string"

    def forward(self, task_id: str, file_name: str) -> str:
        response = requests.get(f"{DEFAULT_API_URL}/files/{task_id}", timeout=15)
        response.raise_for_status()
        with open(file_name, 'wb') as file:
            file.write(response.content)
        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()

# === PROMPTS ===
prompts = yaml.safe_load(
    importlib.resources.files("smolagents.prompts").joinpath("code_agent.yaml").read_text()
)
prompts["system_prompt"] = (
    "You are a general AI assistant. I will ask you a question. Report your thoughts, and finish your answer with the following template: FINAL ANSWER: [YOUR FINAL ANSWER]. YOUR FINAL ANSWER should be 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. "
    + prompts["system_prompt"]
)

# === AGENT FUNCTION ===
def agentRudra():
    model = ChatGroq(
        model_name="llama3-8b-8192",
        api_key=groq_api_key,
        temperature=0
    )
    
    agent = CodeAgent(
        tools=[
            DuckDuckGoSearchTool(),
            VisitWebpageTool(),
            WikipediaSearchTool(),
            GetTaskFileTool(),
            SpeechToTextTool(),
            LoadXlsxFileTool(),
            LoadTextFileTool()
        ],
        model=model,
        prompt_templates=prompts,
        max_steps=15,
        additional_authorized_imports=["pandas"]
    )
    return agent