import gradio as gr
import datetime
from typing import Dict, List, Any, Union, Optional
import random
import os
import json
import time
import matplotlib.pyplot as plt
import numpy as np
from collections import defaultdict

# Import utilities
from utils.storage import load_data, save_data, safe_get
from utils.state import generate_id, get_timestamp, record_activity
from utils.ai_models import generate_motivation_quote
from utils.config import FILE_PATHS
from utils.logging import setup_logger
from utils.error_handling import handle_exceptions
from utils.ui_components import create_card, create_stat_card, create_progress_ring

# Initialize logger
logger = setup_logger(__name__)

# Wellness insights and tips
WELLNESS_TIPS = {
    "stress": [
        "Take 5 deep breaths to activate your parasympathetic nervous system",
        "Try the 4-7-8 breathing technique: inhale for 4, hold for 7, exhale for 8",
        "Step outside for a few minutes to get fresh air and natural light",
        "Practice progressive muscle relaxation starting from your toes",
        "Listen to calming music or nature sounds"
    ],
    "focus": [
        "Use the Pomodoro Technique: 25 minutes focused work, 5 minute break",
        "Eliminate distractions by putting your phone in another room",
        "Try the 2-minute rule: if it takes less than 2 minutes, do it now",
        "Break large tasks into smaller, manageable chunks",
        "Use background noise or instrumental music to enhance concentration"
    ],
    "energy": [
        "Stay hydrated - drink water regularly throughout the day",
        "Take short walks every hour to boost circulation",
        "Eat protein-rich snacks to maintain stable blood sugar",
        "Get 7-9 hours of quality sleep each night",
        "Expose yourself to natural light, especially in the morning"
    ],
    "motivation": [
        "Celebrate small wins to build momentum",
        "Visualize your goals and the benefits of achieving them",
        "Connect with your 'why' - remember your deeper purpose",
        "Surround yourself with positive, supportive people",
        "Review your progress regularly to see how far you've come"
    ]
}

# Meditation types and durations
MEDITATION_TYPES = {
    "mindfulness": "Focus on your breath and present moment awareness",
    "body_scan": "Progressive relaxation through body awareness",
    "loving_kindness": "Cultivate compassion for yourself and others",
    "visualization": "Guided imagery for relaxation and goal achievement",
    "breathing": "Focused breathing exercises for calm and clarity"
}

@handle_exceptions
def initialize_wellness_data(state: Dict[str, Any]) -> None:
    """Initialize wellness-related data structures"""
    if "focus_sessions" not in state:
        state["focus_sessions"] = []
    if "mood_entries" not in state:
        state["mood_entries"] = []
    if "wellness_goals" not in state:
        state["wellness_goals"] = []
    if "meditation_sessions" not in state:
        state["meditation_sessions"] = []
    if "break_reminders" not in state:
        state["break_reminders"] = {"enabled": True, "interval": 60}
    if "stress_levels" not in state:
        state["stress_levels"] = []
    if "wellness_streaks" not in state:
        state["wellness_streaks"] = {}
    if "wellness_insights" not in state:
        state["wellness_insights"] = []

