import streamlit as st
import pandas as pd
import numpy as np
import plotly.express as px
import plotly.graph_objects as go
import os
from datetime import datetime, timedelta, date
import time
import json
import google.generativeai as genai
from google.generativeai.types import HarmCategory, HarmBlockThreshold

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

# Custom CSS
st.markdown("""
<style>
    .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;
    }
    .advisor-card {
        background-color: #f0f7ff;
        border-radius: 10px;
        padding: 20px;
        box-shadow: 0 4px 6px rgba(0,0,0,0.1);
        margin-bottom: 20px;
    }
    .advice-text {
        font-size: 1.1rem;
        line-height: 1.6;
        color: #333;
    }
    .insight-card {
        background-color: #f0f8ff;
        border-left: 4px solid #0066cc;
        padding: 15px;
        margin-bottom: 15px;
        border-radius: 4px;
    }
    .ai-badge {
        background-color: #0066cc;
        color: white;
        padding: 3px 10px;
        border-radius: 10px;
        font-size: 0.8rem;
        margin-bottom: 10px;
        display: inline-block;
    }
    .booking-card {
        background-color: white;
        border-radius: 10px;
        padding: 20px;
        box-shadow: 0 4px 6px rgba(0,0,0,0.1);
        margin-bottom: 20px;
    }
    .session-type {
        font-size: 1.2rem;
        font-weight: bold;
        color: #0066cc;
    }
    .session-duration {
        color: #5c5c5c;
        font-size: 0.9rem;
    }
    .session-price {
        font-size: 1.1rem;
        font-weight: bold;
        color: #28a745;
    }
</style>
""", unsafe_allow_html=True)

# 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)
DEFAULT_MARKETING_BUDGET = 10000  # $10,000 monthly marketing budget

# 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

# Sample data
@st.cache_data
def load_sample_data():
    # 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"},
        {"Date": "2024-01-05", "Category": "Travel", "Vendor": "Delta Airlines", "Amount": 1200, "Description": "Client Meeting Travel", "Flag": "Normal"},
        {"Date": "2024-01-10", "Category": "Marketing", "Vendor": "Facebook Ads", "Amount": 4500, "Description": "Social Media Campaign", "Flag": "Normal"},
        {"Date": "2024-01-15", "Category": "Software", "Vendor": "Atlassian", "Amount": 2800, "Description": "Development Tools", "Flag": "Normal"},
        {"Date": "2024-01-20", "Category": "Payroll", "Vendor": "January Payroll", "Amount": 75000, "Description": "Monthly Payroll", "Flag": "Normal"},
        {"Date": "2024-01-25", "Category": "Office", "Vendor": "WeWork", "Amount": 5000, "Description": "Monthly Rent", "Flag": "Normal"}
    ])
    
    return startup_data, df, transactions

# Setup AI Services
def setup_genai():
    """Initialize and configure Google's Generative AI and list available models"""
    try:
        if 'GOOGLE_API_KEY' in st.secrets:
            genai.configure(api_key=st.secrets['GOOGLE_API_KEY'])
            
            # Get available models and select one for text generation
            models = genai.list_models()
            text_models = [m.name for m in models if 'generateContent' in m.supported_generation_methods]
            
            if text_models:
                # Use first available text generation model
                model_name = text_models[0]
                st.session_state.gemini_model = model_name
                return True
            else:
                st.warning("No appropriate generative AI models available")
                # Use a fallback model name for demonstration
                st.session_state.gemini_model = "gemini-1.5-pro"
                return False
        else:
            st.warning("Google API key not found in secrets. Using simulated AI responses.")
            st.session_state.gemini_model = "gemini-1.5-pro"
            return False
    except Exception as e:
        st.warning(f"Failed to initialize Gemini: {e}. Using simulated AI responses.")
        st.session_state.gemini_model = "gemini-1.5-pro"
        return False

