import gradio as gr
import datetime
import json
import os
from transformers import pipeline

# Load an emotion detection model
# top_k=None to get all emotions and their scores
emotion_pipeline = pipeline("text-classification", model="SamLowe/roberta-base-go_emotions", top_k=None) 

# Custom motivational responses based on various emotions
def get_motivation(label):
    messages = {
        "joy": "That's wonderful! Keep embracing that positive energy! 🌟",
        "optimism": "It's great to see your hopeful outlook! Keep believing in yourself! ✨",
        "love": "Feeling the love is a beautiful thing! Cherish these moments. ❤️",
        "sadness": "I'm here for you. It's okay to feel sad, and remember that tough times don't last. 🌧️",
        "anger": "It sounds like you're feeling angry. Acknowledge your feelings, and remember to be kind to yourself. 😤",
        "fear": "It's natural to feel scared sometimes. Take a deep breath; you're stronger than you think. 🛡️",
        "surprise": "Oh, what an unexpected turn! Hope it's a good one! 😲",
        "neutral": "It's okay to have neutral days. Keep going, one step at a time. 🌱",
        "disappointment": "It's tough when things don't go as planned. Allow yourself to feel it, and remember there's always tomorrow.😔",
        "excitement": "Awesome! Your excitement is contagious! Keep that energy flowing! 🎉",
        "gratitude": "Feeling thankful is a wonderful practice. What a lovely perspective! 🙏",
        "desire": "It's good to have aspirations! What steps can you take towards your goals? 🌠",
        "confusion": "Feeling a bit lost? It's okay to not have all the answers. Take your time to figure things out. 🤔",
        "admiration": "It's inspiring to look up to someone! What qualities do you admire most? 🤩",
        "amusement": "Laughter is the best medicine! Glad to hear you're finding joy. 😂"
    }
    return messages.get(label.lower(), "Thanks for sharing. Keep expressing yourself. 💬")

# Function to get suggestions based on emotion
def get_suggestions(emotion_label):
    suggestions = {
        "sadness": [
            "Relax and listen to calming music.",
            "Take a warm bath or shower.",
            "Connect with a trusted friend or family member.",
            "Go for a gentle walk in nature.",
            "Watch a comforting movie or read a book.",
            "Allow yourself to cry if you need to, it can be a healthy release.",
            "Engage in a favorite hobby that brings you a sense of calm."
        ],
        "anger": [
            "Take a few deep breaths to calm your nervous system.",
            "Step away from the situation if possible.",
            "Engage in physical activity, like a brisk walk or some stretches, to release tension.",
            "Express your feelings in a non-confrontational way, perhaps by writing them down.",
            "Listen to music that helps you relax or channel your energy.",
            "Squeeze a stress ball or do something with your hands like crafting."
        ],
        "fear": [
            "Practice deep breathing exercises (e.g., 4-7-8 breathing).",
            "Identify 5 things you can see, 4 you can touch, 3 you can hear, 2 you can smell, and 1 you can taste (grounding technique).",
            "Talk to someone you trust about your worries.",
            "Engage in a calming activity like meditation or gentle yoga.",
            "Remind yourself that fear is a natural emotion and it will pass.",
            "Challenge anxious thoughts by asking if they are truly factual."
        ],
        "joy": [
            "Savor the moment and appreciate what made you feel this way.",
            "Share your happiness with loved ones.",
            "Engage in activities that bring you more joy, like a favorite hobby.",
            "Practice gratitude: list things you are thankful for.",
            "Help someone else – giving back can amplify joyful feelings.",
            "Capture the moment (e.g., take a photo, write it down) to revisit later."
        ],
        "optimism": [
            "Keep focusing on your goals and take small steps towards them.",
            "Share your positive outlook with others to inspire them.",
            "Continue practicing gratitude to reinforce positive thinking.",
            "Reflect on your strengths and past successes.",
            "Consider setting new positive intentions for your future."
        ],
        "disappointment": [
            "Acknowledge your feelings and allow yourself to feel disappointed.",
            "Reflect on what you can learn from the situation.",
            "Focus on what you can control moving forward.",
            "Talk to someone supportive about how you feel.",
            "Engage in a self-care activity to uplift your spirits.",
            "Remember that it's okay for things not to go as planned sometimes."
        ],
        "gratitude": [
            "Keep a gratitude journal to consistently note things you're thankful for.",
            "Express your thanks to someone in your life.",
            "Perform a random act of kindness for someone else.",
            "Mindfully appreciate the small blessings in your day.",
            "Reflect on how far you've come and what you've overcome."
        ],
        # Add more emotions and their corresponding suggestions here
    }
    
    # Default suggestions if no specific emotion is matched or for neutral
    default_suggestions = [
        "Take a moment to reflect on your day.",
        "Engage in a relaxing activity like reading or listening to music.",
        "Practice deep breathing exercises.",
        "Connect with a friend or family member.",
        "Consider going for a short walk.",
        "Hydrate and consider a healthy snack."
    ]

    chosen_suggestions = suggestions.get(emotion_label.lower(), default_suggestions)
    
    return chosen_suggestions