@handle_exceptions
def create_pomodoro_timer() -> tuple:
    """Create advanced Pomodoro timer with customization"""
    with gr.Group() as pomodoro_group:
        gr.Markdown("### 🍅 Pomodoro Timer")
        
        # Timer display
        timer_display = gr.HTML("""
        <div class="timer-display">
            <div class="timer-circle">
                <div class="timer-time">25:00</div>
                <div class="timer-label">Focus Time</div>
            </div>
        </div>
        """)
        
        # Timer settings
        with gr.Row():
            work_duration = gr.Slider(
                minimum=15, maximum=60, value=25, step=5,
                label="Work Duration (minutes)"
            )
            break_duration = gr.Slider(
                minimum=5, maximum=30, value=5, step=5,
                label="Break Duration (minutes)"
            )
            long_break_duration = gr.Slider(
                minimum=15, maximum=45, value=15, step=5,
                label="Long Break (minutes)"
            )
        
        # Session settings
        with gr.Row():
            sessions_until_long_break = gr.Slider(
                minimum=2, maximum=8, value=4, step=1,
                label="Sessions until long break"
            )
            auto_start_breaks = gr.Checkbox(
                label="Auto-start breaks", value=False
            )
            play_sounds = gr.Checkbox(
                label="Play notification sounds", value=True
            )
        
        # Current session info
        session_info = gr.Markdown("**Session 1 of 4** | Next: 5 min break")
        
        # Timer controls
        with gr.Row():
            start_timer_btn = gr.Button("▶️ Start", variant="primary")
            pause_timer_btn = gr.Button("⏸️ Pause", variant="secondary")
            reset_timer_btn = gr.Button("🔄 Reset", variant="secondary")
            skip_timer_btn = gr.Button("⏭️ Skip", variant="secondary")
        
        # Session description
        session_description = gr.Textbox(
            label="What are you working on?",
            placeholder="Describe your focus session..."
        )
        
        # Background sounds
        background_sound = gr.Dropdown(
            choices=["None", "White Noise", "Rain", "Forest", "Cafe", "Ocean", "Pink Noise"],
            value="None",
            label="Background Sound"
        )
    
    return (pomodoro_group, timer_display, work_duration, break_duration, long_break_duration,
            sessions_until_long_break, auto_start_breaks, play_sounds, session_info,
            start_timer_btn, pause_timer_btn, reset_timer_btn, skip_timer_btn,
            session_description, background_sound)

@handle_exceptions
def create_focus_sessions_tracker() -> tuple:
    """Create focus sessions tracking and analytics"""
    with gr.Group() as focus_tracker:
        gr.Markdown("### 📊 Focus Sessions")
        
        # Today's focus stats
        today_stats = gr.HTML("""
        <div class="focus-stats">
            <div class="stat-item">
                <span class="stat-number">0</span>
                <span class="stat-label">Sessions Today</span>
            </div>
            <div class="stat-item">
                <span class="stat-number">0</span>
                <span class="stat-label">Minutes Focused</span>
            </div>
            <div class="stat-item">
                <span class="stat-number">0</span>
                <span class="stat-label">Current Streak</span>
            </div>
        </div>
        """)
        
        # Recent sessions
        sessions_table = gr.Dataframe(
            headers=["Date", "Duration", "Type", "Description", "Completed"],
            datatype=["str", "str", "str", "str", "str"],
            label="Recent Focus Sessions"
        )
        
        # Focus analytics
        focus_chart = gr.Plot(label="Focus Time Over Time")
        
        # Deep work goals
        gr.Markdown("#### Deep Work Goals")
        with gr.Row():
            daily_goal = gr.Number(
                label="Daily Focus Goal (minutes)",
                value=120, minimum=30, maximum=480
            )
            weekly_goal = gr.Number(
                label="Weekly Focus Goal (hours)",
                value=20, minimum=5, maximum=60
            )
        
        set_goals_btn = gr.Button("Set Goals")
        
        # Progress towards goals
        goals_progress = gr.HTML("""
        <div class="goals-progress">
            <div class="progress-item">
                <span class="progress-label">Daily Goal</span>
                <div class="progress-bar">
                    <div class="progress-fill" style="width: 0%"></div>
                </div>
                <span class="progress-text">0 / 120 minutes</span>
            </div>
            <div class="progress-item">
                <span class="progress-label">Weekly Goal</span>
                <div class="progress-bar">
                    <div class="progress-fill" style="width: 0%"></div>
                </div>
                <span class="progress-text">0 / 20 hours</span>
            </div>
        </div>
        """)
    
    return (focus_tracker, today_stats, sessions_table, focus_chart,
            daily_goal, weekly_goal, set_goals_btn, goals_progress)

