import streamlit as st
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import altair as alt
import google.generativeai as genai
from datetime import datetime
import os
import re
import json

# App title and configuration
st.set_page_config(page_title="Expense Tracker", layout="wide")

# Initialize session state
if 'expenses' not in st.session_state:
    st.session_state.expenses = []
if 'df' not in st.session_state:
    st.session_state.df = pd.DataFrame(columns=['Date', 'Category', 'Amount', 'Description'])
if 'chat_history' not in st.session_state:
    st.session_state.chat_history = []

# Load Gemini API key from secrets
def configure_genai():
    # For local development, use st.secrets
    # For Hugging Face deployment, use environment variables
    if 'GEMINI_API_KEY' in st.secrets:
        api_key = st.secrets['GEMINI_API_KEY']
    else:
        api_key = os.environ.get('GEMINI_API_KEY')
    
    if not api_key:
        st.error("Gemini API key not found. Please add it to the secrets or environment variables.")
        st.stop()
    
    genai.configure(api_key=api_key)
    return genai.GenerativeModel('gemini-2.0-flash')

model = configure_genai()

# Function to extract expense data using Gemini
def extract_expense_data(text):
    prompt = f"""
    Extract expense information from the following text. 
    Return a JSON object with these fields:
    - date: in YYYY-MM-DD format (use today's date if not specified)
    - category: the expense category (e.g., food, transport, entertainment)
    - amount: the numerical amount (just the number, no currency symbol)
    - description: brief description of the expense

    Example output format:
    {{
        "date": "2025-03-19",
        "category": "food",
        "amount": 25.50,
        "description": "lunch at cafe"
    }}

    If multiple expenses are mentioned, return an array of such objects.
    
    Text: {text}
    """
    
    try:
        response = model.generate_content(prompt)
        response_text = response.text
        
        # Extract JSON from the response
        json_match = re.search(r'```json\n(.*?)```', response_text, re.DOTALL)
        if json_match:
            json_str = json_match.group(1)
        else:
            # If no code block, try to find JSON directly
            json_str = response_text
        
        # Parse the JSON
        data = json.loads(json_str)
        return data
    except Exception as e:
        st.error(f"Error extracting expense data: {e}")
        return None

# Function to add expenses to the dataframe
def add_expense_to_df(expense_data):
    if isinstance(expense_data, list):
        # Handle multiple expenses
        for expense in expense_data:
            add_single_expense(expense)
    else:
        # Handle single expense
        add_single_expense(expense_data)
    
    # Sort by date
    st.session_state.df = st.session_state.df.sort_values(by='Date', ascending=False)

def add_single_expense(expense):
    # Convert amount to float
    try:
        amount = float(expense['amount'])
    except:
        amount = 0.0
    
    # Create a new row
    new_row = pd.DataFrame({
        'Date': [expense.get('date', datetime.now().strftime('%Y-%m-%d'))],
        'Category': [expense.get('category', 'Other')],
        'Amount': [amount],
        'Description': [expense.get('description', '')]
    })
    
    # Append to the dataframe
    st.session_state.df = pd.concat([st.session_state.df, new_row], ignore_index=True)

# Function to get AI insights about expenses
def get_expense_insights(query):
    if st.session_state.df.empty:
        return "No expense data available yet. Please add some expenses first."
    
    # Convert dataframe to string representation
    df_str = st.session_state.df.to_string()
    
    prompt = f"""
    Here is a dataset of expenses:
    {df_str}
    
    User query: {query}
    
    Please analyze this expense data and answer the query.
    Provide your analysis in a clear and concise way.
    If the query is about visualizations, describe what kind of chart would be helpful.
    """
    
    try:
        response = model.generate_content(prompt)
        return response.text
    except Exception as e:
        return f"Error getting insights: {e}"

# Function to create visualizations
def create_visualizations():
    if st.session_state.df.empty:
        st.info("Add some expenses to see visualizations")
        return
    
    # Create a copy of the dataframe for visualization
    df = st.session_state.df.copy()
    
    # Ensure Date is datetime
    df['Date'] = pd.to_datetime(df['Date'])
    
    # Create tabs for different visualizations
    tab1, tab2, tab3 = st.tabs(["Expenses by Category", "Expenses Over Time", "Recent Expenses"])
    
    with tab1:
        st.subheader("Expenses by Category")
        category_totals = df.groupby('Category')['Amount'].sum().reset_index()
        
        # Create a pie chart
        fig, ax = plt.subplots(figsize=(8, 8))
        ax.pie(category_totals['Amount'], labels=category_totals['Category'], autopct='%1.1f%%')
        ax.set_title('Expenses by Category')
        st.pyplot(fig)
        
        # Create a bar chart
        category_chart = alt.Chart(category_totals).mark_bar().encode(
            x=alt.X('Category:N', sort='-y'),
            y=alt.Y('Amount:Q'),
            color='Category:N'
        ).properties(
            title='Total Expenses by Category'
        )
        st.altair_chart(category_chart, use_container_width=True)
    
    with tab2:
        st.subheader("Expenses Over Time")
        # Group by date and sum amounts
        daily_totals = df.groupby(df['Date'].dt.date)['Amount'].sum().reset_index()
        
        # Create a line chart
        time_chart = alt.Chart(daily_totals).mark_line(point=True).encode(
            x='Date:T',
            y='Amount:Q',
            tooltip=['Date:T', 'Amount:Q']
        ).properties(
            title='Daily Expenses Over Time'
        )
        st.altair_chart(time_chart, use_container_width=True)
    
    with tab3:
        st.subheader("Recent Expenses")
        # Sort by date and get the last 10 expenses
        recent = df.sort_values('Date', ascending=False).head(10)
        
        # Create a bar chart
        recent_chart = alt.Chart(recent).mark_bar().encode(
            x=alt.X('Description:N', sort='-y'),
            y='Amount:Q',
            color='Category:N',
            tooltip=['Date:T', 'Category:N', 'Amount:Q', 'Description:N']
        ).properties(
            title='Most Recent Expenses'
        )
        st.altair_chart(recent_chart, use_container_width=True)

