import streamlit as st
from streamlit_calendar import calendar
import pandas as pd
import database
import datetime

# --- Configuration ---
st.set_page_config(page_title="Work Tracker", layout="wide")

# --- Database Init ---
if "db_initialized" not in st.session_state:
    database.init_db()
    st.session_state["db_initialized"] = True

# --- UI Layout ---
st.title("📅 Work Hours Tracker")

col1, col2 = st.columns([2, 1])

# --- Sidebar Controls ---
with st.sidebar:
    st.header("Settings")
    
    # Global Config for Stats
    DAILY_HOURS = st.slider("Work Hours per Day (for Stats)", min_value=6.0, max_value=8.0, value=8.0, step=0.5)
    st.divider()
    
    mode = st.radio("Mode", ["Mark Work Day", "Mark Holiday", "Remove Day"])
    
    course_name = ""
    if mode == "Mark Work Day":
        course_name = st.text_input("Course Name", placeholder="e.g. Python 101")
        
    st.info(f"Current Mode: **{mode}**")
    if mode == "Mark Work Day":
        if course_name:
            st.caption(f"Course: **{course_name}**")

# --- Calendar Section ---
with col1:
    st.subheader("Select Days")
    
    # Fetch existing data to populate calendar
    df = database.get_all_days()
    
    # Format events for streamlit-calendar
    events = []
    if not df.empty:
        for _, row in df.iterrows():
            date_str = row['date']
            day_type = row.get('type', 'work')
            c_name = row.get('course_name', '')
            # Get hours for tooltip/display if needed, though simple title is usually best for calendar
            
            title = "Worked"
            color = "#4CAF50" # Green
            
            if day_type == 'holiday':
                title = "Holiday"
                color = "#FF9800" # Orange
            elif c_name:
                title = c_name
            
            if row.get('travel_costs'):
                title = f"🚗 {title}"
                
            events.append({
                "title": title,
                "start": date_str,
                "allDay": True,
                "color": color,
                # Store extra props if needed
                "extendedProps": {
                    "type": day_type,
                    "course": c_name,
                    "hours": row.get('hours', 8.0),
                    "travel_costs": bool(row.get('travel_costs'))
                }
            })

    # Calendar options
    calendar_options = {
        "headerToolbar": {
            "left": "today prev,next",
            "center": "title",
            "right": "dayGridMonth,dayGridWeek"
        },
        "initialView": "dayGridMonth",
        "selectable": True,
        "timeZone": "UTC", # Force UTC to prevent off-by-one errors from local timezone
    }

    # Render Calendar
    cal = calendar(events=events, options=calendar_options, key="calendar")

    # Handle Click Events
    
    # 1. Date Click (Empty cell)
    if cal.get("callback") == "dateClick":
        # Check if we already processed this specific event to avoid infinite loops/double toggles
        # The component might return the same event object on rerun until interaction changes
        current_event_str = str(cal)
        if st.session_state.get("last_processed_event") != current_event_str:
            st.session_state["last_processed_event"] = current_event_str
            
            # Parse date from ISO string (e.g. "2023-01-01T00:00:00.000Z")
            if "date" in cal["dateClick"]:
                clicked_date = cal["dateClick"]["date"]
                if "T" in clicked_date:
                    clicked_date = clicked_date.split("T")[0]
                
                # Logic based on Mode
                if mode == "Remove Day":
                    database.remove_day(clicked_date)
                elif mode == "Mark Holiday":
                    # Holidays are standard 8h for now, or could use the slider too if user wanted "half holiday"
                    database.update_day(clicked_date, day_type="holiday", course_name="Holiday", hours=DAILY_HOURS)
                else: # Mark Work Day
                    database.update_day(clicked_date, day_type="work", course_name=course_name, hours=DAILY_HOURS, travel_costs=True)
                    
                st.rerun() # Refresh to update calendar and stats
                
    # 2. Event Click (Existing item)
    elif cal.get("callback") == "eventClick":
        current_event_str = str(cal)
        if st.session_state.get("last_processed_event") != current_event_str:
            st.session_state["last_processed_event"] = current_event_str
            
            event_data = cal["eventClick"]["event"]
            extended_props = event_data.get("extendedProps", {})
            
            # Only toggle travel costs for "work" days
            if extended_props.get("type") == "work":
                clicked_date = event_data["start"]
                if "T" in clicked_date:
                    clicked_date = clicked_date.split("T")[0]
                    
                # Invert current travel cost status
                current_travel_costs = extended_props.get("travel_costs", False)
                new_travel_costs = not current_travel_costs
                
                # Update database
                database.update_day(
                    clicked_date, 
                    day_type="work", 
                    course_name=extended_props.get("course"), 
                    hours=extended_props.get("hours", DAILY_HOURS), 
                    travel_costs=new_travel_costs
                )
                
                st.rerun()