@handle_exceptions
def create_mood_tracker() -> tuple:
    """Create comprehensive mood tracking system"""
    with gr.Group() as mood_tracker:
        gr.Markdown("### 😊 Mood Tracker")
        
        # Quick mood check-in
        with gr.Row():
            mood_rating = gr.Slider(
                minimum=1, maximum=10, value=5, step=1,
                label="How are you feeling? (1=Very Low, 10=Excellent)"
            )
            energy_level = gr.Slider(
                minimum=1, maximum=10, value=5, step=1,
                label="Energy Level (1=Exhausted, 10=Energized)"
            )
        
        # Mood factors
        mood_factors = gr.CheckboxGroup(
            choices=[
                "😴 Sleep Quality", "🍎 Nutrition", "🏃 Exercise", "👥 Social",
                "💼 Work Stress", "🌤️ Weather", "🎯 Productivity", "💰 Financial",
                "❤️ Relationships", "🎨 Creativity", "📚 Learning", "🧘 Mindfulness"
            ],
            label="What's affecting your mood today?"
        )
        
        # Mood notes
        mood_notes = gr.Textbox(
            label="Notes (optional)",
            placeholder="Any additional thoughts about your mood...",
            lines=3
        )
        
        log_mood_btn = gr.Button("Log Mood Entry", variant="primary")
        
        # Mood visualization
        mood_chart = gr.Plot(label="Mood Trends")
        
        # Mood insights
        mood_insights = gr.Markdown("*Track your mood regularly to see patterns and insights*")
        
        # Mood history
        mood_history = gr.Dataframe(
            headers=["Date", "Mood", "Energy", "Key Factors", "Notes"],
            datatype=["str", "number", "number", "str", "str"],
            label="Mood History"
        )
    
    return (mood_tracker, mood_rating, energy_level, mood_factors, mood_notes,
            log_mood_btn, mood_chart, mood_insights, mood_history)

@handle_exceptions
def create_meditation_timer() -> tuple:
    """Create meditation and mindfulness timer"""
    with gr.Group() as meditation_group:
        gr.Markdown("### 🧘 Meditation Timer")
        
        # Meditation type selection
        meditation_type = gr.Dropdown(
            choices=list(MEDITATION_TYPES.keys()),
            value="mindfulness",
            label="Meditation Type"
        )
        
        meditation_description = gr.Markdown(
            MEDITATION_TYPES["mindfulness"]
        )
        
        # Duration selection
        meditation_duration = gr.Dropdown(
            choices=["5 minutes", "10 minutes", "15 minutes", "20 minutes", "30 minutes", "Custom"],
            value="10 minutes",
            label="Duration"
        )
        
        custom_duration = gr.Number(
            label="Custom Duration (minutes)",
            value=10, minimum=1, maximum=120,
            visible=False
        )
        
        # Guided or silent
        guided_meditation = gr.Checkbox(
            label="Guided meditation (with prompts)",
            value=True
        )
        
        # Background sounds for meditation
        meditation_sound = gr.Dropdown(
            choices=["Silence", "Tibetan Bowls", "Nature Sounds", "Om Chanting", "Soft Music"],
            value="Tibetan Bowls",
            label="Background Sound"
        )
        
        # Meditation timer display
        meditation_timer_display = gr.HTML("""
        <div class="meditation-timer">
            <div class="meditation-circle">
                <div class="meditation-time">10:00</div>
                <div class="meditation-breath">Breathe...</div>
            </div>
        </div>
        """)
        
        # Meditation controls
        with gr.Row():
            start_meditation_btn = gr.Button("🧘 Start Meditation", variant="primary")
            pause_meditation_btn = gr.Button("⏸️ Pause", variant="secondary")
            end_meditation_btn = gr.Button("🛑 End Session", variant="secondary")
        
        # Meditation log
        gr.Markdown("#### Meditation History")
        meditation_log = gr.Dataframe(
            headers=["Date", "Type", "Duration", "Completed", "Notes"],
            datatype=["str", "str", "str", "str", "str"],
            label="Recent Sessions"
        )
        
        # Meditation streak
        meditation_streak = gr.HTML("""
        <div class="streak-display">
            <span class="streak-number">0</span>
            <span class="streak-label">Day Meditation Streak</span>
        </div>
        """)
    
    return (meditation_group, meditation_type, meditation_description, meditation_duration,
            custom_duration, guided_meditation, meditation_sound, meditation_timer_display,
            start_meditation_btn, pause_meditation_btn, end_meditation_btn,
            meditation_log, meditation_streak)

