import gradio as gr
import google.generativeai as genai
import os
from datetime import datetime
import plotly.express as px
import pandas as pd
from langdetect import detect
import random
from google.oauth2.credentials import Credentials
from googleapiclient.discovery import build
from googleapiclient.errors import HttpError
import base64
from email.mime.text import MIMEText
import re
import json

# Configure Gemini API (expects GEMINI_API_KEY in Hugging Face secrets)
genai.configure(api_key=os.getenv("GEMINI_API_KEY"))

# Gmail API configuration (expects GMAIL_ADDRESS, GMAIL_TOKEN_JSON in secrets)
GMAIL_ADDRESS = os.getenv("GMAIL_ADDRESS")
GMAIL_TOKEN_JSON = os.getenv("GMAIL_TOKEN_JSON")

# Initialize Gmail API client
def get_gmail_service():
    try:
        creds = Credentials.from_authorized_user_info(json.loads(GMAIL_TOKEN_JSON))
        return build("gmail", "v1", credentials=creds)
    except Exception as e:
        raise Exception(f"Failed to initialize Gmail API: {str(e)}")

# Coping Strategies Library
coping_strategies = {
    "happy": [
        "Keep the positivity flowing! Try writing down three things you’re grateful for today.",
        "Share your joy! Call a friend or loved one to spread the good vibes."
    ],
    "sad": [
        "It’s okay to feel this way. Try a gentle activity like listening to calming music or taking a short walk.",
        "Write down your thoughts in a journal to process what’s on your mind."
    ],
    "anxious": [
        "Take slow, deep breaths: inhale for 4 seconds, hold for 4, exhale for 4.",
        "Try a grounding exercise: name 5 things you see, 4 you can touch, 3 you hear, 2 you smell, 1 you taste."
    ],
    "stressed": [
        "Pause for a moment and stretch your body to release tension.",
        "Break tasks into smaller steps and tackle one at a time."
    ],
    "other": [
        "Reflect on what’s on your mind with a quick mindfulness moment: focus on your breath for 60 seconds.",
        "Engage in a favorite hobby to lift your spirits."
    ]
}

# Regional Resources
regional_resources = {
    "USA": [
        "National Suicide Prevention Lifeline: 1-800-273-8255",
        "Crisis Text Line: Text HOME to 741741",
        "[MentalHealth.gov](https://www.mentalhealth.gov/)"
    ],
    "India": [
        "Vandrevala Foundation: 1860-2662-345",
        "AASRA Suicide Prevention: +91-9820466726",
        "[iCall Helpline](https://icallhelpline.org/)"
    ],
    "UK": [
        "Samaritans: 116 123",
        "Shout Crisis Text Line: Text SHOUT to 85258",
        "[Mind UK](https://www.mind.org.uk/)"
    ],
    "Global": [
        "Befrienders Worldwide: [Find a helpline](https://befrienders.org/)",
        "[WHO Mental Health Resources](https://www.who.int/health-topics/mental-health)"
    ]
}

# Mock Therapist Database
therapists = {
    "Dr. Jane Smith": {
        "specialty": "Anxiety and Depression",
        "email": "jane.smith@example.com",
        "times": ["09:00", "11:00", "14:00", "16:00"]
    },
    "Dr. Amit Patel": {
        "specialty": "Stress Management",
        "email": "amit.patel@example.com",
        "times": ["10:00", "12:00", "15:00", "17:00"]
    },
    "Dr. Sarah Brown": {
        "specialty": "Trauma and PTSD",
        "email": "sarah.brown@example.com",
        "times": ["08:00", "13:00", "15:30", "18:00"]
    }
}

# Initialize Gemini model
model = genai.GenerativeModel("gemini-1.5-flash", 
                             generation_config={"temperature": 0.7, "max_output_tokens": 200})

