import streamlit as st
import pandas as pd
import yfinance as yf
import plotly.express as px
import plotly.graph_objects as go
from datetime import datetime, timedelta
import requests
from autogen import AssistantAgent, UserProxyAgent, GroupChat, GroupChatManager
import json

# Set Streamlit page configuration
st.set_page_config(
    page_title="AI-Powered Financial Advisor",
    page_icon="💰",
    layout="wide",
)

# Custom CSS styling
st.markdown("""
    <style>
    .stTextInput > label {
        font-weight: 500;
    }
    .stSelectbox > label {
        font-weight: 500;
    }
    .stNumberInput > label {
        font-weight: 500;
    }
    .stButton > button {
        background-color: #4CAF50;
        color: white;
        font-weight: bold;
    }
    .result-box {
        background-color: #f5f5f5;
        padding: 20px;
        border-radius: 10px;
        margin: 20px 0;
        border: 1px solid #ddd;
    }
    .metric-card {
        background-color: white;
        padding: 15px;
        border-radius: 8px;
        box-shadow: 0 2px 4px rgba(0,0,0,0.1);
        text-align: center;
        margin-bottom: 15px;
    }
    .metric-value {
        font-size: 24px;
        font-weight: bold;
        color: #1E88E5;
    }
    .metric-label {
        font-size: 14px;
        color: #757575;
    }
    </style>
""", unsafe_allow_html=True)

# Currency and interest rate data by country
country_data = {
    "India": {
        "currency": "₹",
        "currency_code": "INR",
        "base_interest_rate": 6.5,  # Reserve Bank of India repo rate
        "tax_brackets": [
            {"limit": 250000, "rate": 0},
            {"limit": 500000, "rate": 5},
            {"limit": 1000000, "rate": 20},
            {"limit": float('inf'), "rate": 30}
        ],
        "major_indices": ["^NSEI", "^BSESN"],  # Nifty 50, Sensex
        "popular_funds": ["ICICRED.NS", "HDFCAMC.NS", "KOTAKBANK.NS"],
        "safe_instruments": {"Fixed Deposit": 5.5, "PPF": 7.1, "Government Bonds": 7.0}
    },
    "USA": {
        "currency": "$",
        "currency_code": "USD",
        "base_interest_rate": 5.5,  # Federal Reserve rate
        "tax_brackets": [
            {"limit": 11000, "rate": 10},
            {"limit": 44725, "rate": 12},
            {"limit": 95375, "rate": 22},
            {"limit": 182100, "rate": 24},
            {"limit": 231250, "rate": 32},
            {"limit": 578125, "rate": 35},
            {"limit": float('inf'), "rate": 37}
        ],
        "major_indices": ["^GSPC", "^DJI", "^IXIC"],  # S&P 500, Dow Jones, Nasdaq
        "popular_funds": ["SPY", "VOO", "QQQ"],
        "safe_instruments": {"Treasury Bonds": 4.2, "CD": 4.0, "High-Yield Savings": 3.8}
    },
    "UK": {
        "currency": "£",
        "currency_code": "GBP",
        "base_interest_rate": 5.25,  # Bank of England rate
        "tax_brackets": [
            {"limit": 12570, "rate": 0},  # Personal Allowance
            {"limit": 50270, "rate": 20},  # Basic rate
            {"limit": 125140, "rate": 40},  # Higher rate
            {"limit": float('inf'), "rate": 45}  # Additional rate
        ],
        "major_indices": ["^FTSE"],  # FTSE 100
        "popular_funds": ["CUKX.L", "MIDD.L", "ISF.L"],
        "safe_instruments": {"Premium Bonds": 4.65, "Fixed Rate Bonds": 4.8, "Cash ISA": 4.5}
    }
}

# Function to fetch real-time currency exchange rates
@st.cache_data(ttl=3600)  # Cache for 1 hour
def get_exchange_rates(base_currency):
    try:
        url = f"https://open.er-api.com/v6/latest/{base_currency}"
        response = requests.get(url)
        data = response.json()
        if data["result"] == "success":
            return data["rates"]
        else:
            return {"USD": 1.0, "INR": 82.5, "GBP": 0.79}
    except:
        # Default fallback rates
        return {"USD": 1.0, "INR": 82.5, "GBP": 0.79}

# Function to fetch market index data
@st.cache_data(ttl=3600)  # Cache for 1 hour
def get_market_indices(ticker_symbols):
    end_date = datetime.now()
    start_date = end_date - timedelta(days=365)
    
    data = {}
    for ticker in ticker_symbols:
        try:
            ticker_data = yf.download(ticker, start=start_date, end=end_date)
            if not ticker_data.empty:
                data[ticker] = ticker_data
        except:
            pass
    
    return data