# App layout
st.title("💰 Expense Tracker with AI")

# Sidebar for app navigation
page = st.sidebar.radio("Navigation", ["Add Expenses", "View & Analyze", "Chat with your Data"])

if page == "Add Expenses":
    st.header("Add Your Expenses")
    st.write("Describe your expenses in natural language, and AI will extract the details.")
    
    with st.form("expense_form"):
        user_input = st.text_area(
            "Enter your expenses:", 
            height=100,
            placeholder="Example: I spent $25 on lunch today, $15 on transport yesterday, and $50 on groceries on March 15th"
        )
        submit_button = st.form_submit_button("Add Expenses")
    
    if submit_button and user_input:
        with st.spinner("Processing your expenses..."):
            expense_data = extract_expense_data(user_input)
            
            if expense_data:
                add_expense_to_df(expense_data)
                st.success("Expenses added successfully!")
                st.write("Extracted information:")
                st.json(expense_data)
            else:
                st.error("Failed to extract expense data. Please try again with a clearer description.")
    
    # Show the current expenses
    if not st.session_state.df.empty:
        st.subheader("Your Recent Expenses")
        st.dataframe(st.session_state.df.sort_values(by='Date', ascending=False), use_container_width=True)

elif page == "View & Analyze":
    st.header("Your Expense Data")
    
    # Show the current expenses as a table
    if not st.session_state.df.empty:
        st.dataframe(st.session_state.df.sort_values(by='Date', ascending=False), use_container_width=True)
        
        # Add download button
        csv = st.session_state.df.to_csv(index=False)
        st.download_button(
            label="Download CSV",
            data=csv,
            file_name="expenses.csv",
            mime="text/csv"
        )
        
        # Show summary statistics
        st.subheader("Summary Statistics")
        col1, col2, col3 = st.columns(3)
        with col1:
            st.metric("Total Expenses", f"${st.session_state.df['Amount'].sum():.2f}")
        with col2:
            st.metric("Average Expense", f"${st.session_state.df['Amount'].mean():.2f}")
        with col3:
            st.metric("Number of Expenses", f"{len(st.session_state.df)}")
        
        # Create visualizations
        st.subheader("Visualizations")
        create_visualizations()
    else:
        st.info("No expense data available yet. Please add some expenses first.")

elif page == "Chat with your Data":
    st.header("Chat with Your Expense Data")
    
    if st.session_state.df.empty:
        st.info("No expense data available yet. Please add some expenses first.")
    else:
        st.write("Ask questions about your expenses to get insights.")
        
        # Display chat history
        for message in st.session_state.chat_history:
            with st.chat_message(message["role"]):
                st.write(message["content"])
        
        # Get user input
        user_query = st.chat_input("Ask about your expenses...")
        
        if user_query:
            # Add user message to chat history
            st.session_state.chat_history.append({"role": "user", "content": user_query})
            
            # Display user message
            with st.chat_message("user"):
                st.write(user_query)
            
            # Get AI response
            with st.spinner("Thinking..."):
                response = get_expense_insights(user_query)
            
            # Add AI response to chat history
            st.session_state.chat_history.append({"role": "assistant", "content": response})
            
            # Display AI response
            with st.chat_message("assistant"):
                st.write(response)

# Add instructions for Hugging Face deployment in the sidebar
with st.sidebar.expander("Deployment Instructions"):
    st.write("""
    ### How to deploy to Hugging Face:
    
    1. Save this code as `app.py`
    2. Create a `requirements.txt` file with these dependencies:
       ```
       streamlit
       pandas
       matplotlib
       seaborn
       altair
       google-generativeai
       ```
    3. Create a `README.md` file describing your app
    4. Add your Gemini API key to your Hugging Face Space secrets with the name `GEMINI_API_KEY`
    5. Push your code to a GitHub repository
    6. Create a new Hugging Face Space, select Streamlit as the SDK, and connect your GitHub repository
    """)

# Bottom credits
st.sidebar.markdown("---")
st.sidebar.caption("Built with Streamlit and Gemini AI")