def generate_ai_response(prompt, simulate=False):
    """Generate response from Gemini or simulate one if the API is unavailable"""
    if simulate:
        # Simulate AI response with predefined text based on keywords in prompt
        time.sleep(1)  # Simulate processing time
        
        if "runway" in prompt.lower():
            return "Based on your current spend rate of $85K/month and revenue growth of 8%, your runway is approximately 3.8 months. I recommend reducing non-essential expenses to extend runway to at least 6 months before your next fundraising round."
        elif "hire" in prompt.lower() or "hiring" in prompt.lower():
            return "Adding new hires at this stage would reduce your runway significantly. Consider contracting talent first or postponing hiring until after securing additional funding. Each new engineer costs $10K/month, reducing runway by approximately 3 weeks per hire."
        elif "marketing" in prompt.lower():
            return "Your current CAC to LTV ratio doesn't justify increasing marketing spend. Focus on optimizing current channels and improving conversion rates. Once unit economics improve, gradually increase marketing budget by no more than 20% per month."
        elif "fundraising" in prompt.lower() or "investor" in prompt.lower():
            return "With less than 4 months of runway, you should begin fundraising preparations immediately. Focus on demonstrating product-market fit and improving key metrics like MRR growth, user retention, and unit economics before approaching investors."
        elif "suspicious" in prompt.lower() or "transaction" in prompt.lower():
            return "I've identified several concerning transactions including a $8,500 travel expense and $12,000 in office upgrades. These discretionary expenses represent over 25% of a month's burn and should be reviewed with your finance team immediately."
        else:
            return "Based on your financial data, I recommend prioritizing runway extension and focusing on core metrics that demonstrate product-market fit. Consider reducing non-essential expenses by 15-20% to add 1-2 months to your runway before beginning fundraising conversations."
    else:
        try:
            # Use the actual Gemini model
            model = genai.GenerativeModel(st.session_state.gemini_model)
            
            generation_config = {
                "temperature": 0.7,
                "top_p": 0.95,
                "top_k": 40,
                "max_output_tokens": 1024,
            }
            
            safety_settings = {
                HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,
                HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
                HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
                HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE
            }
            
            response = model.generate_content(
                prompt,
                generation_config=generation_config,
                safety_settings=safety_settings
            )
            
            return response.text
        except Exception as e:
            st.warning(f"Error generating AI response: {e}")
            # Fall back to simulated response
            return generate_ai_response(prompt, simulate=True)

# Financial modeling functions
def calculate_runway(initial_cash, monthly_burn, monthly_revenue, growth_rate, months=24):
    """Calculate runway based on current burn rate and revenue growth."""
    dates = [datetime.now() + timedelta(days=30*i) for i in range(months)]
    df = pd.DataFrame(index=dates, columns=['Cash', 'Revenue', 'Expenses', 'Net_Burn', 'Cumulative_Cash'])
    
    current_cash = initial_cash
    current_revenue = monthly_revenue
    df.iloc[0, df.columns.get_loc('Cash')] = current_cash
    df.iloc[0, df.columns.get_loc('Revenue')] = current_revenue
    df.iloc[0, df.columns.get_loc('Expenses')] = monthly_burn
    df.iloc[0, df.columns.get_loc('Net_Burn')] = monthly_burn - current_revenue
    df.iloc[0, df.columns.get_loc('Cumulative_Cash')] = current_cash
    
    runway_months = months
    for i in range(1, months):
        current_revenue = current_revenue * (1 + growth_rate)
        net_burn = monthly_burn - current_revenue
        current_cash = current_cash - net_burn
        
        df.iloc[i, df.columns.get_loc('Cash')] = current_cash
        df.iloc[i, df.columns.get_loc('Revenue')] = current_revenue
        df.iloc[i, df.columns.get_loc('Expenses')] = monthly_burn
        df.iloc[i, df.columns.get_loc('Net_Burn')] = net_burn
        df.iloc[i, df.columns.get_loc('Cumulative_Cash')] = current_cash
        
        if current_cash <= 0:
            runway_months = i
            break
    
    return runway_months, df

def simulate_decision(initial_cash, monthly_burn, monthly_revenue, growth_rate, 
                     new_expenses=0, new_hires=0, new_marketing=0, growth_impact=0):
    """Simulate the impact of a business decision on runway."""
    # Calculate current runway
    current_runway, current_df = calculate_runway(initial_cash, monthly_burn, monthly_revenue, growth_rate)
    
    # Calculate additional expenses
    additional_expenses = new_expenses + (new_hires * ENGINEER_SALARY) + new_marketing
    
    # Calculate new runway
    new_runway, new_df = calculate_runway(
        initial_cash,
        monthly_burn + additional_expenses,
        monthly_revenue,
        growth_rate + growth_impact
    )
    
    return current_runway, new_runway, current_df, new_df