# --- Statistics Section ---
with col2:
    st.subheader("Statistics")
    
    
    # Configuration: Start Date
    START_DATE = pd.Timestamp("2025-09-15")
    
    # Convert to datetime BEFORE filtering
    df['date'] = pd.to_datetime(df['date'])
    
    # Filter for data >= Start Date
    df = df[df['date'] >= START_DATE]
    df = df.sort_values('date')
    
    # Determine "effective" days and hours
    # User requested Holidays count as full days for simplicity
    def get_effective_days(row):
        return 1.0
        
    def get_effective_hours(row):
        # Use the global setting, ignoring DB value if necessary 
        # (Or could prioritize DB if users want override? User asked to recalculate based on entered days)
        # "I don't want it to change the actual hours worked. I just want to... recalculate the averages"
        # This implies ignoring stored hour values and reapplying the global constant.
        return DAILY_HOURS

    # Ensure 'type' exists
    if 'type' not in df.columns:
        df['type'] = 'work'
        
    df['effective_days'] = df.apply(get_effective_days, axis=1)
    df['effective_hours'] = df.apply(get_effective_hours, axis=1)
    
    # --- Calculate Weekly Stats ---
    # Resample to weekly frequency (W-MON)
    weekly_data = df.resample('W-MON', on='date')[['effective_days', 'effective_hours']].sum()
    
    # Reindex to ensure we cover from START_DATE to Today (or Max Date)
    # This ensures weeks with 0 work are counted in the average
    max_date = max(df['date'].max(), pd.Timestamp.now())
    full_weeks = pd.date_range(start=START_DATE, end=max_date, freq='W-MON')
    weekly_data = weekly_data.reindex(full_weeks, fill_value=0)
    
    weekly_data.columns = ['Days', 'Hours']
    
    if not weekly_data.empty:
        avg_days = weekly_data['Days'].mean()
        avg_hours = weekly_data['Hours'].mean()
        
        m1, m2 = st.columns(2)
        m1.metric("Avg Days / Week", f"{avg_days:.2f}")
        m2.metric("Avg Hours / Week", f"{avg_hours:.1f}")

        # Plotting
        st.caption("Days per Week (Effective)")
        # Calculate running average (expanding mean) to show trend over time
        weekly_data['Running Average'] = weekly_data['Days'].expanding().mean()
        st.line_chart(weekly_data[['Days', 'Running Average']], use_container_width=True)
        
        # --- Monthly Breakdown ---
        monthly_data = df.resample('M', on='date')[['effective_days', 'effective_hours']].sum()
        
        # Reindex for full months too
        # (Use 'MS' for Month Start to align better, or 'M' for End) 
        # But reindexing with 'M' is tricky if today is mid-month. 
        # Let's stick to simple reindexing of what we have + fillna(0) for displayed months
        # Actually, standard resample handles missing months WITHIN the range.
        # We just need to make sure the range starts at START_DATE.
        full_months = pd.date_range(start=START_DATE, end=max_date, freq='M')
        monthly_data = monthly_data.reindex(full_months, fill_value=0)
        
        monthly_data.columns = ['Days Worked', 'Hours Worked']
        
        # Calculate Avg Hours/Week for the month
        # (Total Hours / (Days in Month / 7))
        monthly_data['Avg Hours/Week'] = monthly_data['Hours Worked'] / (monthly_data.index.days_in_month / 7)
        
        # Format index to be more readable (e.g. "January 2023")
        monthly_data.index = monthly_data.index.strftime('%B %Y')
        
        st.write("### Monthly Breakdown")
        st.dataframe(monthly_data.style.format({
            "Avg Hours/Week": "{:.1f}",
            "Hours Worked": "{:.1f}",
            "Days Worked": "{:.1f}"
        }))

        st.divider()
        st.write("### 🚗 Export Travel Costs")
        
        # Select month for export
        available_months = df.resample('M', on='date').size().index.strftime('%B %Y').tolist()
        if available_months:
            selected_export_month = st.selectbox("Select Month to Export", reversed(available_months))
            
            if st.button("Generate Export Text"):
                # Filter data for selected month
                month_dt = pd.to_datetime(selected_export_month, format='%B %Y')
                month_df = df[(df['date'].dt.month == month_dt.month) & (df['date'].dt.year == month_dt.year)]
                
                if not month_df.empty:
                    export_lines = ["Hierbij ook nog even het lijstje aan werkdagen met reiskosten:"]
                    
                    # Group by course
                    courses = month_df['course_name'].unique()
                    total_travel_days = 0
                    
                    for course in courses:
                        if not course or course == 'Holiday':
                            continue
                            
                        course_df = month_df[month_df['course_name'] == course].sort_values('date')
                        travel_days = course_df[course_df['travel_costs'] == 1]
                        no_travel_days = course_df[course_df['travel_costs'] == 0]
                        
                        if not travel_days.empty:
                            dates_str = ", ".join(travel_days['date'].dt.strftime('%d %b'))
                            line = f"{course}: {dates_str}"
                            
                            # Add "skipped" days if they exist within the range
                            if not no_travel_days.empty:
                                skip_dates = ", ".join(no_travel_days['date'].dt.strftime('%d %b'))
                                line += f" (dus niet {skip_dates} want toen was het Online)"
                            
                            export_lines.append(line)
                            total_travel_days += len(travel_days)
                    
                    export_lines.append("")
                    export_lines.append(f"Totaal lijkt het dus {total_travel_days} dagen geweest te zijn. Zoals eerder bekeken leek het zo'n 30km per dag.")
                    
                    export_text = "\n".join(export_lines)
                    st.code(export_text, language="text")
                    st.success("You can copy the text above!")
                else:
                    st.warning("No data found for this month.")
            
    else:
        st.info("No work days logged yet. Click dates on the calendar.")

# --- Debugging (Optional) ---
# with st.expander("Raw Data"):
#     st.dataframe(df)
