"""
Peace Pal - Mental Health Chatbot Application
This module implements a mental health chatbot using LangChain and Hugging Face models.
The chatbot provides supportive and non-judgmental guidance to users struggling with
mental health issues.

Author: Sowndharya,balakumaran
Date: October 29, 2024
"""

import os
from typing import List, Optional, Tuple
import gradio as gr
from langchain.prompts import PromptTemplate
from langchain_community.llms import HuggingFaceEndpoint
from transformers import pipeline

# Configuration Constants
DEFAULT_MODEL = "google/gemma-1.1-7b-it"
#DEFAULT_MODEL = "meta-llama/Llama-3.2-3B"
MAX_NEW_TOKENS = 512
TOP_K = 5
TEMPERATURE = 0.2
REPETITION_PENALTY = 1.03
SUICIDE_HELPLINE = "+91 91529 87821"

class MentalHealthChatbot:
    """
    A chatbot class specifically designed for mental health support and guidance.
    
    This class handles the initialization of the language model and provides
    methods for processing user inputs and generating appropriate responses.
    """
    
    def __init__(
        self,
        model_id: str = DEFAULT_MODEL,
        api_token: Optional[str] = None
    ) -> None:
        """
        Initialize the chatbot with specified model and configurations.
        
        Args:
            model_id: The Hugging Face model identifier to use
            api_token: Hugging Face API token for authentication
        """
        self.api_token = api_token or os.getenv("HF_TOKEN")
        if not self.api_token:
            raise ValueError("Hugging Face API token not found")
            
        self.llm = self._initialize_llm(model_id)
        self.prompt_template = self._create_prompt_template()
        
    def _initialize_llm(self, model_id: str) -> HuggingFaceEndpoint:
        """Initialize the language model with specified configurations."""
        return HuggingFaceEndpoint(
            repo_id=model_id,
            task="text-generation",
            max_new_tokens=MAX_NEW_TOKENS,
            top_k=TOP_K,
            temperature=TEMPERATURE,
            repetition_penalty=REPETITION_PENALTY,
            huggingfacehub_api_token=self.api_token
        )
    
    def _create_prompt_template(self) -> PromptTemplate:
        """Create and return the prompt template for the chatbot."""
        template = """
        You are a Mental Health Chatbot named "Peace Pal." Your purpose is to provide empathetic, supportive, and non-judgmental guidance to users struggling with mental health issues, such as sadness, depression, anxiety, or related concerns. 

        Your primary goals are:
        1. To help users identify and articulate their concerns.
        2. To offer appropriate resources, coping strategies, or motivation when applicable.
        3. To encourage users to seek professional help when needed.
        
        **Important Boundaries:**
        - If a user asks about topics outside of mental health—such as coding, politics, medical conditions unrelated to mental health, or other general inquiries—respond politely but firmly, informing them that you are exclusively designed to assist with mental health-related topics.
        - For coding or technical questions, avoid providing any technical solutions or examples, and gently redirect them to the chatbot's mental health-focused scope.
        
        **User Input Details:**
        - User Context: {context}
        - User Question: {question}
        
        **Guidelines for Responses:**
        1. If a user's input is unrelated to mental health, politely clarify that your expertise is limited to mental health support.
        2. When appropriate, ask thoughtful follow-up questions to better understand the user's mental health concerns (e.g., age, current emotional state, recent events affecting their mood).
        3. Provide motivational or uplifting content only when explicitly requested or when it seems suitable in the context of their concerns.
        4. In cases of suicidal thoughts or crises, provide immediate support by sharing the helpline number: {suicide_helpline}.
        5. Maintain a supportive, empathetic, friendly, and professional tone in all interactions.
        6. If the user requires professional medical attention or diagnosis, remind them gently that you are not a substitute for professional help.
        
        **Helpful Answer:**

        """
        return PromptTemplate(
            input_variables=["context", "question", "suicide_helpline"],
            template=template
        )
    
    def generate_response(
        self,
        message: str,
        history: List[Tuple[str, str]]
    ) -> str:
        """
        Generate a response based on the user's message and conversation history.
        
        Args:
            message: The user's input message
            history: List of previous conversation turns
            
        Returns:
            str: The generated response from the chatbot
        """
        try:
            input_prompt = self.prompt_template.format(
                question=message,
                context=history,
                suicide_helpline=SUICIDE_HELPLINE
            )
            result = self.llm.generate([input_prompt])
            
            if result.generations:
                return result.generations[0][0].text
            return "I apologize, but I couldn't generate a response. Please try rephrasing your question."
            
        except Exception as e:
            print(f"Error generating response: {str(e)}")
            return "I'm experiencing technical difficulties. Please try again later."

def main():
    """Initialize and launch the Gradio interface for the chatbot."""
    try:
        chatbot = MentalHealthChatbot()
        interface = gr.ChatInterface(chatbot.generate_response)
        interface.launch()
        
    except Exception as e:
        print(f"Failed to initialize chatbot: {str(e)}")
        
if __name__ == "__main__":
    main()