# Main journaling function
def analyze_journal(entry, confidence_threshold=0.1): # Added confidence_threshold
    if not entry.strip():
        # When clearing the input, also clear the suggestions output and state
        return "Please write something in your journal entry.", "I'm ready to listen when you are!", "", []

    results = emotion_pipeline(entry)
    
    detected_emotions_display = ""
    suggestions_output_list = []

    # Sort results by score in descending order
    # results is typically a list containing a list of dictionaries, so results[0]
    if not results or not results[0]: # Handle case where results might be empty
        detected_emotions_display = "No emotions detected."
        suggestions_output_list.extend(get_suggestions("neutral"))
        return f"🧠 Detected Emotions:\n{detected_emotions_display}", "Thanks for sharing. Keep expressing yourself. 💬", \
               "\n\n**💡 Here are some suggestions for you:**\n" + "\n".join([f"- {s}" for s in suggestions_output_list]), \
               suggestions_output_list

    sorted_results = sorted(results[0], key=lambda x: x['score'], reverse=True)

    for emotion_data in sorted_results:
        emotion_label = emotion_data['label']
        emotion_score = emotion_data['score']

        if emotion_score >= confidence_threshold:
            detected_emotions_display += f"- **{emotion_label.replace('_', ' ').title()}** (Confidence: {emotion_score:.2f})\n"
            suggestions_output_list.extend(get_suggestions(emotion_label))

    if not detected_emotions_display:
        detected_emotions_display = "No specific emotions detected with high confidence."
        suggestions_output_list.extend(get_suggestions("neutral")) # Fallback to neutral suggestions

    # Remove duplicates from suggestions while preserving order
    unique_suggestions = []
    seen = set()
    for suggestion in suggestions_output_list:
        if suggestion not in seen:
            unique_suggestions.append(suggestion)
            seen.add(suggestion)

    suggestion_text = "\n\n**💡 Here are some suggestions for you:**\n"
    if unique_suggestions:
        for suggestion in unique_suggestions:
            suggestion_text += f"- {suggestion}\n"
    else:
        suggestion_text += "No specific suggestions generated based on your entry."

    # Get a primary emotion for the motivational response (e.g., the top one)
    primary_emotion = sorted_results[0]['label'] if sorted_results else "neutral"
    response = get_motivation(primary_emotion)
    
    return f"🧠 Detected Emotions:\n{detected_emotions_display}", response, suggestion_text, unique_suggestions


# --- Calendar To-Do List Feature Setup ---
TASKS_FILE = "calendar_tasks.json"

def load_calendar_tasks():
    if os.path.exists(TASKS_FILE):
        try:
            with open(TASKS_FILE, 'r') as f:
                return json.load(f)
        except json.JSONDecodeError:
            print(f"Warning: {TASKS_FILE} is corrupted or empty. Starting with empty tasks.")
            return {}
    return {} # { "YYYY-MM-DD": [{"item": "...", "completed": False}, ...] }