@handle_exceptions
def create_wellness_insights() -> tuple:
    """Create AI-powered wellness insights and recommendations"""
    with gr.Group() as wellness_insights:
        gr.Markdown("### 🤖 Wellness Insights")
        
        # Current wellness score
        wellness_score = gr.HTML("""
        <div class="wellness-score">
            <div class="score-circle">
                <span class="score-number">--</span>
                <span class="score-label">Wellness Score</span>
            </div>
        </div>
        """)
        
        # AI recommendations
        ai_recommendations = gr.Markdown("*Complete some wellness activities to get personalized recommendations*")
        
        # Wellness tips
        gr.Markdown("#### Daily Wellness Tips")
        wellness_tip = gr.Markdown("💡 " + random.choice(WELLNESS_TIPS["focus"]))
        new_tip_btn = gr.Button("Get New Tip")
        
        # Stress monitoring
        gr.Markdown("#### Stress Monitoring")
        with gr.Row():
            stress_level = gr.Slider(
                minimum=1, maximum=10, value=5, step=1,
                label="Current Stress Level (1=Very Calm, 10=Very Stressed)"
            )
            stress_factors = gr.CheckboxGroup(
                choices=[
                    "⏰ Time Pressure", "💼 Workload", "👥 Social Situations",
                    "💰 Financial Concerns", "🏠 Personal Life", "🏥 Health Issues",
                    "🌍 World Events", "🎯 Performance Pressure"
                ],
                label="Stress Factors"
            )
        
        log_stress_btn = gr.Button("Log Stress Level")
        
        # Stress relief suggestions
        stress_relief = gr.Markdown("*Log your stress level to get personalized relief suggestions*")
        
        # Wellness patterns
        wellness_patterns = gr.Plot(label="Wellness Patterns Over Time")
    
    return (wellness_insights, wellness_score, ai_recommendations, wellness_tip, new_tip_btn,
            stress_level, stress_factors, log_stress_btn, stress_relief, wellness_patterns)

@handle_exceptions
def create_break_reminders() -> tuple:
    """Create break reminder system"""
    with gr.Group() as break_reminders:
        gr.Markdown("### ⏰ Break Reminders")
        
        # Break reminder settings
        with gr.Row():
            break_interval = gr.Slider(
                minimum=15, maximum=120, value=60, step=15,
                label="Break Reminder Interval (minutes)"
            )
            break_duration_reminder = gr.Slider(
                minimum=2, maximum=15, value=5, step=1,
                label="Suggested Break Duration (minutes)"
            )
        
        # Break types
        break_types = gr.CheckboxGroup(
            choices=[
                "🚶 Movement Break", "👀 Eye Rest", "🧘 Mindfulness",
                "💧 Hydration", "🌬️ Breathing Exercise", "🌿 Nature View"
            ],
            value=["🚶 Movement Break", "👀 Eye Rest"],
            label="Break Activities to Suggest"
        )
        
        # Enable/disable reminders
        enable_reminders = gr.Checkbox(
            label="Enable break reminders",
            value=True
        )
        
        save_reminder_settings_btn = gr.Button("Save Settings")
        
        # Break activity suggestions
        gr.Markdown("#### Quick Break Activities")
        break_suggestions = gr.HTML("""
        <div class="break-activities">
            <div class="activity-card">
                <h4>🚶 Movement Break</h4>
                <p>Stand up, stretch, or take a short walk</p>
            </div>
            <div class="activity-card">
                <h4>👀 Eye Rest</h4>
                <p>Look away from screen, focus on distant objects</p>
            </div>
            <div class="activity-card">
                <h4>🧘 Mindfulness</h4>
                <p>Take 5 deep breaths and center yourself</p>
            </div>
        </div>
        """)
        
        # Break log
        break_log = gr.Dataframe(
            headers=["Time", "Type", "Duration", "Effectiveness"],
            datatype=["str", "str", "str", "str"],
            label="Break History"
        )
    
    return (break_reminders, break_interval, break_duration_reminder, break_types,
            enable_reminders, save_reminder_settings_btn, break_suggestions, break_log)

