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 plotly.express as px
import plotly.graph_objects as go
from datetime import datetime, timedelta, date
import time
import json
import google.generativeai as genai
from google.generativeai.types import HarmCategory, HarmBlockThreshold


# Setup AI Services
def setup_genai():
    """Initialize and configure Google's Generative AI and list available models"""
    try:
        # Try getting API key from Streamlit secrets first
        if 'GOOGLE_API_KEY' in st.secrets:
            api_key = st.secrets['GOOGLE_API_KEY']
        # Fall back to environment variable
        elif 'GOOGLE_API_KEY' in os.environ:
            api_key = os.environ['GOOGLE_API_KEY']
        else:
            st.warning("Google API key not found. Using simulated AI responses.")
            st.session_state.gemini_model = "gemini-1.5-pro"
            return False
            
        genai.configure(api_key=api_key)

# Import pages
from pages.dashboard-page import render_financial_dashboard, get_runway_analysis, get_fundraising_readiness_analysis
from pages.decision-simulator import render_decision_simulator, get_decision_analysis
from pages.fund-monitoring import render_fund_monitoring, get_fraud_analysis
from pages.financial-advisor import render_ai_financial_advisor, get_advisory_guidance, generate_voice_response
from pages.book-session import render_book_session

# Initialize page configuration
st.set_page_config(
    page_title="StartupFinancePilot",
    page_icon="💰",
    layout="wide",
    initial_sidebar_state="expanded"
)


# Initialize session state variables
if 'booked_sessions' not in st.session_state:
    st.session_state.booked_sessions = []
if 'chat_history' not in st.session_state:
    st.session_state.chat_history = []
if 'audio_response' not in st.session_state:
    st.session_state.audio_response = None
if 'insights_cache' not in st.session_state:
    st.session_state.insights_cache = {}
if 'gemini_model' not in st.session_state:
    st.session_state.gemini_model = None
if 'current_page' not in st.session_state:
    st.session_state.current_page = "Financial Dashboard"
    
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')


def main():
    # Load sample data
    startup_data, cash_flow_df, transactions_df = load_sample_data()
    
    # Create sidebar
    st.sidebar.title("StartupFinancePilot")
    st.sidebar.image("https://img.freepik.com/premium-vector/business-finance-analytics-logo-design-vector-template_67715-552.jpg", width=150)
    
    # Company profile
    st.sidebar.header("Company Profile")
    st.sidebar.write(f"**{startup_data['name']}**")
    st.sidebar.write(f"Stage: {startup_data['stage']}")
    st.sidebar.write(f"Founded: {startup_data['founded']}")
    st.sidebar.write(f"Employees: {startup_data['employees']}")
    st.sidebar.write(f"Last Funding: {startup_data['last_funding']}")
    
    
    # Navigation
    st.sidebar.header("Navigation")
    pages = {
        "Financial Dashboard": render_financial_dashboard,
        "Decision Simulator": render_decision_simulator,
        "Fund Monitoring": render_fund_monitoring,
        "AI Financial Advisor": render_ai_financial_advisor,
        "Book a Session": render_book_session
    }
    
    # Page selection
    selected_page = st.sidebar.radio("Go to", list(pages.keys()))
    st.session_state.current_page = selected_page
    
    # Render selected page
    if selected_page == "Financial Dashboard":
        pages[selected_page](startup_data, cash_flow_df)
    elif selected_page == "Decision Simulator":
        pages[selected_page](startup_data)
    elif selected_page == "Fund Monitoring":
        pages[selected_page](transactions_df)
    elif selected_page == "AI Financial Advisor":
        pages[selected_page](startup_data)
    else:  # Book a Session
        pages[selected_page]()

if __name__ == "__main__":
    main()