from smolagents.tools import Tool


# Creating Tool class
# Here we are creating a market data tool using `yfinance`
class ExtractMarketDataTool(Tool):
    name = "extract_market_data"
    description = ("Extracts market data (via `yfinance` or Yahoo! Finance) and "
                   "returns the information about a stock such as the stock's info, "
                   "address, city-state-zip-country, phone, website, industry, sector, "
                   "calendar data such as dividend rate, highs, lows, average, revenue, etc. "
                   "Use this when someone asks information about a market stock. "
                   "Ensure that a digestible or summarized version is the output (let the LLM summarize).")
    inputs = {'stock': {'type': 'string', 'description': 'The market/stock code that the User wants to inspect.'}}
    output_type = "string"

    def forward(self, stock: str) -> str:
        try:
            import yfinance as yf
        except ImportError as e:
            raise ImportError(
                "You must install packages `yfinance` to run this tool: for instance run `pip install yfinance`"
            ) from e

        try:
            data = yf.Ticker(f"{stock}")
            info_data = data.info
            # INFO
            address = info_data['address1']
            city_state_zip_country = f"{info_data['city']}, {info_data['state']}, {info_data['zip']}, {info_data['country']}"
            phone = f"{info_data['phone']}"
            website = f"{info_data['website']}"
            industry = f"{info_data['industry']}"
            sector = f"{info_data['sector']}"

            calendar_data = data.calendar
            # CALENDAR
            dividend_date = calendar_data['Dividend Date']
            ex_dividend_date = calendar_data['Ex-Dividend Date']
            earnings_date = calendar_data['Earnings Date']
            earnings_high = calendar_data['Earnings High']
            earnings_low = calendar_data['Earnings Low']
            earnings_avg = calendar_data['Earnings Average']
            rev_high = calendar_data['Revenue High']
            rev_low = calendar_data['Revenue Low']
            rev_avg = calendar_data['Revenue Average']
            

            output = (f"Here's information about {stock}:\n"
                      f"address: {address}\n"
                      f"city, state, zip, country: {city_state_zip_country}\n"
                      f"phone: {phone}\n"
                      f"website: {website}\n"
                      f"industry: {industry}\n"
                      f"sector: {sector}\n"
                      f"dividend date: {dividend_date}\n"
                      f"ex-dividend date: {ex_dividend_date}\n"
                      f"earnings date: {earnings_date}\n"
                      f"earnings high: {earnings_high}\n"
                      f"earnings low: {earnings_low}\n"
                      f"earnings average: {earnings_avg}\n"
                      f"revenue high: {rev_high}\n"
                      f"revenue low: {rev_low}\n"
                      f"revenue average: {rev_avg}\n")
            return output
        except Exception as e:
            return f"An unexpected error occurred: {str(e)}"

    def __init__(self, *args, **kwargs) -> None:
        self.is_initialized = False