# app.py
import streamlit as st
import json
import os
from datetime import datetime, timedelta
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
from collections import defaultdict
import base64

# Page configuration
st.set_page_config(
    page_title="Atomic Habits Tracker Pro",
    page_icon="🔥",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Categories with emojis
CATEGORIES = {
    "Health": "💪",
    "Learning": "📚",
    "Mindfulness": "🧘",
    "Productivity": "⚡",
    "Creativity": "🎨",
    "Social": "👥",
    "Finance": "💰",
    "Other": "🔹"
}

# Initialize session state for theme
if 'dark_mode' not in st.session_state:
    st.session_state.dark_mode = False

# Theme toggle function
def toggle_theme():
    st.session_state.dark_mode = not st.session_state.dark_mode

# Dynamic CSS based on theme
def get_css():
    if st.session_state.dark_mode:
        return """
        <style>
            .main-header {
                font-size: 3rem;
                font-weight: bold;
                background: linear-gradient(90deg, #FF6B6B, #4ECDC4);
                -webkit-background-clip: text;
                -webkit-text-fill-color: transparent;
                text-align: center;
                margin-bottom: 2rem;
            }
            .stApp {
                background-color: #0e1117;
                color: #fafafa;
            }
            .quote-box {
                background: #1e1e1e;
                border-left: 5px solid #4ECDC4;
                padding: 15px;
                margin: 20px 0;
                font-style: italic;
                color: #fafafa;
                border-radius: 0 10px 10px 0;
            }
            .habit-card {
                background: #1e1e1e;
                padding: 20px;
                border-radius: 15px;
                margin: 10px 0;
                box-shadow: 0 4px 6px rgba(0,0,0,0.3);
                border: 1px solid #333;
            }
            .streak-badge {
                background: #FFD93D;
                color: #333;
                padding: 5px 15px;
                border-radius: 20px;
                font-weight: bold;
                display: inline-block;
                margin: 5px;
            }
            .category-badge {
                background: #333;
                color: #fff;
                padding: 3px 10px;
                border-radius: 15px;
                font-size: 0.8rem;
                display: inline-block;
                margin: 5px 0;
            }
            .stButton>button {
                border-radius: 25px;
                transition: all 0.3s;
            }
            .stTextInput>div>div>input, .stSelectbox>div>div>select {
                background-color: #1e1e1e;
                color: #fafafa;
                border-radius: 10px;
            }
            .stProgress > div > div > div > div {
                background-color: #4ECDC4;
            }
            div[data-testid="stSidebar"] {
                background-color: #1e1e1e;
            }
            h1, h2, h3, h4, h5, h6, p, label {
                color: #fafafa !important;
            }
            .footer-text {
                color: #888 !important;
            }
        </style>
        """
    else:
        return """
        <style>
            .main-header {
                font-size: 3rem;
                font-weight: bold;
                background: linear-gradient(90deg, #FF6B6B, #4ECDC4);
                -webkit-background-clip: text;
                -webkit-text-fill-color: transparent;
                text-align: center;
                margin-bottom: 2rem;
            }
            .quote-box {
                background: #f0f2f6;
                border-left: 5px solid #4ECDC4;
                padding: 15px;
                margin: 20px 0;
                font-style: italic;
                border-radius: 0 10px 10px 0;
            }
            .habit-card {
                background: white;
                padding: 20px;
                border-radius: 15px;
                margin: 10px 0;
                box-shadow: 0 4px 6px rgba(0,0,0,0.1);
            }
            .streak-badge {
                background: #FFD93D;
                color: #333;
                padding: 5px 15px;
                border-radius: 20px;
                font-weight: bold;
                display: inline-block;
                margin: 5px;
            }
            .category-badge {
                background: #e0e0e0;
                color: #333;
                padding: 3px 10px;
                border-radius: 15px;
                font-size: 0.8rem;
                display: inline-block;
                margin: 5px 0;
            }
            .stButton>button {
                border-radius: 25px;
                transition: all 0.3s;
            }
            .stTextInput>div>div>input, .stSelectbox>div>div>select {
                border-radius: 10px;
            }
            .stProgress > div > div > div > div {
                background-color: #4ECDC4;
            }
            .footer-text {
                color: #666 !important;
            }
        </style>
        """

st.markdown(get_css(), unsafe_allow_html=True)

# Data persistence
DATA_FILE = "habits_data.json"

def load_data():
    """Load habits data from file"""
    if os.path.exists(DATA_FILE):
        with open(DATA_FILE, 'r') as f:
            data = json.load(f)
            # Migrate old data to include categories if missing
            for habit in data.get("habits", []):
                if "category" not in habit:
                    habit["category"] = "Other"
                if "reminder_time" not in habit:
                    habit["reminder_time"] = "09:00"
            return data
    return {"habits": [], "completions": {}, "reminders_sent": {}}

def save_data(data):
    """Save habits data to file"""
    with open(DATA_FILE, 'w') as f:
        json.dump(data, f, indent=2)

def get_streak(habit_id, completions, habit_data):
    """Calculate current streak for a habit"""
    if habit_id not in completions or not completions[habit_id]:
        return 0
    
    dates = sorted([datetime.strptime(d, "%Y-%m-%d").date() for d in completions[habit_id]], reverse=True)
    
    if not dates:
        return 0
    
    today = datetime.now().date()
    streak = 0
    check_date = today
    
    if dates[0] == today or dates[0] == today - timedelta(days=1):
        for i, date in enumerate(dates):
            expected_date = today - timedelta(days=i)
            if date == expected_date or (i == 0 and date == today - timedelta(days=1)):
                streak += 1
            else:
                break
    else:
        streak = 0
    
    return streak

def get_weekly_completion(habit_id, completions):
    """Get completion rate for last 7 days"""
    if habit_id not in completions:
        return 0
    
    today = datetime.now().date()
    week_ago = today - timedelta(days=7)
    
    week_completions = sum(
        1 for d in completions[habit_id]
        if week_ago <= datetime.strptime(d, "%Y-%m-%d").date() <= today
    )
    
    return (week_completions / 7) * 100

def get_monthly_stats(habit_id, completions):
    """Get monthly completion stats"""
    if habit_id not in completions:
        return 0, 0
    
    today = datetime.now().date()
    month_ago = today - timedelta(days=30)
    
    month_completions = [
        datetime.strptime(d, "%Y-%m-%d").date()
        for d in completions[habit_id]
        if month_ago <= datetime.strptime(d, "%Y-%m-%d").date() <= today
    ]
    
    return len(month_completions), (len(month_completions) / 30) * 100

def export_to_csv():
    """Export all habit data to CSV"""
    data = st.session_state.data
    rows = []
    
    for habit in data["habits"]:
        habit_id = habit["id"]
        completions = data["completions"].get(habit_id, [])
        
        for date in completions:
            rows.append({
                "Habit Name": habit["name"],
                "Category": habit["category"],
                "Emoji": habit["emoji"],
                "Completion Date": date,
                "Created At": habit["created_at"]
            })
    
    if not rows:
        # Create summary even if no completions
        for habit in data["habits"]:
            rows.append({
                "Habit Name": habit["name"],
                "Category": habit["category"],
                "Emoji": habit["emoji"],
                "Completion Date": "No completions yet",
                "Created At": habit["created_at"]
            })
    
    df = pd.DataFrame(rows)
    return df.to_csv(index=False).encode('utf-8')

def check_reminders():
    """Check and display due reminders"""
    current_time = datetime.now().strftime("%H:%M")
    today = datetime.now().strftime("%Y-%m-%d")
    
    reminders_due = []
    
    for habit in st.session_state.data["habits"]:
        reminder_time = habit.get("reminder_time", "09:00")
        habit_id = habit["id"]
        
        # Check if reminder time matches current time (within 1 minute window)
        if reminder_time == current_time:
            # Check if not already completed today and not already reminded
            completions = st.session_state.data["completions"].get(habit_id, [])
            reminders_sent = st.session_state.data.get("reminders_sent", {})
            
            if today not in completions and reminders_sent.get(habit_id) != today:
                reminders_due.append(habit)
                st.session_state.data["reminders_sent"][habit_id] = today
                save_data(st.session_state.data)
    
    return reminders_due

# Initialize session state
if 'data' not in st.session_state:
    st.session_state.data = load_data()

# Check for reminders
reminders = check_reminders()

# Sidebar
with st.sidebar:
    # Dark mode toggle at top
    st.button("🌙 Dark Mode" if not st.session_state.dark_mode else "☀️ Light Mode", 
              on_click=toggle_theme, use_container_width=True)
    
    st.markdown("---")
    
    st.header("➕ Add New Habit")
    st.markdown("*Make it obvious, make it easy*")
    
    with st.form("new_habit"):
        habit_name = st.text_input("Habit Name", placeholder="e.g., Read 1 page")
        habit_category = st.selectbox("Category", list(CATEGORIES.keys()), 
                                      format_func=lambda x: f"{CATEGORIES[x]} {x}")
        habit_color = st.color_picker("Choose Color", "#4ECDC4")
        habit_emoji = st.selectbox("Icon", ["📚", "🏃", "💧", "🧘", "💪", "🎸", "💻", "🌱", "✍️", "🎯", "🥗", "💤"])
        reminder_time = st.time_input("⏰ Daily Reminder", value=datetime.strptime("09:00", "%H:%M").time())
        
        submitted = st.form_submit_button("Create Habit", use_container_width=True)
        
        if submitted and habit_name:
            habit_id = f"habit_{datetime.now().strftime('%Y%m%d%H%M%S')}"
            new_habit = {
                "id": habit_id,
                "name": habit_name,
                "category": habit_category,
                "color": habit_color,
                "emoji": habit_emoji,
                "reminder_time": reminder_time.strftime("%H:%M"),
                "created_at": datetime.now().strftime("%Y-%m-%d")
            }
            st.session_state.data["habits"].append(new_habit)
            st.session_state.data["completions"][habit_id] = []
            save_data(st.session_state.data)
            st.success(f"Created: {habit_emoji} {habit_name}")
            st.rerun()
    
    st.markdown("---")
    
    # Export Section
    st.header("📥 Export Data")
    if st.session_state.data["habits"]:
        csv = export_to_csv()
        st.download_button(
            label="Download CSV",
            data=csv,
            file_name=f"habits_export_{datetime.now().strftime('%Y%m%d')}.csv",
            mime="text/csv",
            use_container_width=True
        )
        st.caption("Export all your habit data for backup or analysis")
    else:
        st.info("Add habits to enable export")
    
    st.markdown("---")
    
    # Stats Summary
    if st.session_state.data["habits"]:
        st.header("📊 Quick Stats")
        total_habits = len(st.session_state.data["habits"])
        total_completions = sum(len(v) for v in st.session_state.data["completions"].values())
        active_streaks = sum(
            1 for h in st.session_state.data["habits"]
            if get_streak(h["id"], st.session_state.data["completions"], h) > 0
        )
        
        st.metric("Total Habits", total_habits)
        st.metric("Total Completions", total_completions)
        st.metric("Active Streaks", active_streaks)

# Display reminders
if reminders:
    st.markdown("---")
    st.subheader("🔔 Reminders")
    for habit in reminders:
        st.info(f"⏰ Time to complete: {habit['emoji']} **{habit['name']}**!")

# Main Header
st.markdown('<h1 class="main-header">🔥 Atomic Habits Tracker Pro</h1>', unsafe_allow_html=True)

# Daily Quote
quotes = [
    "You do not rise to the level of your goals. You fall to the level of your systems.",
    "Every action you take is a vote for the type of person you wish to become.",
    "Success is the product of daily habits—not once-in-a-lifetime transformations.",
    "The cost of your good habit is in the dozens of moments before you do it.",
    "You should be far more concerned with your current trajectory than with your current results.",
    "Habits are the compound interest of self-improvement.",
    "You don't have to be the victim of your environment. You can also be the architect of it."
]
st.markdown(f'<div class="quote-box">"{quotes[datetime.now().day % len(quotes)]}" — James Clear</div>', unsafe_allow_html=True)

# Category Filter
if st.session_state.data["habits"]:
    all_categories = ["All"] + list(CATEGORIES.keys())
    selected_category = st.selectbox("Filter by Category", all_categories, horizontal=True)
    
    filtered_habits = st.session_state.data["habits"]
    if selected_category != "All":
        filtered_habits = [h for h in filtered_habits if h["category"] == selected_category]

# Main Content
if not st.session_state.data["habits"]:
    st.info("👈 Start by adding your first habit in the sidebar! Remember: start small (2-minute rule).")
else:
    today = datetime.now().strftime("%Y-%m-%d")
    today_display = datetime.now().strftime("%A, %B %d, %Y")
    
    st.subheader(f"📅 {today_display}")
    
    # Group by category if showing all
    if selected_category == "All":
        for category in CATEGORIES.keys():
            category_habits = [h for h in filtered_habits if h["category"] == category]
            if category_habits:
                st.markdown(f"### {CATEGORIES[category]} {category}")
                cols = st.columns(2)
                for idx, habit in enumerate(category_habits):
                    with cols[idx % 2]:
                        render_habit_card(habit, today)
    else:
        cols = st.columns(2)
        for idx, habit in enumerate(filtered_habits):
            with cols[idx % 2]:
                render_habit_card(habit, today)

def render_habit_card(habit, today):
    """Render individual habit card"""
    habit_id = habit["id"]
    completions = st.session_state.data["completions"].get(habit_id, [])
    is_completed_today = today in completions
    streak = get_streak(habit_id, st.session_state.data["completions"], habit)
    weekly_rate = get_weekly_completion(habit_id, st.session_state.data["completions"])
    monthly_count, monthly_rate = get_monthly_stats(habit_id, st.session_state.data["completions"])
    
    with st.container():
        st.markdown(f"""
        <div class="habit-card" style="border-left: 5px solid {habit['color']};">
        """, unsafe_allow_html=True)
        
        # Header with category badge
        col1, col2 = st.columns([3, 1])
        
        with col1:
            st.markdown(f"### {habit['emoji']} {habit['name']}")
            st.markdown(f'<span class="category-badge">{CATEGORIES[habit["category"]]} {habit["category"]}</span>', unsafe_allow_html=True)
            st.caption(f"⏰ Reminder at {habit.get('reminder_time', '09:00')}")
        
        with col2:
            if is_completed_today:
                st.success("✅ Done!")
                if st.button("Undo", key=f"undo_{habit_id}", help="Mark as not done"):
                    st.session_state.data["completions"][habit_id].remove(today)
                    save_data(st.session_state.data)
                    st.rerun()
            else:
                if st.button("Complete ✓", key=f"complete_{habit_id}", type="primary", use_container_width=True):
                    st.session_state.data["completions"][habit_id].append(today)
                    save_data(st.session_state.data)
                    st.balloons()
                    st.rerun()
        
        # Stats row
        stat_col1, stat_col2, stat_col3 = st.columns(3)
        
        with stat_col1:
            if streak > 0:
                st.markdown(f'<span class="streak-badge">🔥 {streak} day streak</span>', unsafe_allow_html=True)
            else:
                st.markdown("🔹 Start today!")
        
        with stat_col2:
            st.progress(weekly_rate / 100, text=f"Week: {weekly_rate:.0f}%")
        
        with stat_col3:
            st.caption(f"30d: {monthly_count} days ({monthly_rate:.0f}%)")
        
        # Edit/Delete
        with st.expander("⚙️ Settings"):
            edit_col1, edit_col2 = st.columns(2)
            
            with edit_col1:
                new_time = st.time_input("Update reminder", 
                                        value=datetime.strptime(habit.get("reminder_time", "09:00"), "%H:%M").time(),
                                        key=f"time_{habit_id}")
                if new_time.strftime("%H:%M") != habit.get("reminder_time", "09:00"):
                    habit["reminder_time"] = new_time.strftime("%H:%M")
                    save_data(st.session_state.data)
                    st.success("Updated!")
            
            with edit_col2:
                if st.button("🗑️ Delete Habit", key=f"del_{habit_id}", type="secondary"):
                    st.session_state.data["habits"].remove(habit)
                    del st.session_state.data["completions"][habit_id]
                    save_data(st.session_state.data)
                    st.rerun()
        
        st.markdown("</div>", unsafe_allow_html=True)

# Analytics Section
st.markdown("---")
st.subheader("📊 Analytics Dashboard")

if st.session_state.data["habits"]:
    # Category breakdown
    cat_data = {}
    for habit in st.session_state.data["habits"]:
        cat = habit["category"]
        completions = len(st.session_state.data["completions"].get(habit["id"], []))
        cat_data[cat] = cat_data.get(cat, 0) + completions
    
    fig_pie = go.Figure(data=[go.Pie(
        labels=[f"{CATEGORIES[k]} {k}" for k in cat_data.keys()],
        values=list(cat_data.values()),
        hole=.4,
        marker_colors=px.colors.qualitative.Set3
    )])
    fig_pie.update_layout(title_text="Completions by Category", showlegend=True)
    
    # Timeline data
    timeline_data = defaultdict(int)
    for habit_id, dates in st.session_state.data["completions"].items():
        for date in dates:
            timeline_data[date] += 1
    
    if timeline_data:
        df_timeline = pd.DataFrame([
            {"Date": datetime.strptime(d, "%Y-%m-%d"), "Total Completions": c}
            for d, c in sorted(timeline_data.items())
        ])
        
        fig_timeline = px.line(df_timeline, x="Date", y="Total Completions", 
                              title="Your Consistency Over Time",
                              markers=True)
        fig_timeline.update_layout(xaxis_title="Date", yaxis_title="Habits Completed")
    else:
        fig_timeline = go.Figure()
        fig_timeline.update_layout(title="Start completing habits to see your timeline!")
    
    col1, col2 = st.columns(2)
    with col1:
        st.plotly_chart(fig_pie, use_container_width=True)
    with col2:
        st.plotly_chart(fig_timeline, use_container_width=True)
    
    # Detailed table
    st.subheader("📋 Habit Performance Table")
    table_data = []
    for habit in st.session_state.data["habits"]:
        habit_id = habit["id"]
        completions = st.session_state.data["completions"].get(habit_id, [])
        streak = get_streak(habit_id, st.session_state.data["completions"], habit)
        total = len(completions)
        
        table_data.append({
            "Habit": f"{habit['emoji']} {habit['name']}",
            "Category": f"{CATEGORIES[habit['category']]} {habit['category']}",
            "Current Streak": f"🔥 {streak}" if streak > 0 else "0",
            "Total Days": total,
            "Success Rate": f"{(total / max(1, (datetime.now() - datetime.strptime(habit['created_at'], '%Y-%m-%d')).days)) * 100:.1f}%"
        })
    
    df_table = pd.DataFrame(table_data)
    st.dataframe(df_table, use_container_width=True, hide_index=True)
    
    # Calendar heatmap for selected habit
    st.subheader("📅 Detailed Activity View")
    selected_habit = st.selectbox(
        "Select habit to view calendar",
        options=[h["name"] for h in st.session_state.data["habits"]],
        format_func=lambda x: next(f"{h['emoji']} {h['name']} ({h['category']})" for h in st.session_state.data["habits"] if h["name"] == x)
    )
    
    selected_habit_data = next(h for h in st.session_state.data["habits"] if h["name"] == selected_habit)
    completions = st.session_state.data["completions"].get(selected_habit_data["id"], [])
    
    if completions:
        # Create comprehensive calendar view
        dates_data = []
        for i in range(90):  # Last 90 days
            date = datetime.now().date() - timedelta(days=i)
            date_str = date.strftime("%Y-%m-%d")
            completed = date_str in completions
            dates_data.append({
                "Date": date,
                "Completed": "Yes" if completed else "No",
                "Week": date.isocalendar()[1],
                "Day": date.strftime("%a"),
                "Month": date.strftime("%b")
            })
        
        cal_df = pd.DataFrame(dates_data)
        
        # Heatmap
        fig_cal = px.density_heatmap(
            cal_df, 
            x="Week", 
            y="Day",
            z=[1 if c == "Yes" else 0 for c in cal_df["Completed"]],
            title=f"Last 90 Days Activity - {selected_habit}",
            color_continuous_scale=[(0, "#e0e0e0"), (1, selected_habit_data["color"])],
            labels={"color": "Completed"}
        )
        st.plotly_chart(fig_cal, use_container_width=True)
    else:
        st.info("No data yet. Complete your habit today to start tracking!")

# Footer
st.markdown("---")
st.markdown(f"""
<div style="text-align: center; padding: 20px;">
    <p class="footer-text">Built with ❤️ using Streamlit | Deployed on Hugging Face Spaces</p>
    <p class="footer-text" style="font-size: 0.8rem;">"Small changes, remarkable results" — James Clear</p>
    <p class="footer-text" style="font-size: 0.7rem;">Current theme: {"Dark" if st.session_state.dark_mode else "Light"} Mode</p>
</div>
""", unsafe_allow_html=True)