from smolagents.tools import Tool
import pandas as pd
from huggingface_hub import hf_hub_download

class ExcelAnalysisTool(Tool):
    def __init__(self):
        self.name = "ExcelAnalysisTool"
        self.description = (
            "Loads an Excel file from the GAIA dataset on Hugging Face and calculates "
            "the total sales for items labeled as 'food', excluding drinks. "
            "Provide input as a string with the filename, e.g., 'sales_data.xlsx'."
        )
        self.inputs = {
            "filename": {
                "type": "string",
                "description": "The name of the Excel file to process (e.g., 'sales_data.xlsx').",
                "nullable": True  # Сделать filename необязательным
            }
        }
        self.output_type = "string"
        self.repo_id = "gaia-benchmark/GAIA"
        super().__init__(
            name=self.name,
            description=self.description,
            inputs=self.inputs,
            output_type=self.output_type
        )

    def forward(self, filename: str = "") -> str:
        """
        Loads and processes the Excel file.

        Args:
            filename (str): The name of the Excel file (e.g., 'sales_data.xlsx').

        Returns:
            str: Total food sales in USD, or an error message.
        """
        if not filename:
            return "Error: 'filename' is required."

        try:
            # Download the file from Hugging Face Hub
            file_path = hf_hub_download(
                repo_id=self.repo_id,
                filename=filename,
                repo_type="dataset"
            )

            # Load the Excel file into a DataFrame
            df = pd.read_excel(file_path)

            # Filter rows: category == 'food' and item != 'drinks'
            food_sales = df[
                (df['category'].str.lower() == 'food') &
                (df['item'].str.lower() != 'drinks')
            ]

            total_sales = food_sales['sales'].sum()
            return f"Total sales for food items: ${total_sales:.2f}"

        except FileNotFoundError:
            return "Error: The specified file was not found."
        except KeyError as e:
            return f"Error: Missing expected column in the Excel file: {str(e)}"
        except Exception as e:
            return f"An unexpected error occurred: {str(e)}"