# Function to get real-time inflation data
@st.cache_data(ttl=86400)  # Cache for 1 day
def get_inflation_rates():
    # This would ideally be from an API, but using static recent data for demo
    return {
        "India": 5.1,
        "USA": 3.3,
        "UK": 3.2
    }

# Function to convert currency
def convert_currency(amount, from_currency, to_currency):
    if from_currency == to_currency:
        return amount
    
    rates = get_exchange_rates(from_currency)
    if to_currency in rates:
        return amount * rates[to_currency]
    return amount  # Fallback to original amount if conversion fails

# Streamlit UI components
st.title("AI-Powered Financial Advisor")

# Sidebar for real-time market information
with st.sidebar:
    st.header("Market Overview")
    
    # Get inflation rates
    inflation_rates = get_inflation_rates()
    
    # Display inflation rates
    st.subheader("Current Inflation Rates")
    for country, rate in inflation_rates.items():
        st.metric(country, f"{rate}%")
    
    st.markdown("---")
    
    # Display current date and time
    st.write(f"Last updated: {datetime.now().strftime('%Y-%m-%d %H:%M')}")

# Main input form
col1, col2 = st.columns(2)

with col1:
    name = st.text_input("Full Name")
    location = st.selectbox("Country", options=["India", "USA", "UK"])
    age = st.number_input("Age", min_value=18, max_value=100)
    marital_status = st.selectbox("Marital Status", ["Single", "Married"])

    # Get currency symbol based on selected country
    currency_symbol = country_data[location]["currency"]
    currency_code = country_data[location]["currency_code"]

with col2:
    assets = st.multiselect("Assets", 
        ["Car", "House", "Bank Balance", "Stocks", "Mutual Funds", "Real Estate", "Gold", "Other"])
    asset_values = {asset: st.number_input(f"{asset} Value ({currency_symbol})", min_value=0) for asset in assets}

    debts = st.multiselect("Debts", 
        ["Education Loan", "Home Loan", "Personal Loan", "Credit Card", "Gold Loan", "Other"])
    debt_values = {debt: st.number_input(f"{debt} Amount ({currency_symbol})", min_value=0) for debt in debts}

monthly_savings = st.number_input(f"Monthly Savings ({currency_symbol})", min_value=0)
target_amount = st.number_input(f"Target Amount ({currency_symbol})", min_value=0)
target_years = st.number_input("Target Time (Years)", min_value=1, max_value=50)

# Market data fetching based on selected country
market_data_loaded = False
if location:
    try:
        with st.expander("View Current Market Data"):
            st.subheader(f"Market Indices - {location}")
            indices_data = get_market_indices(country_data[location]["major_indices"])
            
            if indices_data:
                for ticker, data in indices_data.items():
                    # Calculate percentage change
                    if not data.empty:
                        current = data['Close'].iloc[-1]
                        previous = data['Close'].iloc[-2]
                        change_pct = (current - previous) / previous * 100
                        
                        # Display the index name and its current value
                        index_name = {
                            "^NSEI": "Nifty 50", "^BSESN": "Sensex",
                            "^GSPC": "S&P 500", "^DJI": "Dow Jones", "^IXIC": "Nasdaq",
                            "^FTSE": "FTSE 100"
                        }.get(ticker, ticker)
                        
                        st.metric(
                            index_name, 
                            f"{current:.2f}", 
                            f"{change_pct:.2f}%",
                            delta_color="normal"
                        )
                        
                        # Plot the index trend
                        fig = px.line(data, y='Close', title=f"{index_name} - Past Year")
                        st.plotly_chart(fig, use_container_width=True)
                market_data_loaded = True
            else:
                market_data_loaded = False
    except:
        # Silently handle the exception without showing error to user
        market_data_loaded = False

