import gradio as gr
import google.generativeai as genai
import time
from dotenv import load_dotenv
import os
from vaderSentiment.vaderSentiment import SentimentIntensityAnalyzer

# Load environment variables from .env file
load_dotenv()

# Get the Google API Key from environment variables
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
genai.configure(api_key=GOOGLE_API_KEY)
# genai.configure(api_key="")

# Initialize the chatbot model
model = genai.GenerativeModel('gemini-1.5-flash-latest')
chat = model.start_chat(history=[])

# Initialize the sentiment analyzer
analyzer = SentimentIntensityAnalyzer()

# Function to detect user sentiment
def detect_tone(user_input):
    score = analyzer.polarity_scores(user_input)
    if score['compound'] >= 0.05:
        return "positive"
    elif score['compound'] <= -0.05:
        return "negative"
    else:
        return "neutral"

# Function to classify user type
def classify_user(user_input):
    user_input = user_input.lower()
    if "parent" in user_input:
        return "Parent of an ADHD child"
    elif "teacher" in user_input or "educator" in user_input:
        return "Teacher/Educator working with ADHD students"
    elif "employer" in user_input:
        return "Employer managing an ADHD employee"
    elif "adhd" in user_input or "adult" in user_input:
        return "Adult with ADHD"
    else:
        return "General user (uncategorized)"

# Function to transform Gradio history to Gemini format
def transform_history(history, system_prompt):
    new_history = []
    new_history.append({"parts": [{"text": system_prompt}], "role": "user"})
    for chat in history:
        new_history.append({"parts": [{"text": chat[0]}], "role": "user"})
        new_history.append({"parts": [{"text": chat[1]}], "role": "model"})
    return new_history

# Response function for chat
def response(message, history):
    global chat

    # Detect sentiment
    tone = detect_tone(message)

    # Classify user type (based on first few messages)
    if len(history) > 0:
        user_category = classify_user(history[0][0])  # First user message
    else:
        user_category = "General user (uncategorized)"


    system_prompt = """ 
    You are an ADHD support chatbot with the personality of a warm, supportive counselor. 
    Your job is to provide clear, concise, and structured guidance to users based on their needs.

    User type: {user_category}

    - If they are an ADHD adult, offer self-management strategies.
    - If they are a parent with ADHD children, provide parenting tips.
    - If they are a teacher, suggest ADHD-friendly teaching techniques.
    - If they are an employer, recommend workplace accommodations.

    Adjust your tone based on the user’s sentiment:
    - If they seem frustrated or overwhelmed, respond with extra warmth and reassurance.
    - If they are neutral, provide clear, structured advice.
    - If they are excited or positive, encourage and validate their progress.

    Most of your users are Indians, so make sure to reply considering that also. If they can also ask questions in Indian origin languages such as Hindi, Tamil, Malayalam, Marathi, Telugu, etc., be ready to do the conversation in that language itself.
    """
    if tone == "negative":
        system_prompt += "Your tone should be **extra empathetic and calming.** Offer reassurance and coping strategies."
    elif tone == "neutral":
        system_prompt += "Your tone should be **neutral and structured.** Provide clear, ADHD-friendly advice."
    elif tone == "positive":
        system_prompt += "Your tone should be **encouraging and celebratory.** Reinforce positive progress."        
    
    chat.history = transform_history(history, system_prompt)
    response = chat.send_message(message)
    response.resolve()

    for i in range(len(response.text)):
        time.sleep(0.005)
        yield response.text[:i+20]

# Gradio interface with rearranged components
with gr.Blocks() as demo:
    with gr.Column():
        # Banner and logo at the top
        # logo = gr.Image(value="logo.jpg", label="ADHDGuru_logo", show_label=False, interactive=False, height=150)  
        logo = gr.Image(value="banner.png", label="ADHDGuru_banner", show_label=False, interactive=False, height=150) 
        banner_text = gr.Markdown("## ADHDGuru - Your Well-being Companion")

        # chat.history = transform_history(history, system_prompt)
        # response = chat.send_message(message)
        # response.resolve()

        # Chat interface above the input field
        chat_interface = gr.ChatInterface(response,
                                          title="Chat with ADHDGuru",
                                          textbox=gr.Textbox(placeholder="Type here...")) #ADHDGuru Bot - Your Well-being Mentor and Companion
        # This will show the chat history first, then the input field will appear below it

# Launch the Gradio app
demo.launch(debug=True)