# Chatbot function
def chatbot_function(message, mood, conversation_mode, region, state):
    # Initialize chat history if not present
    if "chat_history" not in state:
        state["chat_history"] = []
    
    history = state["chat_history"]
    
    # Multilingual Support: Detect input language
    try:
        lang = detect(message) if message.strip() else "en"
    except:
        lang = "en"
    
    # Append mood to history if provided
    if mood and mood != "Select mood (optional)":
        history.append([f"I'm feeling {mood.lower()}.", None])
    
    # Append user message
    history.append([message, None])
    
    # Themed Conversation Modes: Adjust tone
    tone_instruction = {
        "Calm": "Respond in a soothing, gentle tone to promote relaxation.",
        "Motivational": "Use an uplifting, encouraging tone to inspire confidence.",
        "Neutral": "Maintain a balanced, empathetic tone."
    }.get(conversation_mode, "Maintain a balanced, empathetic tone.")
    
    # Prepare prompt
    prompt = f"""
    You are a compassionate mental health support chatbot. Engage in a supportive conversation with the user based on their input: {message}. 
    - Provide empathetic, sensitive responses in the user's language (detected as {lang}).
    - {tone_instruction}
    - If signs of distress are detected, suggest coping strategies relevant to their mood or input.
    - Recommend professional resources tailored to the user's region ({region}).
    - Keep responses concise, warm, and encouraging.
    """
    
    # Generate response
    try:
        response = model.generate_content(prompt)
        response_text = response.text
    except Exception:
        response_text = "I'm here for you. Could you share a bit more so I can support you better?"
    
    # Add coping strategy if mood is provided
    if mood and mood != "Select mood (optional)":
        mood_key = mood.lower()
        if mood_key in coping_strategies:
            strategy = random.choice(coping_strategies[mood_key])
            response_text += f"\n\n**Coping Strategy**: {strategy}"
    
    # Add regional resources
    region_key = region if region in regional_resources else "Global"
    resources = "\n\n**Recommended Resources**:\n" + "\n".join(regional_resources[region_key])
    response_text += resources
    
    # Update history
    history[-1][1] = response_text
    state["chat_history"] = history
    
    # Format history for display
    chat_display = ""
    for user_msg, bot_msg in history:
        if user_msg:
            chat_display += f"**You**: {user_msg}\n\n"
        if bot_msg:
            chat_display += f"**Bot**: {bot_msg}\n\n"
    
    return chat_display, state

# Clear chat history
def clear_chat(state):
    state["chat_history"] = []
    return "", state

# Mood journal function
def log_mood(mood, state):
    if mood and mood != "Select mood (optional)":
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        state["mood_journal"].append({"timestamp": timestamp, "mood": mood.lower()})
        return "Mood logged successfully!", state
    return "Please select a mood to log.", state

# Mood trend visualization
def show_mood_trends(state):
    if not state["mood_journal"]:
        return "No moods logged yet.", state
    df = pd.DataFrame(state["mood_journal"])
    fig = px.line(df, x="timestamp", y="mood", title="Mood Trends Over Time", markers=True)
    return fig, state

# Feedback function
def submit_feedback(feedback_text, rating, state):
    if feedback_text or rating:
        state["feedback"].append({"text": feedback_text, "rating": rating, "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")})
        return "Thank you for your feedback!", state
    return "Please provide feedback or a rating.", state

# Emergency resources
def show_emergency_resources():
    return """
    **Crisis Support:**
    - National Suicide Prevention Lifeline: 1-800-273-8255
    - Crisis Text Line: Text HOME to 741741
    - [MentalHealth.gov](https://www.mentalhealth.gov/)
    """

# Get available times for selected therapist
def get_available_times(therapist):
    if therapist and therapist in therapists:
        return gr.update(choices=therapists[therapist]["times"], value=None), f"Available times for {therapist} loaded."
    return gr.update(choices=[], value=None), "Please select a therapist."

# Create MIME message for Gmail API
def create_message(to, subject, message_text):
    message = MIMEText(message_text)
    message["to"] = to
    message["from"] = GMAIL_ADDRESS
    message["subject"] = subject
    raw = base64.urlsafe_b64encode(message.as_bytes()).decode()
    return {"raw": raw}

