import sys
import subprocess
subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'websockets>=13.0'])

import gradio as gr
import yfinance as yf
import pandas as pd
import matplotlib.pyplot as plt
from io import BytesIO
import base64

# Define the portfolio data
portfolio = {
    'Ticker': ['TSLA', 'PLTR', 'SOUN'],
    'Shares': [200, 490, 652],
    'Avg Cost': [245.0, 52.0, 7.5]
}

def get_current_prices(tickers):
    """
    Fetch the latest closing prices for the given tickers.
    """
    data = yf.download(tickers, period='1d')['Close']
    return data.iloc[-1]

def generate_pie_chart(values, labels):
    """
    Generate a pie chart as a base64-encoded image.
    """
    fig, ax = plt.subplots()
    ax.pie(values, labels=labels, autopct='%1.1f%%', startangle=90)
    ax.set_title('Portfolio Allocation by Current Value')
    ax.axis('equal')  # Equal aspect ratio ensures the pie is circular.
    buf = BytesIO()
    plt.savefig(buf, format='png')
    buf.seek(0)
    img_base64 = base64.b64encode(buf.read()).decode('utf-8')
    plt.close(fig)
    return f"data:image/png;base64, {img_base64}"

def display_portfolio():
    """
    Compute portfolio metrics and generate outputs.
    """
    df = pd.DataFrame(portfolio)
    try:
        prices = get_current_prices(df['Ticker'].tolist())
        df['Current Price'] = prices.values
    except Exception as e:
        return f"Error fetching prices: {str(e)}", "", "N/A", "N/A"
    df['Cost Basis'] = df['Shares'] * df['Avg Cost']
    df['Current Value'] = df['Shares'] * df['Current Price']
    df['Gain/Loss'] = df['Current Value'] - df['Cost Basis']
    df['Gain/Loss %'] = (df['Gain/Loss'] / df['Cost Basis']) * 100
    total_cost = df['Cost Basis'].sum()
    total_value = df['Current Value'].sum()
    total_gain = total_value - total_cost
    # Format the table as HTML
    table_html = df.to_html(index=False, float_format=lambda x: f'{x:.2f}')
    # Generate pie chart
    pie_img = generate_pie_chart(df['Current Value'], df['Ticker'])
    total_value_str = f"${total_value:.2f}"
    total_gain_str = f"${total_gain:.2f} ({(total_gain / total_cost * 100):.2f}%)"
    return table_html, pie_img, total_value_str, total_gain_str

# Create the Gradio interface
with gr.Blocks(title="Investment Portfolio Dashboard") as demo:
    gr.Markdown("# Investment Portfolio Dashboard")
    gr.Markdown("View your portfolio metrics with real-time stock prices.")
    refresh_btn = gr.Button("Refresh Data")
    table_output = gr.HTML(label="Portfolio Table")
    pie_output = gr.Image(label="Portfolio Allocation", type="pil")  # Will display base64 image
    total_value_output = gr.Textbox(label="Total Portfolio Value")
    total_gain_output = gr.Textbox(label="Total Gain/Loss")
    refresh_btn.click(
        fn=display_portfolio,
        inputs=[],
        outputs=[table_output, pie_output, total_value_output, total_gain_output]
    )

# Launch the demo (automatically handled in Hugging Face Spaces)
if __name__ == "__main__":
    demo.launch()