def save_calendar_tasks(tasks_data):
    with open(TASKS_FILE, 'w') as f:
        json.dump(tasks_data, f, indent=4)

# Load tasks when the app starts
global_calendar_tasks = load_calendar_tasks()

def format_date_input(selected_date):
    """Ensures the date is in "YYYY-MM-DD" format from various Gradio date inputs."""
    if isinstance(selected_date, datetime.date):
        return selected_date.strftime("%Y-%m-%d")
    elif isinstance(selected_date, datetime.datetime):
        return selected_date.strftime("%Y-%m-%d")
    elif isinstance(selected_date, str):
        try:
            # Handle potential ISO format strings from gr.DateTime
            # Gradio's gr.DateTime often returns ISO formatted strings like "YYYY-MM-DDTHH:MM:SS.sssZ"
            dt_obj = datetime.datetime.fromisoformat(selected_date.replace('Z', '+00:00'))
            return dt_obj.strftime("%Y-%m-%d")
        except ValueError:
            # Assume it's already "YYYY-MM-DD" if other conversions fail
            return selected_date
    return datetime.date.today().strftime("%Y-%m-%d") # Default fallback

def get_tasks_for_date(selected_date):
    selected_date_str = format_date_input(selected_date)
    return global_calendar_tasks.get(selected_date_str, [])

def display_calendar_tasks(selected_date):
    selected_date_str = format_date_input(selected_date)
    tasks = get_tasks_for_date(selected_date_str)
    
    display_title = f"### To-Do List for {selected_date_str}:"
    if not tasks:
        display_title += "\nNo tasks for this date. Add one!"
    
    # We will no longer generate HTML directly.
    # Instead, we will configure the CheckboxGroup to reflect the current tasks and their completion status.
    
    task_labels = [task['item'] for task in tasks]
    completed_task_labels = [task['item'] for task in tasks if task['completed']]

    # The CheckboxGroup's choices are all task items.
    # Its value is the list of items that are completed.
    return (
        display_title, 
        gr.CheckboxGroup(choices=task_labels, value=completed_task_labels, visible=True, label="Tasks"),
        "" # Clear the new item textbox
    )


def add_calendar_todo_item(selected_date, item):
    selected_date_str = format_date_input(selected_date)

    if not item.strip():
        # Re-display tasks without adding empty item, clear input
        title, checkbox_group_output, _ = display_calendar_tasks(selected_date_str)
        return title, checkbox_group_output, "" # Clear input
    
    tasks_for_today = global_calendar_tasks.get(selected_date_str, [])
    # Check if the item already exists to prevent duplicates if you want to.
    # For now, we'll allow duplicates to simplify.
    
    tasks_for_today.append({"item": item.strip(), "completed": False})
    global_calendar_tasks[selected_date_str] = tasks_for_today
    save_calendar_tasks(global_calendar_tasks)
    
    # Re-display tasks and clear input
    title, checkbox_group_output, _ = display_calendar_tasks(selected_date_str)
    return title, checkbox_group_output, "" # Clear input

def toggle_calendar_todo_item(selected_date, selected_tasks_labels_from_checkbox_group):
    selected_date_str = format_date_input(selected_date)
    tasks = global_calendar_tasks.get(selected_date_str, [])

    # Update the completed status for each task based on the CheckboxGroup's value
    for task in tasks:
        task["completed"] = (task["item"] in selected_tasks_labels_from_checkbox_group)

    save_calendar_tasks(global_calendar_tasks)
    
    # Re-display the tasks to update the visual representation (e.g., strikethrough)
    # The display_calendar_tasks function will now return a new gr.CheckboxGroup
    # with updated choices and values.
    return display_calendar_tasks(selected_date)


