import os
import re
from dotenv import load_dotenv
import torch
from transformers import RobertaForSequenceClassification, RobertaTokenizerFast, pipeline as text_pipeline
import gradio as gr
from openai import OpenAI

# Load environment variables from .env file
load_dotenv()

# Get API key from environment
API_KEY = os.getenv("API_KEY")

# Initialize OpenAI client
client = OpenAI(
    base_url="https://integrate.api.nvidia.com/v1",
    api_key=API_KEY
)

# Load classification model
def load_emotion_model(model_path):
    model = RobertaForSequenceClassification.from_pretrained(model_path)
    tokenizer = RobertaTokenizerFast.from_pretrained(model_path)
    return model, tokenizer

# Map prediction to readable labels
def map_to_labels(label):
    return "Happy/Positive Mindset" if label.lower() == "positive" else "Depressed/Negative Mindset"

# Classify mental state based on user input
def classify_emotion(user_input, model, tokenizer, device):
    nlp = text_pipeline("text-classification", model=model, tokenizer=tokenizer, device=device)
    result = nlp(user_input)
    return map_to_labels(result[0]['label'])

# Analyze emotion using the LLM
def emotion_analysis(user_input):

    # # Validate input
    # if not user_input.strip():  # Check for empty or blank input
    #     progress_callback("Please provide valid input before submitting.", False)
    #     return "No input provided.", ""

    # Load model
    model_path = "mentalhealth-roberta-base_nemotron_model"  # Replace with your model path
    model, tokenizer = load_emotion_model(model_path)
    device = 0 if torch.cuda.is_available() else -1

    # Step 1: Classify emotion
    predicted_emotion = classify_emotion(user_input, model, tokenizer, device)

    # Step 2: Generate LLM response
    prompt = f"""
    Task: You are a social psychologist specializing in Roy Baumeister's six-stage theory of emotional progression. Your task is to analyze emotional states based on user input while adhering strictly to specific response boundaries.

    [Input Information]:
    **User Input**: "{user_input}"
    **Model Output**: "{predicted_emotion}"

    Specifics:
    1. Respond **only** to questions or input related to mental health or emotional well-being. For unrelated input, reply strictly with: "Not a valid question."
    2. Use the **User Input** as the primary source for determining the emotional state, while considering the **Model Output** ("happy" or "depressed") as a secondary reference.
    3. Categorize the user’s emotional state into one of Roy Baumeister’s six stages of emotional progression:
       - Stage 1: Falling short of expectations
       - Stage 2: Attributions to self
       - Stage 3: High self-awareness
       - Stage 4: Negative affect
       - Stage 5: Cognitive deconstruction
       - Stage 6: Disinhibition
    4. Provide specific stage-based recommendations:
       - If **depressed**, suggest remedies for improving emotional state.
       - If **happy**, suggest strategies for maintaining or enhancing happiness.
    5. Maintain clarity, empathy, and practicality in all analyses and suggestions.

    [Response Rules]:
    - If input is unrelated to mental health, reply only with: "Not a valid question."
    - Always analyze user input independently, even if it conflicts with the model's predicted output.

    [Desired Output Format for valid input]:
    Emotional Analysis:
    I'd say you're feeling: <Happy/Depressed>
    Emotional Stage: <Stage and brief reasoning>
    Suggested Remedies/Strategies: <Practical advice based on the stage>

    For invalid input:
    "Not a valid question."
    """

    try:
        completion = client.chat.completions.create(
            model="nvidia/nemotron-4-340b-instruct",
            messages=[{"role": "user", "content": prompt}],
            temperature=0.5,
            top_p=0.7,
            max_tokens=512,
            stream=True
        )

        # Iterate over the streaming response
        response = ""
        for chunk in completion:
            if chunk.choices[0].delta.content is not None:
                print(chunk.choices[0].delta.content, end="")
                # response = chunk.choices[0].delta.content
                response_chunk = chunk.choices[0].delta.content
                response += response_chunk
            else:
                print(f"Unexpected chunk format: {chunk}")

    except Exception as e:
        response = f"An error occurred while processing the response: {e}"
    response= str(response).replace("*", '')
    return response

def extract_analysis_details(analysis_text):
    feelings_match = re.search(r"I'd say you're feeling:\s*([^\n]+)", analysis_text)
    feelings = feelings_match.group(1).strip() if feelings_match else "Not Found"
    print(feelings)
    if feelings.lower() == "happy":
        feelings = feelings + " with Positive Mindset"
    elif feelings.lower() == "depressed":
        feelings = feelings + " with Negative Mindset"
    else:
        feelings

    # Extract emotional stage
    # stage_match = re.search(r"Emotional Stage:\s*([^\n.]+)", analysis_text)
    # print(stage_match)
    stage_match = re.search(r"Emotional Stage:\s*(.*?)(?=\n[A-Z])", analysis_text, re.DOTALL)
    emotional_stage = stage_match.group(1).strip() if stage_match else "Not Found"

    # Regex to match the section header and capture from there to the end
    pattern = r"(Suggested Remedies|Suggested Remedies/Strategies|Suggested Strategies):.*"
    match = re.search(pattern, analysis_text, re.DOTALL)
    suggestions = match.group(0).strip() if match else "No matching section found."
    # print(suggestions)

    if feelings == "Not Found" or feelings == "Not a valid question.":
        feelings = "Not a valid question."
    return feelings, emotional_stage, suggestions


# Gradio interface with input validation
def validate_and_run(user_input):
    if not user_input.strip():  # Check if the input is empty or just spaces
        return "Please provide valid input before submitting.", "Not Applicable", "Not Applicable"
    else:
        response = emotion_analysis(user_input)
        return extract_analysis_details(response)


# Gradio interface
iface = gr.Interface(
    fn=validate_and_run,
    inputs=gr.Textbox(#lines=2,
        label="How are you feeling today?",
        placeholder="Share your thoughts here...!"),
    outputs=[
        # gr.Textbox(label="Analysing Your State of Mind...."),
        # gr.Textbox(label="Providing Best Strategies")
        # gr.Textbox(label="Original"),
        gr.Textbox(label="Feelings"),
        gr.Textbox(label="Emotional Stage"),
        gr.Textbox(label="Providing Best Strategies")
    ],
    # live=True,
    title="Analyze your emotions and generate stage-specific psychological insights\n",
    # title = "Emotion Analysis and Dynamic Response Generator"
    # description="Analyze your emotions and receive dynamic psychological insights."
)

# Launch the app
if __name__ == "__main__":
    iface.launch()