def detect_suspicious_transactions(transactions_df):
    """AI-enhanced suspicious transaction detection."""
    df = transactions_df.copy()
    
    # Define thresholds for each category
    category_thresholds = {
        "Travel": 3000,
        "Marketing": 10000,
        "Office": 7000,
        "Software": 6000,
        "Consulting": 5000,
        "Legal": 6000
    }
    
    # Define suspicious terms
    suspicious_terms = ['luxury', 'cruise', 'premium', 'personal', 'gift']
    
    # Add suspicious column
    df['Suspicious'] = False
    df['Reason'] = ""
    df['Risk_Score'] = 0
    
    # Check for suspicious patterns
    for idx, row in df.iterrows():
        reasons = []
        risk_score = 0
        
        # Check if amount exceeds category threshold
        if row['Category'] in category_thresholds:
            if row['Amount'] > category_thresholds[row['Category']]:
                reasons.append(f"Amount exceeds typical spending for {row['Category']}")
                risk_score += 30
                
                # Higher risk for significantly exceeding threshold
                excess_percentage = (row['Amount'] - category_thresholds[row['Category']]) / category_thresholds[row['Category']] * 100
                if excess_percentage > 100:  # More than double the threshold
                    risk_score += 20
        
        # Check for suspicious vendors or descriptions
        if any(term in str(row['Vendor']).lower() for term in suspicious_terms):
            reasons.append(f"Vendor name contains suspicious term")
            risk_score += 25
        
        if any(term in str(row['Description']).lower() for term in suspicious_terms):
            reasons.append(f"Description contains suspicious term")
            risk_score += 20
        
        # Check for rounded amounts (potential indicator of estimation/fabrication)
        if row['Amount'] % 1000 == 0 and row['Amount'] > 3000:
            reasons.append(f"Suspiciously round amount")
            risk_score += 15
        
        # Mark as suspicious if risk score is high enough
        if risk_score >= 30:
            df.at[idx, 'Suspicious'] = True
            df.at[idx, 'Reason'] = "; ".join(reasons)
            df.at[idx, 'Risk_Score'] = risk_score
    
    # Sort by risk score
    df = df.sort_values(by='Risk_Score', ascending=False)
    
    return df

# Import page functions
from dashboard_page import render_financial_dashboard, get_runway_analysis, get_fundraising_readiness_analysis
from decision_simulator import render_decision_simulator, get_decision_analysis
from fund_monitoring import render_fund_monitoring, get_fraud_analysis
from financial_advisor import render_ai_financial_advisor, get_advisory_guidance, generate_voice_response
from book_session import render_book_session

# UI Components
def create_sidebar():
    """Create sidebar with company profile and filters."""
    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
    startup_data, _, _ = load_sample_data()
    
    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']}")
    
    # AI Status
    has_api = setup_genai()
    ai_status = "🟢 Connected" if has_api else "🟡 Demo Mode"
    st.sidebar.write(f"AI Status: {ai_status}")
    if not has_api:
        st.sidebar.info("Running in demo mode with simulated AI responses. Add GOOGLE_API_KEY to secrets for full functionality.")
    
    # App navigation
    st.sidebar.header("Navigation")
    page = st.sidebar.radio("Go to", [
        "Financial Dashboard", 
        "Decision Simulator", 
        "Fund Monitoring", 
        "AI Financial Advisor",
        "Book a Session"
    ])
    
    return page

# Main application
def main():
    # Load sample data
    startup_data, cash_flow_df, transactions_df = load_sample_data()
    
    # Create sidebar and get selected page
    page = create_sidebar()
    
    # Render selected page
    if page == "Financial Dashboard":
        render_financial_dashboard(startup_data, cash_flow_df)
    elif page == "Decision Simulator":
        render_decision_simulator(startup_data)
    elif page == "Fund Monitoring":
        render_fund_monitoring(transactions_df)
    elif page == "AI Financial Advisor":
        render_ai_financial_advisor(startup_data)
    elif page == "Book a Session":
        render_book_session()

if __name__ == "__main__":
    main()