import streamlit as st
import pandas as pd
from utils import *

st.title("Stock Portfolio and Tax Analysis")

# Initialize empty dataframes for stock sales and portfolio stocks
if 'stock_sales' not in st.session_state:
    st.session_state.stock_sales = pd.DataFrame(columns=['Name', 'Sale Price', 'Purchase Price', 'Quantity', 'Holding Period (Months)'])
if 'portfolio_stocks' not in st.session_state:
    st.session_state.portfolio_stocks = pd.DataFrame(columns=['Name', 'Purchase Price', 'Current Price', 'Quantity'])

# Define function to add stock sale
def add_stock_sale():
    with st.form(key="stock_sale_form"):
        col1, col2, col3, col4, col5 = st.columns(5)
        with col1:
            name = st.text_input("Stock Name", key="stock_sale_name")
        with col2:
            sale_price = st.number_input("Sale Price", min_value=1.0, key="stock_sale_price")
        with col3:
            purchase_price = st.number_input("Purchase Price", min_value=1.0, key="stock_sale_purchase_price")
        with col4:
            sell_quantity = st.number_input("Quantity", min_value=10, key="stock_sale_qty")
        with col5:
            holding_period_months = st.number_input("Holding Period (Months)", min_value=1, key="stock_sale_holding_period")
        
        submit_button = st.form_submit_button("Add Sold Stock")
        if submit_button:
            new_stock = pd.DataFrame([{'Name': name, 'Sale Price': sale_price, 'Purchase Price': purchase_price, 'Quantity': sell_quantity, 'Holding Period (Months)': holding_period_months}])
            st.session_state.stock_sales = pd.concat([st.session_state.stock_sales, new_stock], ignore_index=True)
            st.success("Sold Stock Added!")

# Define function to add portfolio stock
def add_portfolio_stock():
    with st.form(key="portfolio_stock_form"):
        col1, col2, col3, col4 = st.columns(4)
        with col1:
            name = st.text_input("Stock Name", key="portfolio_stock_name")
        with col2:
            purchase_price = st.number_input("Purchase Price", min_value=0.0, key="portfolio_stock_purchase_price")
        with col3:
            current_price = st.number_input("Current Price", min_value=0.0, key="portfolio_stock_current_price")
        with col4:
            quantity = st.number_input("Quantity", min_value=0, key="portfolio_stock_quantity")

        submit_button = st.form_submit_button("Add Portfolio Stock")
        if submit_button:
            new_stock = pd.DataFrame([{'Name': name, 'Purchase Price': purchase_price, 'Current Price': current_price, 'Quantity': quantity}])
            st.session_state.portfolio_stocks = pd.concat([st.session_state.portfolio_stocks, new_stock], ignore_index=True)
            st.success("Stock Added to Portfolio!")

# Define functions to edit and delete stocks
def delete_stock_sale(index):
    st.session_state.stock_sales = st.session_state.stock_sales.drop(index).reset_index(drop=True)
    st.success("Sold Stock Deleted!")

def delete_portfolio_stock(index):
    st.session_state.portfolio_stocks = st.session_state.portfolio_stocks.drop(index).reset_index(drop=True)
    st.success("Stock Deleted from Portfolio!")

def edit_stock_sale(index, name, sale_price, purchase_price, holding_period_months):
    st.session_state.stock_sales.at[index, 'Name'] = name
    st.session_state.stock_sales.at[index, 'Sale Price'] = sale_price
    st.session_state.stock_sales.at[index, 'Purchase Price'] = purchase_price
    st.session_state.stock_sales.at[index, 'Quantity'] = sell_quantity
    st.session_state.stock_sales.at[index, 'Holding Period (Months)'] = holding_period_months
    st.success("Sold Stock Updated!")

def edit_portfolio_stock(index, name, purchase_price, current_price, quantity):
    st.session_state.portfolio_stocks.at[index, 'Name'] = name
    st.session_state.portfolio_stocks.at[index, 'Purchase Price'] = purchase_price
    st.session_state.portfolio_stocks.at[index, 'Current Price'] = current_price
    st.session_state.portfolio_stocks.at[index, 'Quantity'] = quantity
    st.success("Portfolio Stock Updated!")

# UI Tabs
tabs = st.tabs(["Add Stocks", "Manage Stocks", "Tax Analysis"])

with tabs[0]:
    st.header("Add New Stock Sale and Portfolio Stock")
    st.subheader("Sold Stock")
    add_stock_sale()
    st.subheader("Portfolio Stock")
    add_portfolio_stock()