# Send emails using Gmail API
def send_emails(therapist, time_slot, date, user_email, therapist_email, state):
    if "failed_emails" not in state:
        state["failed_emails"] = []
    
    # Therapist email
    therapist_body = f"""
Dear {therapist},

You have a new appointment:
- Date: {date}
- Time: {time_slot}
- Client Email: {user_email}

Please confirm with the client.

Best,
Mental Health Chatbot
    """
    therapist_message = create_message(therapist_email, "New Appointment", therapist_body)
    
    # User email
    user_body = f"""
Dear User,

Your appointment is booked:
- Therapist: {therapist}
- Date: {date}
- Time: {time_slot}

Expect a confirmation from {therapist}.

Best,
Mental Health Chatbot
    """
    user_message = create_message(user_email, "Appointment Confirmation", user_body)
    
    try:
        service = get_gmail_service()
        
        # Send therapist email
        therapist_result = service.users().messages().send(userId="me", body=therapist_message).execute()
        # Send user email
        user_result = service.users().messages().send(userId="me", body=user_message).execute()
        
        return True, ""
    except HttpError as e:
        error_msg = f"Gmail API error: {str(e)}"
        state["failed_emails"].append({
            "therapist_email": {"to": therapist_email, "subject": "New Appointment", "body": therapist_body},
            "user_email": {"to": user_email, "subject": "Appointment Confirmation", "body": user_body},
            "error": error_msg,
            "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        })
        return False, error_msg
    except Exception as e:
        error_msg = f"Unexpected error: {str(e)}"
        state["failed_emails"].append({
            "therapist_email": {"to": therapist_email, "subject": "New Appointment", "body": therapist_body},
            "user_email": {"to": user_email, "subject": "Appointment Confirmation", "body": user_body},
            "error": error_msg,
            "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        })
        return False, error_msg

# Schedule appointment
def schedule_appointment(therapist, time_slot, date, user_email, state):
    if not therapist or therapist not in therapists:
        return "Please select a therapist.", state
    if not time_slot or time_slot not in therapists[therapist]["times"]:
        return "Please select a valid time slot.", state
    if not date:
        return "Please select a date.", state
    if not user_email or not re.match(r"[^@]+@[^@]+\.[^@]+", user_email):
        return "Please enter a valid email address.", state
    
    try:
        appointment_date = datetime.strptime(date, "%Y-%m-%d")
        if appointment_date < datetime.now():
            return "Please select a future date.", state
    except ValueError:
        return "Invalid date format. Use YYYY-MM-DD.", state
    
    # Store appointment
    appointment = {
        "therapist": therapist,
        "time": time_slot,
        "date": date,
        "user_email": user_email,
        "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    }
    state["appointments"].append(appointment)
    
    # Send emails
    therapist_email = therapists[therapist]["email"]
    success, error_msg = send_emails(therapist, time_slot, date, user_email, therapist_email, state)
    if success:
        return f"Appointment booked with {therapist} on {date} at {time_slot}. Emails sent from {GMAIL_ADDRESS} to {therapist_email} and {user_email}!", state
    else:
        return (f"Appointment booked with {therapist} on {date} at {time_slot}. "
                f"Emails from {GMAIL_ADDRESS} could not be sent. "
                f"Please email {therapist_email} with your details (date: {date}, time: {time_slot}) "
                f"and check {user_email} for confirmation (spam/junk)."), state

# Gradio Interface
with gr.Blocks(title="Mental Health Support Chatbot") as demo:
    # Initialize state
    state = gr.State({"chat_history": [], "mood_journal": [], "feedback": [], "appointments": [], "failed_emails": []})
    
    gr.Markdown("# Mental Health Support Chatbot")
    gr.Markdown("I'm here to listen and support you. Feel free to share how you're feeling.")
    
    # Mood Indicator
    mood = gr.Dropdown(
        choices=["Select mood (optional)", "Happy", "Sad", "Anxious", "Stressed", "Other"],
        label="How are you feeling today? (Optional)"
    )
    
    # Conversation Mode
    conversation_mode = gr.Radio(
        choices=["Neutral", "Calm", "Motivational"],
        label="Conversation Style",
        value="Neutral"
    )
    
    # Region Selector
    region = gr.Dropdown(
        choices=["USA", "India", "UK", "Global"],
        label="Select your region for tailored resources",
        value="Global"
    )
    
    # Chat Interface
    chatbot = gr.Textbox(
        label="Conversation",
        value="",
        interactive=False,
        lines=10
    )
    user_input = gr.Textbox(
        placeholder="Type your message here...",
        label="Your Message"
    )
    
    # Buttons
    with gr.Row():
        submit_btn = gr.Button("Send")
        clear_btn = gr.Button("Clear Chat")
    
    # Emergency Resources
    emergency_btn = gr.Button("Emergency Resources")
    emergency_output = gr.Markdown("")
    
    # Mood Journal
    with gr.Accordion("Mood Journal"):
        log_mood_btn = gr.Button("Log Mood")
        mood_log_output = gr.Textbox(label="Mood Log Status", interactive=False)
        mood_trend_btn = gr.Button("Show Mood Trends")
        mood_trend_output = gr.Plot()
    
    # Feedback
    with gr.Accordion("Provide Feedback"):
        feedback_text = gr.Textbox(label="Your Feedback (Optional)", placeholder="How can we improve?")
        feedback_rating = gr.Slider(minimum=1, maximum=5, step=1, label="Rate this interaction")
        feedback_btn = gr.Button("Submit Feedback")
        feedback_output = gr.Textbox(label="Feedback Status", interactive=False)
    
    # Appointment Scheduling
    with gr.Accordion("Schedule Appointment"):
        therapist = gr.Dropdown(
            choices=list(therapists.keys()),
            label="Select a Therapist"
        )
        time_slot = gr.Dropdown(
            choices=[],
            label="Select a Time Slot",
            value=None,
            interactive=True
        )
        date = gr.Textbox(
            label="Appointment Date (YYYY-MM-DD)",
            placeholder="e.g., 2025-04-20"
        )
        user_email = gr.Textbox(
            label="Your Email Address",
            placeholder="e.g., user@example.com"
        )
        schedule_btn = gr.Button("Book Appointment")
        schedule_output = gr.Textbox(label="Booking Status", interactive=False)
    
    # Event Handlers
    submit_btn.click(
        fn=chatbot_function,
        inputs=[user_input, mood, conversation_mode, region, state],
        outputs=[chatbot, state]
    )
    user_input.submit(
        fn=chatbot_function,
        inputs=[user_input, mood, conversation_mode, region, state],
        outputs=[chatbot, state]
    )
    clear_btn.click(
        fn=clear_chat,
        inputs=[state],
        outputs=[chatbot, state]
    )
    emergency_btn.click(
        fn=show_emergency_resources,
        inputs=None,
        outputs=emergency_output
    )
    log_mood_btn.click(
        fn=log_mood,
        inputs=[mood, state],
        outputs=[mood_log_output, state]
    )
    mood_trend_btn.click(
        fn=show_mood_trends,
        inputs=[state],
        outputs=[mood_trend_output, state]
    )
    feedback_btn.click(
        fn=submit_feedback,
        inputs=[feedback_text, feedback_rating, state],
        outputs=[feedback_output, state]
    )
    therapist.change(
        fn=get_available_times,
        inputs=therapist,
        outputs=[time_slot, schedule_output]
    )
    schedule_btn.click(
        fn=schedule_appointment,
        inputs=[therapist, time_slot, date, user_email, state],
        outputs=[schedule_output, state]
    )

    # Resource Links
    gr.Markdown("""
    ---
    **Helpful Resources:**
    - [National Suicide Prevention Lifeline](https://suicidepreventionlifeline.org/)
    - [MentalHealth.gov](https://www.mentalhealth.gov/)
    - [Crisis Text Line](https://www.crisistextline.org/)
    """)

# Launch (for local testing)
if __name__ == "__main__":
    demo.launch()