def clear_completed_tasks_today(selected_date):
    selected_date_str = format_date_input(selected_date)
    
    if selected_date_str in global_calendar_tasks:
        # Filter out completed tasks
        global_calendar_tasks[selected_date_str] = [
            task for task in global_calendar_tasks[selected_date_str] if not task['completed']
        ]
        save_calendar_tasks(global_calendar_tasks)
    
    # Re-display the tasks for today after clearing completed ones
    # This will now correctly return a new gr.CheckboxGroup with fewer items if tasks were cleared
    return display_calendar_tasks(selected_date)


# This function is called to update the choices of the gr.Radio component
def populate_suggestions_radio(suggestions_list):
    """
    Populates the gr.Radio component with suggestions.
    This function should be directly mapped to the outputs of an event.
    """
    if suggestions_list:
        # If there are suggestions, make it visible and populate choices
        return gr.Radio(choices=suggestions_list, value=None, visible=True)
    else:
        # If no suggestions, ensure choices are empty and it's hidden
        return gr.Radio(choices=[], value=None, visible=False)

def add_suggestion_to_todo(selected_date, suggestion_to_add):
    if suggestion_to_add:
        # Call the existing add_calendar_todo_item function
        title, checkbox_group_output, new_item_textbox_val = add_calendar_todo_item(selected_date, suggestion_to_add)
        # Clear the selected radio button after adding
        return title, checkbox_group_output, new_item_textbox_val, gr.Radio(value=None, visible=True) # Reset radio selection
    else:
        # If no suggestion selected, just refresh the display and clear input field
        title, checkbox_group_output, _ = display_calendar_tasks(selected_date)
        return title, checkbox_group_output, gr.Textbox(value="", interactive=True), gr.Radio(value=None, visible=True)