if st.button("Calculate"):
    # Display a loading spinner
    with st.spinner("Analyzing financial data and generating recommendations..."):
        # Calculate total assets and debts
        total_assets = sum(asset_values.values())
        total_debts = sum(debt_values.values())
        net_worth = total_assets - total_debts
        
        # Dashboard Metrics
        st.markdown('<div class="result-box">', unsafe_allow_html=True)
        
        # Financial overview section
        st.header("Financial Overview")
        
        # Display key metrics
        metric_cols = st.columns(4)
        with metric_cols[0]:
            st.markdown(f"""
            <div class="metric-card">
                <div class="metric-value">{currency_symbol}{net_worth:,.2f}</div>
                <div class="metric-label">Net Worth</div>
            </div>
            """, unsafe_allow_html=True)
            
        with metric_cols[1]:
            debt_to_asset = 0 if total_assets == 0 else (total_debts / total_assets) * 100
            st.markdown(f"""
            <div class="metric-card">
                <div class="metric-value">{debt_to_asset:.1f}%</div>
                <div class="metric-label">Debt-to-Asset Ratio</div>
            </div>
            """, unsafe_allow_html=True)
            
        with metric_cols[2]:
            st.markdown(f"""
            <div class="metric-card">
                <div class="metric-value">{currency_symbol}{monthly_savings:,.2f}</div>
                <div class="metric-label">Monthly Savings</div>
            </div>
            """, unsafe_allow_html=True)
            
        with metric_cols[3]:
            # FIX 2: Use the user-entered target_years value directly instead of calculating
            st.markdown(f"""
            <div class="metric-card">
                <div class="metric-value">{target_years}</div>
                <div class="metric-label">Years to Goal</div>
            </div>
            """, unsafe_allow_html=True)
        
        # Net Worth Breakdown Chart
        st.subheader("Net Worth Breakdown")
        
        # FIX 1: Improve pie chart data preparation to properly show both assets and debts
        if total_assets > 0 or total_debts > 0:
            # Create two separate traces for a better visualization
            fig = go.Figure()
            
            # Group assets together (positive values)
            asset_labels = []
            asset_values_list = []
            for asset, value in asset_values.items():
                if value > 0:
                    asset_labels.append(asset)
                    asset_values_list.append(value)
            
            # Group debts together (use absolute values for display)
            debt_labels = []
            debt_values_list = []
            for debt, value in debt_values.items():
                if value > 0:
                    debt_labels.append(debt)
                    debt_values_list.append(value)  # Using positive values for better visualization
            
            # Create combined labels and values for the pie chart
            combined_labels = asset_labels + debt_labels
            combined_values = asset_values_list + [-v for v in debt_values_list]  # Make debt values negative
            
            if combined_labels and combined_values:
                # Use abs(val) for sizing the pie segments but keep colors based on sign
                fig = go.Figure(data=[go.Pie(
                    labels=combined_labels,
                    values=[abs(val) for val in combined_values],  # Use absolute values for segment size
                    hole=.4,
                    textinfo='label+percent',
                    marker=dict(colors=[
                        '#4CAF50' if val > 0 else '#F44336' for val in combined_values
                    ])
                )])
                
                # Add a color legend
                fig.update_layout(
                    title_text="Assets and Debts",
                    legend_title="Items",
                    annotations=[
                        dict(text="Assets", x=0.85, y=1.1, showarrow=False, font=dict(color='#4CAF50', size=12)),
                        dict(text="Debts", x=0.95, y=1.1, showarrow=False, font=dict(color='#F44336', size=12))
                    ]
                )
                st.plotly_chart(fig, use_container_width=True)
        else:
            st.info("Please enter asset and debt values to see breakdown")
        
        # Create AutoGen agents
        financial_planner = AssistantAgent(
            name="Financial_Planner",
            llm_config={"model": "gpt-4o"},
            system_message=f"""
            You are a certified financial planner specializing in {location}-based financial planning.
            Use the following real-time market data for your analysis:
            - Current inflation rate in {location}: {inflation_rates.get(location, 5.0)}%
            - Base interest rate: {country_data[location]['base_interest_rate']}%
            - Safe investment returns: {json.dumps(country_data[location]['safe_instruments'])}
            - Tax brackets: {json.dumps(country_data[location]['tax_brackets'])}
            
            Your task is to:
            1. Calculate user's net worth and analyze financial health
            2. Assess feasibility of financial goals
            3. Provide detailed investment recommendations specific to {location}
            """
        )

        market_analyst = AssistantAgent(
            name="Market_Analyst",
            llm_config={"model": "gpt-4o"},
            system_message=f"""
            You are a market analyst specializing in {location} financial markets.
            Use the following real-time market data:
            - Current inflation rate in {location}: {inflation_rates.get(location, 5.0)}%
            - Popular market indices in {location}: {country_data[location]['major_indices']}
            - Popular funds in {location}: {country_data[location]['popular_funds']}
            
            Your task is to:
            1. Analyze current market conditions in {location}
            2. Recommend specific investment vehicles appropriate for the user's situation
            3. Provide a realistic forecast of expected returns in {location}'s market
            """
        )

        tax_advisor = AssistantAgent(
            name="Tax_Advisor",
            llm_config={"model": "gpt-4o"},
            system_message=f"""
            You are a tax advisor specializing in {location} tax law.
            Use the following real-time data:
            - Tax brackets in {location}: {json.dumps(country_data[location]['tax_brackets'])}
            - Available tax-saving instruments in {location}: {json.dumps(country_data[location]['safe_instruments'])}
            
            Your task is to:
            1. Calculate potential tax liability based on income and assets
            2. Suggest specific tax-saving strategies available in {location}
            3. Recommend tax-efficient investment vehicles for the user's goals
            """
        )

        user_proxy = UserProxyAgent(
            name="User",
            human_input_mode="NEVER",
            system_message="You represent the user and relay their financial goals.",
            code_execution_config={"use_docker": False}
        )

        # Group chat setup
        group_chat = GroupChat(
            agents=[user_proxy, financial_planner, market_analyst, tax_advisor],
            messages=[],
            max_round=10
        )

        manager = GroupChatManager(groupchat=group_chat, llm_config={"model": "gpt-4o"})

        # Start conversation
        user_proxy.initiate_chat(
            manager,
            message=f"""
            User profile:
            - Name: {name}
            - Location: {location}
            - Age: {age}
            - Marital Status: {marital_status}
            - Assets: {asset_values}
            - Debts: {debt_values}
            - Monthly Savings: {currency_symbol}{monthly_savings}
            - Target Amount: {currency_symbol}{target_amount}
            - Target Time: {target_years} years
            
            Task:
            1. Analyze feasibility of achieving the target amount of {currency_symbol}{target_amount} in {target_years} years.
            2. Provide investment recommendations specific to {location} market.
            3. Suggest tax-saving strategies available in {location}.
            """
        )

        # Modified message filtering logic
        if len(group_chat.messages) > 0:
            # Create a placeholder for each agent
            output = {}
            for agent in ["Financial_Planner", "Market_Analyst", "Tax_Advisor"]:
                output[agent] = []
            
            # Collect messages by agent
            for msg in group_chat.messages:
                if 'name' in msg and msg['name'] in output:
                    content = msg['content'].strip()
                    if content and not content.startswith("Next speaker:"):
                        output[msg['name']].append(content)
            
            # Display messages from each agent
            for agent, messages in output.items():
                if messages:
                    st.subheader(f"{agent.replace('_', ' ')} Analysis")
                    for msg in messages:
                        st.markdown(msg)
                    st.markdown("---")
            
            # Investment Growth Projection Chart
            st.subheader("Investment Growth Projection")
            
            # Simplified projection calculation
            years = list(range(1, target_years + 1))
            
            # Conservative scenario (lower return rate)
            conservative_rate = country_data[location]['base_interest_rate'] - 1.0
            conservative_values = [
                monthly_savings * 12 * (((1 + conservative_rate/100) ** y) - 1) / (conservative_rate/100)
                for y in years
            ]
            
            # Moderate scenario
            moderate_rate = country_data[location]['base_interest_rate'] + 1.0
            moderate_values = [
                monthly_savings * 12 * (((1 + moderate_rate/100) ** y) - 1) / (moderate_rate/100)
                for y in years
            ]
            
            # Aggressive scenario
            aggressive_rate = country_data[location]['base_interest_rate'] + 3.0
            aggressive_values = [
                monthly_savings * 12 * (((1 + aggressive_rate/100) ** y) - 1) / (aggressive_rate/100)
                for y in years
            ]
            
            # Target line
            target_line = [target_amount] * len(years)
            
            # Create figure
            fig = go.Figure()
            
            # Add traces
            fig.add_trace(go.Scatter(
                x=years, y=conservative_values,
                mode='lines',
                name=f'Conservative ({conservative_rate:.1f}%)',
                line=dict(color='blue', dash='dash')
            ))
            
            fig.add_trace(go.Scatter(
                x=years, y=moderate_values,
                mode='lines',
                name=f'Moderate ({moderate_rate:.1f}%)',
                line=dict(color='green')
            ))
            
            fig.add_trace(go.Scatter(
                x=years, y=aggressive_values,
                mode='lines',
                name=f'Aggressive ({aggressive_rate:.1f}%)',
                line=dict(color='red', dash='dot')
            ))
            
            fig.add_trace(go.Scatter(
                x=years, y=target_line,
                mode='lines',
                name='Target Amount',
                line=dict(color='black', dash='dash')
            ))
            
            # Update layout
            fig.update_layout(
                title=f'Projected Growth of Monthly Investment ({currency_symbol}{monthly_savings}/month)',
                xaxis_title='Years',
                yaxis_title=f'Value ({currency_symbol})',
                legend=dict(y=0.5, traceorder='reversed'),
                hovermode='x unified'
            )
            
            # Format y-axis with appropriate currency
            fig.update_layout(yaxis=dict(
                tickprefix=currency_symbol,
                tickformat=",."
            ))
            
            st.plotly_chart(fig, use_container_width=True)
            
            # If no valid messages were found, show a more user-friendly message
            if all(len(msgs) == 0 for msgs in output.values()):
                st.info("""
                Our advisors are still analyzing your financial situation. 
                Please ensure you've entered all required information and try again.
                """)
                
        else:
            st.info("Our advisors are preparing your personalized financial analysis. Please try again in a moment.")

        st.markdown('</div>', unsafe_allow_html=True)