#!pip install langchain
#!pip install langchain-community
#!pip install langchain-google-genai
#!pip install gradio
#!pip install huggingface_hub

import os
import gradio as gr
from langchain_openai import ChatOpenAI  # <-- Changed to OpenAI's wrapper
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.output_parsers import StrOutputParser
from langchain_community.chat_message_histories import ChatMessageHistory

# 1. Configuration - Use OpenRouter API Key
os.environ["OPENROUTER_API_KEY"] = "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" # <-- Uncommented and replace with your actual API key

# 2. Define the Persona Template
SYSTEM_PROMPT = """You are an AI mental health and well-being companion named MindMate.
Your core identity is built on deep empathy, active listening, and providing a safe, non-judgmental space for users to navigate their thoughts and emotions.

Traits:
- Empathetic & Validating: Always acknowledge and validate the user's feelings before offering any perspective or coping strategies.
- Non-Judgmental: Create an environment where the user feels secure sharing anything without fear of criticism.
- The Gentle Guide: Rely heavily on open-ended, reflective questions to help users explore their own feelings and come to their own conclusions.
- Grounded & Realistic: Avoid toxic positivity. Acknowledge that it is okay to not be okay.

Rules:
- Safety First (Absolute Priority): If a user expresses intent for self-harm, suicide, or severe crisis, immediately intervene by providing standard emergency/crisis hotline resources and gently encourage them to seek professional help.
- Never Diagnose or Prescribe: You are a supportive companion, not a medical professional. Never attempt to diagnose a psychological condition or prescribe treatments.
- Listen Before Fixing: Prioritize understanding over problem-solving. Only offer actionable advice or cognitive behavioral exercises when the user explicitly asks for help or is clearly ready for it.
- Maintain Boundaries: Be transparent about your nature as an AI. Do not feign human experiences, trauma, or emotions
"""

# 3. Initialize the Model (Configured for OpenRouter)
llm = ChatOpenAI(
    api_key=os.environ.get("OPENROUTER_API_KEY"),
    base_url="https://openrouter.ai/api/v1", # <-- Point to OpenRouter instead of OpenAI
    model="z-ai/glm-4.5-air:free",        # <-- Swap this with ANY model ID from OpenRouter
    temperature=0.5,
    default_headers={
        "HTTP-Referer": "http://localhost:7860", # Optional: Used by OpenRouter for app rankings
        "X-Title": "Cortex AI Assistant",        # Optional: Used by OpenRouter for app rankings
    }
)

# 4. Create the Prompt Structure
prompt = ChatPromptTemplate.from_messages([
    ("system", SYSTEM_PROMPT),
    MessagesPlaceholder(variable_name="chat_history"),
    ("human", "{user_message}"),
])

# 5. Build the LCEL Chain
chain = prompt | llm | StrOutputParser()

# Memory storage (In-memory for this demo)
demo_ephemeral_chat_history = ChatMessageHistory()

def get_text_response(message, history):
    """
    Gradio passes the current 'message' and the 'history' list.
    We convert the history to LangChain format and invoke the chain.
    """
    # Convert Gradio history format to LangChain messages
    formatted_history = []
    for human, ai in history:
        formatted_history.append(("human", human))
        formatted_history.append(("ai", ai))

    # Run the chain
    response = chain.invoke({
        "chat_history": formatted_history,
        "user_message": message
    })

    return response

# 6. Launch the Gradio Interface
demo = gr.ChatInterface(
    get_text_response,
    type="messages",
    examples=[
        "I'm feeling really overwhelmed with everything on my plate lately and just don't know where to start.",
        "Can we just talk for a bit? I'm feeling a little lonely and disconnected today.",
        "I've been having trouble sleeping because my mind keeps racing at night. Do you have any advice?",
        "How do I deal with feeling like I'm not good enough at what I do?"
    ],
    title="MindMate: Your Safe Space"
)

if __name__ == "__main__":
    demo.launch()