# --- Combined Gradio Interface (All in one frame) ---
with gr.Blocks(title="✨ Your Multi-purpose AI Companion") as demo:
    gr.Markdown("# ✨ Your Multi-purpose AI Companion")
    gr.Markdown("Welcome! This AI buddy offers journaling and a to-do list all in one place.")

    # --- Journal Buddy Section ---
    gr.HTML("<hr>")
    gr.Markdown("## 📝 Journal Buddy")
    gr.Markdown("Reflect, vent, or write freely. Get feedback based on your emotional tone.")
    with gr.Column():
        journal_input = gr.Textbox(label="Your Journal Entry", placeholder="Write how your day went or what's on your mind...", lines=5)
        journal_analyze_button = gr.Button("Analyze Journal")
        journal_output_mood = gr.Markdown(label="Detected Mood")
        journal_output_response = gr.Text(label="Journal Buddy's Response")
        journal_output_suggestions_markdown = gr.Markdown(label="Suggestions for You") # Markdown for display
        
        # Hidden state to pass suggestions to the calendar tab
        journal_suggestions_list_state = gr.State(value=[]) 

        journal_analyze_result = journal_analyze_button.click( # Store the event object
            fn=analyze_journal,
            inputs=journal_input,
            outputs=[
                journal_output_mood,
                journal_output_response,
                journal_output_suggestions_markdown,
                journal_suggestions_list_state # This updates the gr.State
            ]
        )
        gr.Examples(
            examples=[
                ["I had a fantastic day, everything went perfectly and I feel so happy!"],
                ["I'm really struggling today, nothing seems to be going right and I feel so sad and frustrated."],
                ["I'm so angry about what happened at work, I can't believe it. I also feel a bit disappointed."],
                ["Today was just a regular day, nothing special happened."],
                ["I am so grateful for the sunny weather and a good cup of coffee this morning. Feeling very content."]
            ],
            inputs=journal_input
        )

    # --- Calendar To-Do List Section ---
    gr.HTML("<hr>")
    gr.Markdown("## 📅 Calendar To-Do List")
    gr.Markdown("This To-Do list is automatically set to today's date.")
    
    with gr.Tabs() as calendar_tabs: # Use gr.Tabs for multiple sections within the To-Do list
        with gr.TabItem("Manage To-Dos", id="manage_todos_tab"): # Added ID for clarity
            with gr.Column():
                # Automatically set to today's date and make it non-interactive
                calendar_date_picker = gr.DateTime(
                    label="Today's Date",
                    value=datetime.datetime.now(), 
                    type="date", 
                    info="Tasks are managed for today's date.",
                    interactive=False # Key change: Make it non-interactive
                )
                    
                calendar_todo_title = gr.Markdown() # To display the date title
                
                # IMPORTANT CHANGE: Use gr.CheckboxGroup directly for displaying tasks
                calendar_todo_checkboxes_display = gr.CheckboxGroup(
                    label="Your Tasks",
                    choices=[], # Will be populated by display_calendar_tasks
                    value=[],   # Will be populated by display_calendar_tasks
                    interactive=True # Allow users to check/uncheck
                )
                
                with gr.Row():
                    new_item_textbox_calendar = gr.Textbox(label="New To-Do Item", placeholder="e.g., Team meeting", scale=4)
                    add_button_calendar = gr.Button("Add Item", scale=1)

                clear_completed_button_calendar = gr.Button("Clear Done Tasks Today", variant="secondary")


                # This ensures tasks are displayed when the app initially loads, for today's date
                demo.load(
                    fn=lambda: display_calendar_tasks(datetime.date.today()),
                    outputs=[calendar_todo_title, calendar_todo_checkboxes_display, new_item_textbox_calendar]
                )

                add_button_calendar.click(
                    fn=add_calendar_todo_item,
                    inputs=[calendar_date_picker, new_item_textbox_calendar],
                    outputs=[calendar_todo_title, calendar_todo_checkboxes_display, new_item_textbox_calendar]
                )

                # IMPORTANT CHANGE: Wire the `change` event of the visible CheckboxGroup
                calendar_todo_checkboxes_display.change(
                    fn=toggle_calendar_todo_item,
                    inputs=[calendar_date_picker, calendar_todo_checkboxes_display],
                    outputs=[calendar_todo_title, calendar_todo_checkboxes_display, new_item_textbox_calendar] # Ensure all outputs are updated
                )
                
                clear_completed_button_calendar.click(
                    fn=clear_completed_tasks_today,
                    inputs=calendar_date_picker,
                    outputs=[calendar_todo_title, calendar_todo_checkboxes_display, new_item_textbox_calendar]
                )
                    
        with gr.TabItem("AI Suggestions", id="ai_suggestions_tab"): # Added ID for clarity
            gr.Markdown("### Add Journal Suggestions to your To-Do List")
            gr.Markdown("Select a suggestion from your last journal entry to add it directly to the To-Do list for today.")
            
            # Display suggestions as radio buttons for selection
            ai_suggestions_radio = gr.Radio(label="Choose a suggestion to add:", choices=[], value=None, visible=False)
            add_suggestion_button = gr.Button("Add Selected Suggestion to To-Do List")

            # Corrected logic: Use the .then() method to ensure sequential updates
            # First, analyze journal and update state. Then, use the state to populate the radio.
            journal_analyze_result.then( # Use the stored event object from journal_analyze_button.click
                fn=populate_suggestions_radio,
                inputs=journal_suggestions_list_state, # Input the state that was just updated
                outputs=ai_suggestions_radio # Output to the radio button to update its choices
            ).then( # Chain another .then() to switch tabs
                fn=lambda: gr.Tabs(selected="ai_suggestions_tab"),
                outputs=calendar_tabs
            )

            # Action to add selected suggestion to the to-do list
            add_suggestion_button.click(
                fn=add_suggestion_to_todo,
                inputs=[calendar_date_picker, ai_suggestions_radio],
                outputs=[calendar_todo_title, calendar_todo_checkboxes_display, new_item_textbox_calendar, ai_suggestions_radio]
            )
            
            # Clear suggestions when switching back to the "Manage To-Dos" tab
            # We don't need to clear suggestions on date_picker.change since it's non-interactive
            calendar_tabs.select(
                fn=lambda selected_tab: gr.Radio(choices=[], value=None, visible=False) if selected_tab == "manage_todos_tab" else gr.Noop(),
                inputs=calendar_tabs,
                outputs=ai_suggestions_radio
            )

demo.launch()