import streamlit as st
import pandas as pd
import numpy as np
import os
from datetime import datetime, timedelta, date
import time
import io
import base64

# Import your existing pages (make sure they're in the same directory or in a 'pages' directory)
# If you have these in separate files, you'll need to adjust the imports
# For example:
# from pages.dashboard import render_financial_dashboard
# from pages.decision_simulator import render_decision_simulator
# from pages.fund_monitoring import render_fund_monitoring
# from pages.advisor import render_ai_financial_advisor

# Constants
DEFAULT_GROWTH_RATE = 0.08  # 8% monthly growth
DEFAULT_BURN_RATE = 85000   # $85,000 monthly burn
ENGINEER_SALARY = 10000     # $10,000 monthly cost per engineer ($120K/year)

# Initialize session state variables
if 'startups' not in st.session_state:
    st.session_state.startups = {}  # Dictionary to store multiple startup data
if 'current_startup' not in st.session_state:
    st.session_state.current_startup = None  # Currently selected startup
if 'current_page' not in st.session_state:
    st.session_state.current_page = 'upload'  # Default page
if 'insights_cache' not in st.session_state:
    st.session_state.insights_cache = {}

def switch_page(page_name):
    """Function to switch between pages"""
    st.session_state.current_page = page_name
    st.rerun()

