import streamlit as st
import json
import os
from datetime import datetime, timedelta
import pandas as pd
import plotly.express as px
from plotly.subplots import make_subplots
import plotly.graph_objects as go

# Page config
st.set_page_config(
    page_title="Atomic Habits Tracker", 
    page_icon="🎯",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS for Atomic Habits styling
st.markdown("""
<style>
    .identity-card {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        padding: 20px;
        border-radius: 15px;
        color: white;
        margin: 10px 0;
        box-shadow: 0 4px 6px rgba(0,0,0,0.1);
    }
    .streak-badge {
        background-color: #ff6b6b;
        color: white;
        padding: 5px 15px;
        border-radius: 20px;
        font-weight: bold;
        display: inline-block;
        margin: 5px;
    }
    .two-minute-win {
        background-color: #4ecdc4;
        color: white;
        padding: 3px 10px;
        border-radius: 15px;
        font-size: 0.8em;
        margin-left: 10px;
    }
    .missed-alert {
        background-color: #ffe66d;
        color: #333;
        padding: 10px;
        border-radius: 8px;
        border-left: 5px solid #ff6b6b;
        margin: 10px 0;
    }
    .chain-intact {
        color: #4ecdc4;
        font-size: 24px;
    }
    .chain-broken {
        color: #ff6b6b;
        font-size: 24px;
    }
</style>
""", unsafe_allow_html=True)

# Data persistence
HABITS_FILE = "habits_data.json"
LOGS_FILE = "habits_log.json"

def load_data():
    """Load habits and logs from JSON"""
    habits = {}
    logs = {}
    
    if os.path.exists(HABITS_FILE):
        with open(HABITS_FILE, 'r') as f:
            habits = json.load(f)
    
    if os.path.exists(LOGS_FILE):
        with open(LOGS_FILE, 'r') as f:
            logs = json.load(f)
    
    return habits, logs

def save_data(habits, logs):
    """Save to JSON files"""
    with open(HABITS_FILE, 'w') as f:
        json.dump(habits, f, indent=2)
    with open(LOGS_FILE, 'w') as f:
        json.dump(logs, f, indent=2)

# Initialize data
if 'habits' not in st.session_state:
    st.session_state.habits, st.session_state.logs = load_data()

habits = st.session_state.habits
logs = st.session_state.logs

# Helper functions
def get_today():
    return datetime.now().strftime("%Y-%m-%d")

def get_yesterday():
    return (datetime.now() - timedelta(days=1)).strftime("%Y-%m-%d")

def calculate_streak(habit_id):
    """Calculate current streak for a habit"""
    if habit_id not in logs:
        return 0
    
    dates = sorted(logs[habit_id].keys(), reverse=True)
    if not dates:
        return 0
    
    today = get_today()
    yesterday = get_yesterday()
    
    # If nothing logged today, check if we can continue from yesterday
    if today not in logs[habit_id]:
        if yesterday not in logs[habit_id]:
            return 0
    
    streak = 0
    check_date = datetime.now()
    
    # If missed today, start counting from yesterday
    if today not in logs[habit_id]:
        check_date = datetime.now() - timedelta(days=1)
    
    while True:
        date_str = check_date.strftime("%Y-%m-%d")
        if date_str in logs[habit_id]:
            status = logs[habit_id][date_str].get('status')
            if status in ['completed', 'two_minute']:
                streak += 1
                check_date -= timedelta(days=1)
            else:
                break
        else:
            break
    
    return streak

def missed_yesterday(habit_id):
    """Check if habit was missed yesterday"""
    yesterday = get_yesterday()
    if habit_id not in logs or yesterday not in logs[habit_id]:
        return True
    return logs[habit_id][yesterday].get('status') == 'missed'

# Sidebar - Add New Habit (Identity Design)
with st.sidebar:
    st.header("🆕 Design New Habit")
    st.markdown("**Remember:** Focus on *identity*, not outcomes.")
    
    with st.form("new_habit"):
        habit_name = st.text_input("Habit Name", placeholder="e.g., Read 30 pages")
        identity = st.text_input("Identity Label", placeholder="e.g., Reader")
        
        st.markdown("**Habit Stack** (Implementation Intention)")
        anchor = st.text_input("After I...", placeholder="pour my morning coffee")
        behavior = st.text_input("I will...", placeholder="read one page")
        
        st.markdown("**Environment Design**")
        environment = st.text_input("Environment Cue", placeholder="Book on pillow")
        
        st.markdown("**The 2-Minute Rule**")
        two_min = st.text_input("2-Minute Version", placeholder="Open the book")
        
        submitted = st.form_submit_button("Create Habit", use_container_width=True)
        
        if submitted and habit_name and identity:
            habit_id = f"habit_{datetime.now().timestamp()}"
            habits[habit_id] = {
                "name": habit_name,
                "identity": identity,
                "anchor": anchor,
                "behavior": behavior,
                "environment": environment,
                "two_minute": two_min,
                "created": get_today()
            }
            save_data(habits, logs)
            st.success(f"Created identity: I am a {identity}")
            st.rerun()

# Main Layout
st.title("🎯 Atomic Habits Tracker")
st.markdown("*Every action you take is a vote for the type of person you wish to become.*")

# Today's Check-in Section
st.header("📅 Today's Vote")

today = get_today()
cols = st.columns(2)

if not habits:
    st.info("👈 Create your first identity-based habit in the sidebar")
else:
    for idx, (habit_id, habit) in enumerate(habits.items()):
        with cols[idx % 2]:
            # Identity Card
            streak = calculate_streak(habit_id)
            st.markdown(f"""
            <div class="identity-card">
                <h3>I am a {habit['identity']}</h3>
                <p><em>"After I {habit['anchor']}, I will {habit['behavior']}"</em></p>
                <span class="streak-badge">🔥 {streak} day streak</span>
            </div>
            """, unsafe_allow_html=True)
            
            # Never Miss Twice Alert
            if missed_yesterday(habit_id):
                st.markdown("""
                <div class="missed-alert">
                    ⚠️ <b>Never Miss Twice!</b> You missed yesterday. 
                    Do the 2-minute version today to break the fall.
                </div>
                """, unsafe_allow_html=True)
            
            # Check-in interface
            current_status = logs.get(habit_id, {}).get(today, {}).get('status', 'pending')
            
            col1, col2 = st.columns([2, 1])
            
            with col1:
                status = st.radio(
                    "Today's action:",
                    ['completed', 'two_minute', 'missed'],
                    key=f"radio_{habit_id}",
                    format_func=lambda x: {
                        'completed': f"✅ Completed ({habit['behavior']})",
                        'two_minute': f"🌱 2-Minute Win ({habit['two_minute']})",
                        'missed': "❌ Missed"
                    }[x],
                    index=['completed', 'two_minute', 'missed'].index(current_status) if current_status in ['completed', 'two_minute', 'missed'] else 0
                )
            
            with col2:
                if st.button("Log Vote", key=f"log_{habit_id}", use_container_width=True):
                    if habit_id not in logs:
                        logs[habit_id] = {}
                    
                    logs[habit_id][today] = {
                        'status': status,
                        'timestamp': datetime.now().isoformat(),
                        'note': ''
                    }
                    save_data(habits, logs)
                    st.success("Vote recorded!")
                    st.rerun()
            
            # Environment reminder
            st.caption(f"💡 Reminder: {habit['environment']}")
            st.divider()

# Dashboard Section
if habits:
    st.header("📊 Identity Dashboard")
    
    # Calculate aggregate stats
    total_votes = 0
    identity_votes = {}
    
    for habit_id, habit in habits.items():
        identity = habit['identity']
        if identity not in identity_votes:
            identity_votes[identity] = 0
        
        if habit_id in logs:
            for date, entry in logs[habit_id].items():
                if entry['status'] in ['completed', 'two_minute']:
                    total_votes += 1
                    identity_votes[identity] += 1
    
    # Metrics row
    col1, col2, col3 = st.columns(3)
    with col1:
        st.metric("Total Identity Votes", total_votes)
    with col2:
        st.metric("Active Identities", len(habits))
    with col3:
        # Completion rate today
        today_completed = sum(1 for h in habits if h in logs and today in logs[h] and logs[h][today]['status'] != 'missed')
        rate = (today_completed / len(habits) * 100) if habits else 0
        st.metric("Today's Completion", f"{rate:.0f}%")
    
    # Visual Chain Calendar
    st.subheader("🔗 Don't Break the Chain")
    
    for habit_id, habit in habits.items():
        if habit_id not in logs:
            continue
            
        # Get last 30 days of data
        dates = []
        statuses = []
        
        for i in range(29, -1, -1):
            date = (datetime.now() - timedelta(days=i)).strftime("%Y-%m-%d")
            dates.append(date)
            
            if date in logs[habit_id]:
                statuses.append(logs[habit_id][date]['status'])
            else:
                statuses.append('none')
        
        # Create color map
        color_map = {'completed': '#4ecdc4', 'two_minute': '#95e1d3', 'missed': '#ff6b6b', 'none': '#f0f0f0'}
        colors = [color_map[s] for s in statuses]
        
        # Plotly calendar heatmap
        fig = go.Figure(data=[go.Bar(
            x=dates,
            y=[1]*len(dates),
            marker_color=colors,
            hovertemplate='%{x}<br>Status: %{text}<extra></extra>',
            text=statuses
        )])
        
        fig.update_layout(
            title=f"I am a {habit['identity']}",
            height=150,
            showlegend=False,
            xaxis_showgrid=False,
            yaxis_showgrid=False,
            yaxis_visible=False,
            plot_bgcolor='rgba(0,0,0,0)'
        )
        
        st.plotly_chart(fig, use_container_width=True, key=f"chart_{habit_id}")

# Export/Backup
with st.expander("💾 Backup Data"):
    st.download_button(
        "Download Habits JSON",
        data=json.dumps(habits, indent=2),
        file_name="atomic_habits_backup.json",
        mime="application/json"
    )