with tabs[1]:
    st.header("Manage Stocks")

    st.subheader("Sold Stock")
    st.dataframe(st.session_state.stock_sales)

    if st.session_state.stock_sales.shape[0] > 0:
        selected_sale_index = st.selectbox("Select Sale to Edit", range(st.session_state.stock_sales.shape[0]), key="select_sale_to_edit")
        selected_sale = st.session_state.stock_sales.iloc[selected_sale_index]
        with st.form(key="edit_stock_sale_form"):
            col1, col2, col3, col4, col5 = st.columns(5)
            with col1:
                name = st.text_input("Edit Stock Sale Name", selected_sale['Name'], key="edit_stock_sale_name")
            with col2:
                sale_price = st.number_input("Edit Sale Price", min_value=0.0, value=selected_sale['Sale Price'], key="edit_stock_sale_price")
            with col3:
                purchase_price = st.number_input("Edit Purchase Price", min_value=0.0, value=selected_sale['Purchase Price'], key="edit_stock_sale_purchase_price")
            with col4:
                sell_quantity = st.number_input("Quantity", min_value=0.0, value=selected_sale['Quantity'], key="edit_stock_sale_qty")
            with col5:
                holding_period_months = st.number_input("Edit Holding Period (Months)", min_value=0, value=selected_sale['Holding Period (Months)'], key="edit_stock_sale_holding_period")

            update_button = st.form_submit_button("Update Stock Sale")
            if update_button:
                edit_stock_sale(selected_sale_index, name, sale_price, purchase_price, holding_period_months)

        delete_button = st.button("Delete Stock Sale", key="delete_stock_sale_button")
        if delete_button:
            delete_stock_sale(selected_sale_index)

    st.subheader("Portfolio Stocks")
    st.dataframe(st.session_state.portfolio_stocks)

    if st.session_state.portfolio_stocks.shape[0] > 0:
        selected_stock_index = st.selectbox("Select Portfolio Stock to Edit", range(st.session_state.portfolio_stocks.shape[0]), key="select_stock_to_edit_portfolio")
        selected_stock = st.session_state.portfolio_stocks.iloc[selected_stock_index]
        with st.form(key="edit_portfolio_stock_form"):
            col1, col2, col3, col4 = st.columns(4)
            with col1:
                name = st.text_input("Edit Portfolio Stock Name", selected_stock['Name'], key="edit_portfolio_stock_name")
            with col2:
                purchase_price = st.number_input("Edit Purchase Price", min_value=0.0, value=selected_stock['Purchase Price'], key="edit_portfolio_stock_purchase_price")
            with col3:
                current_price = st.number_input("Edit Current Price", min_value=0.0, value=selected_stock['Current Price'], key="edit_portfolio_stock_current_price")
            with col4:
                quantity = st.number_input("Edit Quantity", min_value=0, value=selected_stock['Quantity'], key="edit_portfolio_stock_quantity")

            update_button = st.form_submit_button("Update Portfolio Stock")
            if update_button:
                edit_portfolio_stock(selected_stock_index, name, purchase_price, current_price, quantity)

        delete_button = st.button("Delete Portfolio Stock", key="delete_portfolio_stock_button")
        if delete_button:
            delete_portfolio_stock(selected_stock_index)

with tabs[2]:
    st.header("Tax Analysis")

    if st.session_state.stock_sales.shape[0] > 0:
        portfolio = Portfolio()
        for _, row in st.session_state.stock_sales.iterrows():
            portfolio.add_stock(StockSale(row['Name'], row['Sale Price'], row['Purchase Price'], row['Quantity'], row['Holding Period (Months)']))

        total_stcg, total_ltcg = portfolio.analyze_portfolio()
        stcg_tax, ltcg_tax, total_tax = portfolio.apply_tax(total_stcg, total_ltcg)

        stocks = []
        for _, row in st.session_state.portfolio_stocks.iterrows():
            stocks.append(Stock(row['Name'], row['Purchase Price'], row['Current Price'], row['Quantity']))

        suggested_stocks = optimize_taxes_selling_stocks(stocks)
        total_loss = calculate_total_loss(suggested_stocks)
        recalculated_tax = max(0, total_tax - total_loss)

        st.subheader("Analysis Results")
        if total_tax > total_loss:
            
            st.write(f"Total Short-Term Capital Gains: ₹{total_stcg:.2f}")
            st.write(f"Total Long-Term Capital Gains: ₹{total_ltcg:.2f}")
            st.write(f"Short-Term Capital Gains Tax: ₹{stcg_tax:.2f}")
            st.write(f"Long-Term Capital Gains Tax: ₹{ltcg_tax:.2f}")
            st.write(f"Total Tax Liability: ₹{total_tax:.2f}")
            st.write(f"Total Loss of loss making stockes: ₹{total_loss:.2f}")

           # Add a vertical spacing for better separation
            st.write("\n")
            st.subheader("Suggested Stocks to Sell to Offset Gains:")
            # Use Streamlit's dataframe for a tabular vies
            suggested_stocks_data = {
            'Name': [stock.name for stock in suggested_stocks],
            'Current Price': [stock.current_price for stock in suggested_stocks],
            'Purchase Price': [stock.purchase_price for stock in suggested_stocks]
            }
            
            st.dataframe(suggested_stocks_data)
            
            st.write(f"Total Tax Liability before selling loss making stocks: ₹{total_tax:.2f}")
            st.write(f"Adjusted Tax Liability after Selling Loss-Making Stocks and buy again: ₹{recalculated_tax:.2f}")
        
        elif total_loss >= 0:
            st.write(f"Total Short-Term Capital Gains: ₹{total_stcg:.2f}")
            st.write(f"Total Long-Term Capital Gains: ₹{total_ltcg:.2f}")
            st.write(f"Short-Term Capital Gains Tax: ₹{stcg_tax:.2f}")
            st.write(f"Long-Term Capital Gains Tax: ₹{ltcg_tax:.2f}")
            st.write(f"Total Tax Liability: ₹{total_tax:.2f}")
            
        else:
            st.write("No suggested stocks for selling to offset gains.")
            

# Ensure you have the 'calculate_total_loss' function implemented in utils.py
def calculate_total_loss(suggested_stocks):
    total_loss = 0.0
    for stock in suggested_stocks:
        total_loss += (stock.purchase_price - stock.current_price) * stock.quantity
    return total_loss