# Page config
st.set_page_config(
    page_title="StartupFinancePilot",
    page_icon="💰",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS
st.markdown("""
<style>
    #MainMenu {visibility: hidden;}
    footer {visibility: hidden;}
    .stDeployButton {display:none;}
    
    .main-header {
        font-size: 2.5rem;
        color: #0066cc;
        margin-bottom: 0.5rem;
    }
    .sub-header {
        font-size: 1.5rem;
        color: #5c5c5c;
        margin-bottom: 1.5rem;
    }
    .metric-card {
        background-color: #f8f9fa;
        border-radius: 10px;
        padding: 20px;
        box-shadow: 0 4px 6px rgba(0,0,0,0.1);
    }
    .metric-label {
        font-size: 1rem;
        color: #5c5c5c;
    }
    .metric-value {
        font-size: 1.8rem;
        color: #0066cc;
        font-weight: bold;
    }
    .good-metric {
        color: #28a745;
    }
    .warning-metric {
        color: #ffc107;
    }
    .danger-metric {
        color: #dc3545;
    }
    
    /* Style for sidebar buttons */
    div.stButton > button {
        width: 100%;
        padding: 10px 10px;
        border: none;
        background-color: #E6F3FF;
        color: #0066cc;
        border-radius: 10px;
        text-align: left;
        margin: 5px 0;
        font-weight: bold;
    }
    
    div.stButton > button:hover {
        background-color: #CCE5FF;
        color: #004080;
    }
    
    /* Style for title box */
    .title-box {
        background: linear-gradient(45deg, #0066cc, #66b3ff);
        padding: 20px;
        border-radius: 10px;
        margin-bottom: 20px;
        text-align: center;
        color: white;
        cursor: pointer;
    }
</style>
""", unsafe_allow_html=True)

# Sample data
def load_sample_data():
    """Load sample data for demonstration"""
    # TechHealth AI data
    startup_data = {
        "name": "TechHealth AI",
        "stage": "Seed",
        "founded": "18 months ago",
        "employees": 12,
        "last_funding": "$1.2M seed round 10 months ago",
        "cash": 320000,
        "burn_rate": 85000,
        "revenue": 15000,
        "growth_rate": 0.08
    }
    
    # Cash flow history
    cash_flow_data = {
        "Month": [f"Month {i}" for i in range(1, 11)],
        "Revenue": [8000, 8500, 9200, 10000, 10800, 11700, 12600, 13600, 14700, 15800],
        "Payroll": [60000, 60000, 62000, 62000, 65000, 65000, 70000, 70000, 75000, 75000],
        "Marketing": [8000, 9000, 10000, 12000, 15000, 18000, 15000, 12000, 10000, 8000],
        "Office": [5000, 5000, 5000, 5000, 5000, 5000, 5000, 5000, 5000, 5000],
        "Software": [3000, 3200, 3500, 3800, 4000, 4200, 4500, 4800, 5000, 5200],
        "Travel": [2000, 1800, 2500, 3000, 4000, 4500, 3500, 3000, 2500, 2000],
        "Legal": [1500, 1000, 800, 1200, 800, 2000, 1500, 1000, 3000, 1200],
        "Misc": [1000, 1200, 1300, 1500, 1700, 1800, 2000, 2200, 2500, 2800]
    }
    
    # Add calculated fields
    df = pd.DataFrame(cash_flow_data)
    df["Total_Expenses"] = df[["Payroll", "Marketing", "Office", "Software", "Travel", "Legal", "Misc"]].sum(axis=1)
    df["Net_Burn"] = df["Total_Expenses"] - df["Revenue"]
    
    # Transaction data
    transactions = pd.DataFrame([
        {"Date": "2023-11-05", "Category": "Travel", "Vendor": "Caribbean Cruises", "Amount": 8500, "Description": "Team Retreat Planning", "Flag": "Suspicious"},
        {"Date": "2023-11-12", "Category": "Marketing", "Vendor": "LuxuryGifts Inc", "Amount": 4200, "Description": "Client Appreciation", "Flag": "Suspicious"},
        {"Date": "2023-11-22", "Category": "Office", "Vendor": "Premium Furniture", "Amount": 12000, "Description": "Office Upgrades", "Flag": "Suspicious"},
        {"Date": "2023-11-28", "Category": "Consulting", "Vendor": "Strategic Vision LLC", "Amount": 7500, "Description": "Strategy Consulting", "Flag": "Suspicious"},
        {"Date": "2023-12-05", "Category": "Software", "Vendor": "Personal Apple Store", "Amount": 3200, "Description": "Development Tools", "Flag": "Suspicious"},
        {"Date": "2023-12-12", "Category": "Legal", "Vendor": "Anderson Brothers", "Amount": 5800, "Description": "Legal Services", "Flag": "Normal"},
        {"Date": "2023-12-20", "Category": "Payroll", "Vendor": "November Payroll", "Amount": 75000, "Description": "Monthly Payroll", "Flag": "Normal"},
        {"Date": "2023-12-22", "Category": "Marketing", "Vendor": "Google Ads", "Amount": 8000, "Description": "Ad Campaign", "Flag": "Normal"},
        {"Date": "2023-12-25", "Category": "Office", "Vendor": "WeWork", "Amount": 5000, "Description": "Monthly Rent", "Flag": "Normal"},
        {"Date": "2023-12-28", "Category": "Software", "Vendor": "AWS", "Amount": 5200, "Description": "Cloud Services", "Flag": "Normal"}
    ])
    
    return startup_data, df, transactions

# Parse CSV file to dataframe
def parse_csv_to_df(file):
    """Parse uploaded CSV file to Pandas DataFrame"""
    try:
        df = pd.read_csv(file)
        return df, None
    except Exception as e:
        return None, f"Error parsing CSV: {e}"

# Upload and process financial data files
def render_upload_page():
    """Render the upload page for startup data"""
    st.markdown("<h1 class='main-header'>Upload Your Startup Data</h1>", unsafe_allow_html=True)
    st.markdown("<p class='sub-header'>Upload CSV files or use sample data to get started</p>", unsafe_allow_html=True)
    
    with st.expander("Upload Instructions", expanded=False):
        st.markdown("""
        ### How to Upload Your Startup Data
        
        You can upload three types of files:
        
        1. **Company Profile** - A CSV with basic information about your startup including:
           - name, stage, founded, employees, last_funding, cash, burn_rate, revenue, growth_rate
        
        2. **Cash Flow Data** - A CSV with monthly cash flow data with columns:
           - Month, Revenue, Payroll, Marketing, Office, Software, Travel, Legal, Misc
        
        3. **Transaction Data** - A CSV with transaction details:
           - Date, Category, Vendor, Amount, Description, Flag
           
        If you don't have these files ready, you can use our sample data.
        """)
    
    col1, col2 = st.columns(2)
    
    with col1:
        startup_name = st.text_input("Startup Name", value="My Startup")
        
        profile_file = st.file_uploader("Upload Company Profile (CSV)", type=['csv'])
        cash_flow_file = st.file_uploader("Upload Cash Flow Data (CSV)", type=['csv'])
        transactions_file = st.file_uploader("Upload Transactions Data (CSV)", type=['csv'])
    
    with col2:
        st.markdown("""
        <div style="background-color: #f0f7ff; padding: 15px; border-radius: 10px; height: 90%;">
            <h4>Why Upload Your Data?</h4>
            <p>By uploading your actual financial data, you'll get:</p>
            <ul>
                <li>Personalized AI insights tailored to your startup</li>
                <li>Accurate runway projections based on your real spending patterns</li>
                <li>Custom recommendations to optimize your burn rate</li>
                <li>More realistic decision simulations</li>
            </ul>
            <p>All data is processed securely and never stored permanently.</p>
        </div>
        """, unsafe_allow_html=True)
    
    # Process the files if uploaded
    if st.button("Process Data"):
        # Initialize with default values
        startup_data = {
            "name": startup_name,
            "stage": "Seed",
            "founded": "12 months ago",
            "employees": 5,
            "last_funding": "Not specified",
            "cash": 100000,
            "burn_rate": 20000,
            "revenue": 5000,
            "growth_rate": 0.05
        }
        
        cash_flow_df = None
        transactions_df = None
        
        # Parse company profile
        if profile_file:
            try:
                profile_df, error = parse_csv_to_df(profile_file)
                if error:
                    st.error(error)
                else:
                    # Get the first row as a dictionary
                    if len(profile_df) > 0:
                        startup_data.update(profile_df.iloc[0].to_dict())
                        st.success(f"Successfully loaded company profile for {startup_data['name']}")
            except Exception as e:
                st.error(f"Error processing company profile: {e}")
        
        # Parse cash flow data
        if cash_flow_file:
            cash_flow_df, error = parse_csv_to_df(cash_flow_file)
            if error:
                st.error(error)
            else:
                # Add calculated fields if not present
                if "Total_Expenses" not in cash_flow_df.columns:
                    expense_columns = [col for col in cash_flow_df.columns if col not in ["Month", "Revenue", "Total_Expenses", "Net_Burn"]]
                    cash_flow_df["Total_Expenses"] = cash_flow_df[expense_columns].sum(axis=1)
                
                if "Net_Burn" not in cash_flow_df.columns:
                    cash_flow_df["Net_Burn"] = cash_flow_df["Total_Expenses"] - cash_flow_df["Revenue"]
                
                st.success("Successfully loaded cash flow data")
        
        # Parse transactions data
        if transactions_file:
            transactions_df, error = parse_csv_to_df(transactions_file)
            if error:
                st.error(error)
            else:
                # Ensure transactions data has required columns
                required_columns = ["Date", "Category", "Vendor", "Amount", "Description"]
                if all(col in transactions_df.columns for col in required_columns):
                    if "Flag" not in transactions_df.columns:
                        transactions_df["Flag"] = "Normal"  # Default flag
                    
                    st.success("Successfully loaded transactions data")
                else:
                    st.error("Transactions file is missing required columns")
        
        # If any files were processed, save the data to session state
        if profile_file or cash_flow_file or transactions_file:
            if cash_flow_df is None:
                # Create a sample cash flow dataframe if none was uploaded
                cash_flow_data = {
                    "Month": [f"Month {i}" for i in range(1, 7)],
                    "Revenue": [startup_data['revenue'] * (1 + startup_data['growth_rate'])**i for i in range(6)],
                    "Payroll": [startup_data['burn_rate'] * 0.7] * 6,
                    "Marketing": [startup_data['burn_rate'] * 0.15] * 6,
                    "Office": [startup_data['burn_rate'] * 0.05] * 6,
                    "Software": [startup_data['burn_rate'] * 0.03] * 6,
                    "Travel": [startup_data['burn_rate'] * 0.02] * 6,
                    "Legal": [startup_data['burn_rate'] * 0.01] * 6,
                    "Misc": [startup_data['burn_rate'] * 0.04] * 6
                }
                cash_flow_df = pd.DataFrame(cash_flow_data)
                cash_flow_df["Total_Expenses"] = cash_flow_df[["Payroll", "Marketing", "Office", "Software", "Travel", "Legal", "Misc"]].sum(axis=1)
                cash_flow_df["Net_Burn"] = cash_flow_df["Total_Expenses"] - cash_flow_df["Revenue"]
            
            if transactions_df is None:
                # Create a sample transactions dataframe if none was uploaded
                transactions_data = {
                    "Date": [(datetime.now() - timedelta(days=i*5)).strftime("%Y-%m-%d") for i in range(10)],
                    "Category": ["Payroll", "Marketing", "Office", "Software", "Travel", "Legal", "Misc", "Payroll", "Marketing", "Office"],
                    "Vendor": ["Payroll Provider", "Facebook Ads", "Office Rent", "AWS", "Travel Agency", "Legal Firm", "Miscellaneous", "Payroll Provider", "Google Ads", "Office Supplies"],
                    "Amount": [startup_data['burn_rate'] * 0.7, startup_data['burn_rate'] * 0.15, startup_data['burn_rate'] * 0.05, startup_data['burn_rate'] * 0.03, startup_data['burn_rate'] * 0.02, startup_data['burn_rate'] * 0.01, startup_data['burn_rate'] * 0.04, startup_data['burn_rate'] * 0.7, startup_data['burn_rate'] * 0.15, startup_data['burn_rate'] * 0.05],
                    "Description": ["Monthly Payroll", "Ad Campaign", "Monthly Rent", "Cloud Services", "Business Travel", "Legal Services", "Miscellaneous Expenses", "Monthly Payroll", "Ad Campaign", "Office Supplies"],
                    "Flag": ["Normal", "Normal", "Normal", "Normal", "Normal", "Normal", "Normal", "Normal", "Normal", "Normal"]
                }
                transactions_df = pd.DataFrame(transactions_data)
            
            # Store in session state
            st.session_state.startups[startup_data['name']] = {
                'profile': startup_data,
                'cash_flow': cash_flow_df,
                'transactions': transactions_df
            }
            
            # Set as current startup
            st.session_state.current_startup = startup_data['name']
            
            st.success(f"Successfully added {startup_data['name']} to your startups")
            st.info("You can now analyze this startup's data in the dashboard")
            
            # Redirect to dashboard
            switch_page('dashboard')
    
    # Sample data options
    st.subheader("Or Use Sample Data")
    
    sample_col1, sample_col2 = st.columns(2)
    
    with sample_col1:
        if st.button("Use TechHealth AI Sample"):
            # Load sample data
            startup_data, cash_flow_df, transactions_df = load_sample_data()
            
            # Store in session state
            st.session_state.startups["TechHealth AI"] = {
                'profile': startup_data,
                'cash_flow': cash_flow_df,
                'transactions': transactions_df
            }
            
            # Set as current startup
            st.session_state.current_startup = "TechHealth AI"
            
            st.success("Successfully loaded TechHealth AI sample data")
            # Redirect to dashboard
            switch_page('dashboard')
    
    with sample_col2:
        if st.button("Use Custom Sample"):
            # Create a custom sample
            startup_data = {
                "name": "GreenTech Innovations",
                "stage": "Series A",
                "founded": "3 years ago",
                "employees": 25,
                "last_funding": "$4.5M Series A 8 months ago",
                "cash": 2800000,
                "burn_rate": 220000,
                "revenue": 75000,
                "growth_rate": 0.12
            }
            
            # Sample cash flow data
            cash_flow_data = {
                "Month": [f"Month {i}" for i in range(1, 11)],
                "Revenue": [45000, 48000, 52000, 57000, 62000, 66000, 70000, 72000, 74000, 75000],
                "Payroll": [140000, 142000, 145000, 150000, 160000, 165000, 175000, 180000, 185000, 190000],
                "Marketing": [25000, 28000, 30000, 32000, 35000, 32000, 30000, 28000, 26000, 25000],
                "Office": [12000, 12000, 12000, 12000, 12000, 12000, 12000, 12000, 12000, 12000],
                "Software": [8000, 8500, 9000, 9500, 10000, 10500, 11000, 11500, 12000, 12500],
                "Travel": [5000, 6000, 7000, 8000, 7000, 6000, 5000, 6000, 7000, 8000],
                "Legal": [4000, 3000, 3500, 2500, 3000, 4000, 3500, 3000, 2500, 3000],
                "Misc": [3000, 3500, 4000, 4500, 5000, 5500, 6000, 6500, 7000, 7500]
            }
            
            cash_flow_df = pd.DataFrame(cash_flow_data)
            cash_flow_df["Total_Expenses"] = cash_flow_df[["Payroll", "Marketing", "Office", "Software", "Travel", "Legal", "Misc"]].sum(axis=1)
            cash_flow_df["Net_Burn"] = cash_flow_df["Total_Expenses"] - cash_flow_df["Revenue"]
            
            # Sample transaction data
            transactions_df = pd.DataFrame([
                {"Date": "2023-11-10", "Category": "Travel", "Vendor": "First Class Flights", "Amount": 12000, "Description": "Executive Retreat", "Flag": "Suspicious"},
                {"Date": "2023-11-18", "Category": "Marketing", "Vendor": "VIP Events Co", "Amount": 18000, "Description": "Investor Dinner", "Flag": "Suspicious"},
                {"Date": "2023-12-01", "Category": "Office", "Vendor": "Luxury Furniture", "Amount": 25000, "Description": "Executive Office Upgrade", "Flag": "Suspicious"},
                {"Date": "2023-12-15", "Category": "Legal", "Vendor": "Premium Law Group", "Amount": 35000, "Description": "Legal Consultation", "Flag": "Normal"},
                {"Date": "2023-12-20", "Category": "Payroll", "Vendor": "December Payroll", "Amount": 190000, "Description": "Monthly Payroll", "Flag": "Normal"}
            ])
            
            # Store in session state
            st.session_state.startups["GreenTech Innovations"] = {
                'profile': startup_data,
                'cash_flow': cash_flow_df,
                'transactions': transactions_df
            }
            
            # Set as current startup
            st.session_state.current_startup = "GreenTech Innovations"
            
            st.success("Successfully loaded GreenTech Innovations sample data")
            # Redirect to dashboard
            switch_page('dashboard')

# Create sidebar navigation
def create_sidebar():
    with st.sidebar:
        # Title box that works as home button
        st.markdown("""
            <div class="title-box" onclick="window.location.href='#'">
                <h1>💰 StartupFinancePilot</h1>
                <p>AI-powered financial assistant for startups</p>
            </div>
        """, unsafe_allow_html=True)

        # Startup selector (if there are startups in the session state)
        if st.session_state.startups:
            st.subheader("Selected Startup")
            startup_names = list(st.session_state.startups.keys())
            selected_startup = st.selectbox(
                "Choose Startup", 
                startup_names,
                index=startup_names.index(st.session_state.current_startup) if st.session_state.current_startup in startup_names else 0
            )
            st.session_state.current_startup = selected_startup
            
            # Show basic startup info
            if selected_startup in st.session_state.startups:
                startup_data = st.session_state.startups[selected_startup]['profile']
                st.markdown(f"""
                **Stage:** {startup_data['stage']}  
                **Cash:** ${startup_data['cash']:,}  
                **Monthly Burn:** ${startup_data['burn_rate']:,}  
                **Monthly Revenue:** ${startup_data['revenue']:,}  
                """)

        st.markdown("<hr>", unsafe_allow_html=True)  # Divider
        
        # Upload data button at the top
        if st.button("📤 Upload Startup Data", use_container_width=True):
            switch_page('upload')

        # Navigation buttons
        if st.button("📊 Financial Dashboard", use_container_width=True):
            switch_page('dashboard')
            
        if st.button("🔮 Decision Simulator", use_container_width=True):
            switch_page('simulator')
            
        if st.button("🕵️ Fund Monitoring", use_container_width=True):
            switch_page('monitoring')
            
        if st.button("🤖 AI Financial Advisor", use_container_width=True):
            switch_page('advisor')

# Main function
def main():
    # Create sidebar navigation
    create_sidebar()
    
    # Check if we have any data to work with
    if not st.session_state.startups and st.session_state.current_page != 'upload':
        st.warning("No startup data found. Please upload your data or use a sample dataset.")
        render_upload_page()
        return
    
    # Render the correct page based on session state
    if st.session_state.current_page == 'upload':
        render_upload_page()
    elif st.session_state.current_page == 'dashboard':
        # Replace with your dashboard page function
        from dashboard import render_financial_dashboard
        render_financial_dashboard(
            st.session_state.startups[st.session_state.current_startup]['profile'],
            st.session_state.startups[st.session_state.current_startup]['cash_flow']
        )
    elif st.session_state.current_page == 'simulator':
        # Replace with your simulator page function
        from decision_simulator import render_decision_simulator
        render_decision_simulator(
            st.session_state.startups[st.session_state.current_startup]['profile']
        )
    elif st.session_state.current_page == 'monitoring':
        # Replace with your monitoring page function
        from fund_monitoring import render_fund_monitoring
        render_fund_monitoring(
            st.session_state.startups[st.session_state.current_startup]['transactions']
        )
    elif st.session_state.current_page == 'advisor':
        # Replace with your advisor page function
        from advisor import render_ai_financial_advisor
        render_ai_financial_advisor(
            st.session_state.startups[st.session_state.current_startup]['profile']
        )

if __name__ == "__main__":
    main()