@handle_exceptions
def create_focus_page(state: Dict[str, Any]) -> None:
    """
    Create the comprehensive Focus & Wellness page
    
    Args:
        state: Application state
    """
    logger.info("Creating comprehensive focus and wellness page")
    
    # Initialize data
    initialize_wellness_data(state)
    
    # Create the focus and wellness page layout
    with gr.Column(elem_id="focus-wellness-page"):
        gr.Markdown("# 🧘 Focus & Wellness Tools")
        gr.Markdown("*Enhance your focus, track your well-being, and maintain healthy habits*")
        
        # Wellness overview
        with gr.Row():
            wellness_overview = gr.HTML("""
            <div class="wellness-overview">
                <div class="wellness-stat">
                    <h3>🎯 Today's Focus</h3>
                    <p class="stat-number">0 min</p>
                </div>
                <div class="wellness-stat">
                    <h3>😊 Mood Score</h3>
                    <p class="stat-number">--</p>
                </div>
                <div class="wellness-stat">
                    <h3>🧘 Meditation</h3>
                    <p class="stat-number">0 min</p>
                </div>
                <div class="wellness-stat">
                    <h3>🔥 Streak</h3>
                    <p class="stat-number">0 days</p>
                </div>
            </div>
            """)
        
        # Main tabs
        with gr.Tabs(elem_id="focus-wellness-tabs"):
            # Pomodoro Timer
            with gr.TabItem("🍅 Pomodoro Timer", elem_id="pomodoro-tab"):
                pomodoro_group, timer_display, work_duration, break_duration, long_break_duration, sessions_until_long_break, auto_start_breaks, play_sounds, session_info, start_timer_btn, pause_timer_btn, reset_timer_btn, skip_timer_btn, session_description, background_sound = create_pomodoro_timer()
            
            # Focus Sessions
            with gr.TabItem("📊 Focus Sessions", elem_id="focus-sessions-tab"):
                focus_tracker, today_stats, sessions_table, focus_chart, daily_goal, weekly_goal, set_goals_btn, goals_progress = create_focus_sessions_tracker()
            
            # Mood Tracking
            with gr.TabItem("😊 Mood Tracker", elem_id="mood-tab"):
                mood_tracker, mood_rating, energy_level, mood_factors, mood_notes, log_mood_btn, mood_chart, mood_insights, mood_history = create_mood_tracker()
            
            # Meditation Timer
            with gr.TabItem("🧘 Meditation", elem_id="meditation-tab"):
                meditation_group, meditation_type, meditation_description, meditation_duration, custom_duration, guided_meditation, meditation_sound, meditation_timer_display, start_meditation_btn, pause_meditation_btn, end_meditation_btn, meditation_log, meditation_streak = create_meditation_timer()
            
            # Break Reminders
            with gr.TabItem("⏰ Break Reminders", elem_id="breaks-tab"):
                break_reminders, break_interval, break_duration_reminder, break_types, enable_reminders, save_reminder_settings_btn, break_suggestions, break_log = create_break_reminders()
            
            # Wellness Insights
            with gr.TabItem("🤖 Wellness Insights", elem_id="insights-tab"):
                wellness_insights, wellness_score, ai_recommendations, wellness_tip, new_tip_btn, stress_level, stress_factors, log_stress_btn, stress_relief, wellness_patterns = create_wellness_insights()
        
        # Event handlers and functionality will be added here
        # This includes all the interactive functions for:
        # - Pomodoro timer functionality
        # - Focus session tracking
        # - Mood logging and analysis
        # - Meditation timer and guidance
        # - Break reminder system
        # - Wellness insights and recommendations
                           value="Exercise" ,
                            label="Goal Type"
                        )
                        
                        goal_target = gr.Number(
                            value=30,
                            label="Target (minutes/amount)"
                        )
                        
                        goal_frequency = gr.Dropdown(
                            choices=[
                                "Daily",
                                "Weekly",
                                "Monthly"
                            ],
                            value="Daily",
                            label="Frequency"
                        )
                        
                        add_goal_btn = gr.Button("Add Goal")
                    
                    # Goals progress
                    with gr.Column(scale=2):
                        goals_progress = gr.Dataframe(
                            headers=["Goal", "Target", "Progress", "Streak"],
                            datatype=["str", "str", "str", "str"],
                            label="Current Goals"
                        )
                
                # Wellness insights
                wellness_insights = gr.Markdown(
                    "*Add wellness goals to see insights*"
                )
        
        # Function to update timer display
        @handle_exceptions
        def update_timer(minutes_left: int, seconds_left: int) -> str:
            """Update the timer display"""
            return f"{minutes_left:02d}:{seconds_left:02d}"
        
        # Function to toggle custom length input
        @handle_exceptions
        def toggle_custom_length(choice: str) -> bool:
            """Show/hide custom length input based on dropdown selection"""
            return choice == "Custom"
        
        # Function to start focus session
        @handle_exceptions
        def start_session(
            length: str,
            custom: int,
            description: str,
            sound: str
        ) -> None:
            """Start a new focus session"""
            logger.info(f"Starting focus session: {description}")
            
            # Get session length in minutes
            if length == "Custom":
                minutes = custom
            else:
                minutes = int(length.split()[0])
            
            # Create new session
            session = {
                "id": generate_id(),
                "start_time": get_timestamp(),
                "duration": minutes,
                "description": description,
                "sound": sound,
                "completed": False
            }
            
            # Add to state
            state["focus_sessions"].append(session)
            save_data(FILE_PATHS["focus_sessions"], state["focus_sessions"])
            
            # Record activity
            record_activity(state, "Started Focus Session", {
                "duration": minutes,
                "description": description
            })
        
        # Function to complete focus session
        @handle_exceptions
        def complete_session(session_id: str) -> None:
            """Mark a focus session as completed"""
            logger.info(f"Completing focus session: {session_id}")
            
            # Find and update session
            for session in state["focus_sessions"]:
                if session["id"] == session_id:
                    session["completed"] = True
                    session["end_time"] = get_timestamp()
                    break
            
            # Save updated sessions
            save_data(FILE_PATHS["focus_sessions"], state["focus_sessions"])
            
            # Record activity
            record_activity(state, "Completed Focus Session", {
                "session_id": session_id
            })
        
        # Function to log mood
        @handle_exceptions
        def log_mood(rating: int, notes: str) -> None:
            """Log a new mood entry"""
            logger.info(f"Logging mood: {rating}")
            
            # Create mood entry
            entry = {
                "id": generate_id(),
                "timestamp": get_timestamp(),
                "rating": rating,
                "notes": notes
            }
            
            # Add to state
            state["mood_entries"].append(entry)
            save_data(FILE_PATHS["mood_entries"], state["mood_entries"])
            
            # Record activity
            record_activity(state, "Logged Mood", {
                "rating": rating
            })
        
        # Function to add wellness goal
        @handle_exceptions
        def add_wellness_goal(
            goal_type: str,
            target: int,
            frequency: str
        ) -> None:
            """Add a new wellness goal"""
            logger.info(f"Adding wellness goal: {goal_type}")
            
            # Create goal
            goal = {
                "id": generate_id(),
                "type": goal_type,
                "target": target,
                "frequency": frequency,
                "created_at": get_timestamp(),
                "progress": 0,
                "streak": 0
            }
            
            # Add to state
            state["wellness_goals"].append(goal)
            save_data(FILE_PATHS["wellness_goals"], state["wellness_goals"])
            
            # Record activity
            record_activity(state, "Added Wellness Goal", {
                "type": goal_type,
                "target": target,
                "frequency": frequency
            })
        
        # Connect UI components to functions
        session_length.change(toggle_custom_length, inputs=[session_length], outputs=[custom_length])
        
        start_btn.click(
            start_session,
            inputs=[session_length, custom_length, session_description, background_sound],
            outputs=[]
        )
        
        log_mood_btn.click(
            log_mood,
            inputs=[mood_rating, mood_notes],
            outputs=[]
        )
        
        add_goal_btn.click(
            add_wellness_goal,
            inputs=[goal_type, goal_target, goal_frequency],
            outputs=[]
        )
        
        # Record page visit in activity
        record_activity(state, "Viewed Focus & Wellness Page")