from typing import Dict
from smolagents.tools import Tool
import yfinance as yf
import matplotlib.pyplot as plt
from io import BytesIO
import base64
import matplotlib

class StockChartTool(Tool):
    name = "stock_chart"
    description = "Generates a stock price chart for the given ticker symbol.  Returns the chart as a base64 encoded PNG image."
    inputs = {'ticker': {'type': 'string', 'description': 'The stock ticker symbol (e.g., AAPL, MSFT).'}}
    output_type = "string"  # Base64 encoded image string

    def forward(self, ticker: str) -> str:
        try:
            data = yf.download(ticker, period="30d")  # Download 30 days of data
            if data.empty:
                return "Error: No data found for that ticker."

            plt.figure(figsize=(10, 6))
            plt.plot(data['Close'])
            plt.title(f"{ticker} Stock Price (Last 30 Days)")
            plt.xlabel("Date")
            plt.ylabel("Closing Price")
            plt.grid(True)

            img_buf = BytesIO()
            plt.savefig(img_buf, format='png')
            img_buf.seek(0)
            img_base64 = base64.b64encode(img_buf.getvalue()).decode('utf-8')
            plt.close()  # Important to close the plot to release memory

            return f"<img src='data:image/png;base64,{img_base64}' alt='Stock Price Plot'>"

        except Exception as e:
            return f"Error generating chart: {e}"

    def __init__(self, *args, **kwargs):
        self.is_initialized = False # Add this line to avoid the warning