import streamlit as st
from huggingface_hub import InferenceClient
from datetime import datetime
import pandas as pd
import os
import uuid
import json
import re

# Page configuration
st.set_page_config(
    page_title="Mental Wellness Diary Analyzer",
    page_icon="🧠",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS for better appearance
st.markdown("""
<style>
    .main .block-container {
        padding-top: 2rem;
        padding-bottom: 2rem;
    }
    .stTextArea textarea {
        font-size: 16px;
    }
    .journal-history {
        border-left: 2px solid #f0f2f6;
        padding-left: 20px;
    }
    .css-1v3fvcr {
        background-color: #fcfcff;
    }
    .stButton button {
        background-color: #4e89ae;
        color: white;
        border-radius: 5px;
        padding: 0.5rem 1rem;
        font-weight: bold;
    }
    .stButton button:hover {
        background-color: #3d7092;
    }
</style>
""", unsafe_allow_html=True)

# App title and description
st.title("🧠 Mental Wellness Diary Analyzer")
st.markdown("""
Track your mental wellness journey with AI-powered insights. Write your daily thoughts, 
and receive personalized analysis to help you understand your emotional patterns.
""")

# Initialize session state for journal entries
if 'journal_entries' not in st.session_state:
    # Try to load from file if it exists
    if os.path.exists('journal_entries.json'):
        try:
            with open('journal_entries.json', 'r') as f:
                st.session_state.journal_entries = json.load(f)
        except:
            st.session_state.journal_entries = []
    else:
        st.session_state.journal_entries = []

# Hugging Face API setup
def get_hf_api_token():
    # First check if token exists in secrets
    if hasattr(st, 'secrets') and 'huggingface' in st.secrets:
        return st.secrets["huggingface"]["api_token"]
    
    # For local development, check environment variable
    return os.environ.get("HF_API_TOKEN")

# Initialize the InferenceClient
@st.cache_resource
def get_inference_client():
    api_token = get_hf_api_token()
    if not api_token:
        st.error("No Hugging Face API token found. Please configure it in your Spaces secrets.")
        return None
    
    return InferenceClient(
        provider="nebius",
        api_key=api_token,
    )

# Clean the output to remove thinking process
def clean_output(text):
    # Extract only the parts after the headings
    cleaned_text = ""
    
    # Extract Emotional tone section
    if "Emotional tone:" in text:
        match = re.search(r'Emotional tone:(.*?)(?:Recurring themes:|$)', text, re.DOTALL)
        if match:
            emotional_tone = match.group(1).strip()
            cleaned_text += f"😊 **Emotional tone:** {emotional_tone}\n\n"
    
    # Extract Recurring themes section
    if "Recurring themes:" in text:
        match = re.search(r'Recurring themes:(.*?)(?:Advice:|$)', text, re.DOTALL)
        if match:
            themes = match.group(1).strip()
            cleaned_text += f"🔍 **Recurring themes:** {themes}\n\n"
    
    # Extract Advice section
    if "Advice:" in text:
        match = re.search(r'Advice:(.*?)$', text, re.DOTALL)
        if match:
            advice = match.group(1).strip()
            cleaned_text += f"💡 **Advice:** {advice}"
    
    return cleaned_text if cleaned_text else text

# Function to analyze journal entries
def analyze_journal(entry, model_name="deepseek-ai/DeepSeek-R1"):
    client = get_inference_client()
    if not client:
        return "Error: Could not initialize the Inference Client."
    
    try:
        prompt = f"""You are a compassionate mental wellness assistant. Analyze the following journal entry with empathy and provide helpful insights.

Journal entry: "{entry}"

Please provide the following analysis directly, without explaining your thinking process:

1. Emotional tone: Identify the primary emotions expressed in the entry
2. Recurring themes: Note any patterns or important topics mentioned
3. Advice: Offer kind, supportive guidance tailored to what was shared

Format your response with only these three sections, using the exact headings: "Emotional tone:", "Recurring themes:", and "Advice:". Do not include any additional explanations or thinking.
"""
        
        completion = client.chat.completions.create(
            model=model_name,
            messages=[
                {
                    "role": "user",
                    "content": prompt
                }
            ],
            max_tokens=512,
            temperature=0.7,
        )
        
        output = completion.choices[0].message.content
        
        # Clean the output to remove any thinking process
        return clean_output(output)
        
    except Exception as e:
        st.error(f"API Error: {str(e)}")
        return f"Error analyzing journal entry: {str(e)}"

# Function to save an entry
def save_entry(entry, analysis):
    entry_id = str(uuid.uuid4())
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    
    # Extract emotional tone for tagging (simplified)
    emotional_tone = "neutral"
    if "Emotional tone:" in analysis:
        tone_text = analysis.split("Emotional tone:")[1].split("\n\n")[0].strip()
        if "anxious" in tone_text.lower() or "stress" in tone_text.lower():
            emotional_tone = "anxious"
        elif "happy" in tone_text.lower() or "joy" in tone_text.lower():
            emotional_tone = "happy"
        elif "sad" in tone_text.lower() or "depress" in tone_text.lower():
            emotional_tone = "sad"
        elif "angry" in tone_text.lower() or "frustrat" in tone_text.lower():
            emotional_tone = "angry"
    
    new_entry = {
        "id": entry_id,
        "date": timestamp,
        "entry": entry,
        "analysis": analysis,
        "emotion": emotional_tone
    }
    
    st.session_state.journal_entries.append(new_entry)
    
    # Save to file (for spaces persistence)
    try:
        with open('journal_entries.json', 'w') as f:
            json.dump(st.session_state.journal_entries, f)
    except Exception as e:
        st.warning(f"Could not save entries to file: {str(e)}")

# Sidebar for navigation and settings
with st.sidebar:
    st.header("Navigation")
    page = st.radio("", ["Write New Entry", "View Journal History"])
    
    st.header("Settings")
    model_choice = st.selectbox(
        "Language Model",
        ["deepseek-ai/DeepSeek-R1", "meta-llama/Llama-3-8b-chat", "microsoft/phi-3-mini-4k-instruct"],
        index=0,
        help="Select which AI model to use for analysis"
    )
    
    st.markdown("---")
    st.markdown("### About")
    st.info("""
    This Mental Wellness Diary Analyzer helps you track your emotional wellbeing.
    All entries are stored locally in your browser session.
    
    **Disclaimer**: This app is not a substitute for professional mental health support.
    """)

# Main content based on selected page
if page == "Write New Entry":
    st.header("📝 New Journal Entry")
    
    # Current date display
    st.markdown(f"**Today**: {datetime.now().strftime('%A, %B %d, %Y')}")
    
    # Journal input
    journal_entry = st.text_area(
        "How are you feeling today?",
        height=200,
        placeholder="Write about your day, feelings, thoughts, or anything on your mind..."
    )
    
    # Mood tracker (optional quick selection)
    cols = st.columns(5)
    with cols[0]:
        mood_happy = st.button("😊 Happy")
    with cols[1]:
        mood_calm = st.button("😌 Calm")
    with cols[2]:
        mood_anxious = st.button("😰 Anxious") 
    with cols[3]:
        mood_sad = st.button("😢 Sad")
    with cols[4]:
        mood_angry = st.button("😠 Angry")
    
    # Handle mood button clicks
    mood_prefix = ""
    if mood_happy:
        mood_prefix = "I'm feeling happy today. "
    elif mood_calm:
        mood_prefix = "I'm feeling calm and relaxed today. "
    elif mood_anxious:
        mood_prefix = "I'm feeling anxious today. "
    elif mood_sad:
        mood_prefix = "I'm feeling sad today. "
    elif mood_angry:
        mood_prefix = "I'm feeling angry and frustrated today. "
    
    if mood_prefix and not journal_entry.startswith(mood_prefix):
        journal_entry = mood_prefix + journal_entry
    
    # Analyze button
    col1, col2, col3 = st.columns([1, 1, 1])
    with col2:
        analyze_button = st.button("🧠 Analyze My Entry", use_container_width=True)
    
    # Sample entries
    with st.expander("Need inspiration? Try a sample entry"):
        sample_1 = st.button("Sample: Stressed at work")
        sample_2 = st.button("Sample: Good day")
        sample_3 = st.button("Sample: Relationship issues")
        
        if sample_1:
            journal_entry = "I feel really stressed and anxious about my work. I haven't been sleeping well, and I'm worried about meeting deadlines. My boss keeps adding more tasks, and I don't know how to say no."
        elif sample_2:
            journal_entry = "Today was a great day! I finished a project I've been working on for weeks, and my team was really impressed. I feel proud of myself and motivated to keep up the good work."
        elif sample_3:
            journal_entry = "I had an argument with my partner today. We've been fighting a lot lately about small things. I'm not sure if it's just stress or if there's a bigger issue we need to address."
    
    # Analysis section
    if analyze_button and journal_entry.strip():
        with st.spinner("Analyzing your thoughts..."):
            analysis = analyze_journal(journal_entry, model_choice)
        
        st.markdown("### 💭 Your Wellness Insights")
        st.markdown(analysis)
        
        # Save option
        if st.button("📋 Save this entry to your journal"):
            save_entry(journal_entry, analysis)
            st.success("Entry saved successfully!")
        
        # Supportive message
        st.markdown("---")
        st.markdown("💪 **Remember**: Writing about your feelings is already a positive step for your mental wellness!")
    
    elif analyze_button:
        st.warning("Please write in your journal entry first.")

else:  # View Journal History page
    st.header("📚 Your Journal History")
    
    if not st.session_state.journal_entries:
        st.info("You haven't saved any journal entries yet. Head to 'Write New Entry' to get started!")
    else:
        # Filter options
        col1, col2 = st.columns(2)
        with col1:
            filter_option = st.selectbox(
                "Filter by emotion:",
                ["All", "happy", "anxious", "sad", "angry", "neutral"]
            )
        with col2:
            sort_option = st.selectbox(
                "Sort by:",
                ["Newest first", "Oldest first"]
            )
        
        # Filter and sort entries
        filtered_entries = st.session_state.journal_entries
        if filter_option != "All":
            filtered_entries = [e for e in filtered_entries if e.get("emotion") == filter_option]
        
        if sort_option == "Newest first":
            filtered_entries = sorted(filtered_entries, key=lambda x: x.get("date", ""), reverse=True)
        else:
            filtered_entries = sorted(filtered_entries, key=lambda x: x.get("date", ""))
        
        # Display entries
        for i, entry in enumerate(filtered_entries):
            with st.expander(f"📝 {entry.get('date', 'Unknown date')}"):
                st.markdown("#### Journal Entry")
                st.write(entry.get("entry", ""))
                
                st.markdown("#### Analysis")
                st.markdown(entry.get("analysis", "No analysis available"))
                
                # Delete option
                if st.button(f"🗑️ Delete this entry", key=f"delete_{i}"):
                    st.session_state.journal_entries.remove(entry)
                    with open('journal_entries.json', 'w') as f:
                        json.dump(st.session_state.journal_entries, f)
                    st.experimental_rerun()
        
        # Add insights section if there are enough entries
        if len(st.session_state.journal_entries) >= 3:
            st.markdown("---")
            st.header("📊 Emotional Trends")
            
            # Create a dataframe for visualization
            entries_df = pd.DataFrame(st.session_state.journal_entries)
            entries_df['date'] = pd.to_datetime(entries_df['date'])
            
            # Count emotions
            emotion_counts = entries_df['emotion'].value_counts()
            st.bar_chart(emotion_counts)
            
            # Most common themes - this would require more sophisticated analysis
            st.markdown("### Common Themes")
            st.info("This feature will analyze common themes across your entries in a future update.")

# Footer
st.markdown("---")
st.caption("""
**Disclaimer**: This app is for educational purposes only and is not a substitute for professional mental health advice, diagnosis, or treatment. 
Always seek the advice of your physician or other qualified health provider with any questions regarding a medical condition.
""")