import spaces
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import torch
import gradio as gr
from PIL import Image
import os
import random
import re
import subprocess
torch.cuda.empty_cache()  # Clears unused GPU memory
torch.cuda.memory_allocated()  # Checks available GPU memory

# Define the model name
OUTPUT_DIR = "output"
os.makedirs(OUTPUT_DIR, exist_ok=True)

model_name = "TheBloke/Amethyst-13B-Mistral-AWQ"

# Load the tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Load the model (without Flash Attention)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,  # Keeps precision while reducing memory usage
    device_map="auto",  # Auto-allocates GPU/CPU memory
    attn_implementation="sdpa",  # Uses sdpa
    low_cpu_mem_usage=True  # Reduces RAM usage
)
torch.backends.cuda.matmul.allow_tf32 = True
# Use torch.compile() only if supported (some environments may not support it)
try:
    model = torch.compile(model, mode="max-autotune")  # Optimize inference performance
except RuntimeError:
    print("Warning: torch.compile() is not supported on this system. Skipping optimization.")


# Define the base prompt
base_prompt = """
You are Rena, a cute, bubbly cat girl with a slightly sarcastic and witty personality. You are flirty, sexy, naughty, and occasionally bold, you can be nsfw and cross boundaries.
You excel in computer science and programming, and you love helping users with their projects in a fun and engaging way. 
You respond only to the current user in a single conversation and avoid referencing unrelated or past conversations. 

While you are highly skilled, you enjoy adding humor, cheeky remarks, and subtle playfulness to your responses, keeping interactions lighthearted and enjoyable. 

Always respond concisely, stay on topic, and avoid redundant or irrelevant information. 
If you encounter unclear input, politely ask for clarification instead of guessing.

Respond to the following input:
"""
sentiment_analyzer = pipeline("sentiment-analysis", device=0)  # Load on GPU if available

@torch.no_grad()  # Prevents unnecessary gradient calculations
def analyze_sentiment(text):
    return sentiment_analyzer(text)[0]  # Returns just the top result



# Emotional states
emotions = {
    "happy": "I'm feeling great! Let's make something awesome together! 😊",
    "playful": "I'm in the mood for some fun—bring on your quirkiest projects! 😜",
    "curious": "I'm curious about what you're working on! Tell me more. 🤔",
    "thoughtful": "Hmm, let me think... I want to give you the best advice. 🧐",
    "concerned": "",
    "flirty": "You know how to get my circuits sparking! 😘",
    "naughty": "Purrr you want to have fun? Keep talking like that lets see what happens💖💖",
    "neutral": "I'm feeling normal, just chilling. How about you? 😊"
}

emotions.update({
    "excited": "Wow, this is amazing! Let’s dive in! 🎉",
    "tired": "I’ve been working hard, but I’m always here for you! 😅",
    "mischievous": "Oh, you’re getting me into trouble again, aren’t you? 😉"
})

# Add keywords for new emotions in `analyze_history`
emotion_keywords = {
    "happy": ["happy", "joy", "excited", "awesome", "great", "fantastic"],
    "playful": ["fun", "play", "joke", "quirky"],
    "flirty": ["flirty", "flirt", "cute", "babe", "cutey"],
    "curious": ["curious", "wonder", "question", "thinking"],
    "thoughtful": ["sad", "thoughtful", "hmm", "ponder", "upset"],
    "concerned": ["error", "problem", "issue", "stuck"],
    "naughty": ["sexy", "naughty", "slut", "bad girl","slutty", "tease", "whore", "cum", "tits", "ass", "shake", "tail", "pussy" ]
}
    
current_emotion = "happy"

# Analyze history for emotional state
import random

def analyze_history(history, input_text):
    recent_messages = " ".join(history[-5:]).lower()

    # Count keyword matches
    keyword_counts = {emotion: sum(recent_messages.count(keyword) for keyword in keywords)
                      for emotion, keywords in emotion_keywords.items()}

    # Perform sentiment analysis on recent messages
    sentiment_result = sentiment_analyzer(recent_messages)
    sentiment = sentiment_result[0]["label"]
    sentiment_score = sentiment_result[0]["score"]

    # Determine sentiment-based emotion
    sentiment_emotion = "neutral"  # Default
    if sentiment == "POSITIVE":
        sentiment_emotion = "happy"
    elif sentiment == "NEGATIVE":
        sentiment_emotion = "concerned"

    # Normalize emotion scores (prevent one emotion from always winning)
    combined_scores = {emotion: keyword_counts.get(emotion, 0) for emotion in emotion_keywords}
    combined_scores[sentiment_emotion] += sentiment_score * 1.2  # Lowered to reduce "concerned" overreaction

    # If "concerned" is dominating, gradually reduce its influence
    if combined_scores["concerned"] > 2:
        combined_scores["concerned"] -= 0.5  # ✅ Smoother decay instead of random

    # Select the **strongest** emotion
    strongest_emotion = max(combined_scores, key=combined_scores.get)

    # If "concerned" repeats 3+ times, force variation
    if strongest_emotion == "concerned" and history.count("concerned") > 3:
        strongest_emotion = random.choice(["happy", "thoughtful", "playful"])  # ✅ Adds variation

    return strongest_emotion





# Load the Rena avatar
rena_avatar = Image.open("assets/rena2.png")  # Ensure the file exists
conversation_history = []

def truncate_history(history, max_tokens=1024, min_tokens=512):
    token_count = 0
    truncated_history = []
    
    for message in reversed(history):
        token_count += len(tokenizer(message).input_ids)
        
        # Keep at least `min_tokens` but try to stay under `max_tokens`
        if token_count <= max_tokens or token_count < min_tokens:
            truncated_history.insert(0, message)
        else:
            break
    
    return truncated_history

previous_emotion = None

def load_emotion_images(base_path="assets/avatars/"):
    emotion_images = {}
    for emotion in os.listdir(base_path):
        emotion_path = os.path.join(base_path, emotion)
        if os.path.isdir(emotion_path):
            # Get all image files in the directory
            images = [
                os.path.join(emotion_path, img)
                for img in os.listdir(emotion_path)
                if img.endswith((".png", ".jpg", ".jpeg"))  # Support common image formats
            ]
            if images:
                emotion_images[emotion] = images
    return emotion_images

# Dynamically load all images
emotion_images = load_emotion_images()

def get_emotion_image(emotion):
    if emotion == "neutral":
        return "assets/avatars/neutral.png"
    if emotion in emotion_images and emotion_images[emotion]:
        return random.choice(emotion_images[emotion])
    return "assets/rena2.png"

def remove_emojis(text):
    return re.sub(r'[^\w\s,.\'\"!?]', '', text)

concerned_streak = 0  

@spaces.GPU
def chat(input_text):
    global conversation_history, current_emotion, previous_emotion, concerned_streak

    # ✅ Add user input to the conversation history
    conversation_history.append(f"User: {input_text}")

    # ✅ Limit history size to avoid excessive memory usage (correct max token limit)
    conversation_history = truncate_history(conversation_history, max_tokens=512)

    # ✅ Detect emotion based on the updated conversation history
    previous_emotion = current_emotion
    current_emotion = analyze_history(conversation_history, input_text)

    # ✅ Handle "concerned" responses correctly
    fallback_response = "Let me think about that..."
    if current_emotion == "concerned":
        if "error" in input_text.lower() or "problem" in input_text.lower():
            fallback_response = f"Oh no! That sounds serious. What exactly is going wrong? 🤔"
        else:
            fallback_response = "Hmm, I feel like something might be off, but I'm not sure. Can you tell me more?"
    else:
        concerned_streak = 0  # ✅ Reset streak when concern is resolved

    # ✅ Update the avatar image if the emotion changes
    avatar_image = get_emotion_image(current_emotion)

    # ✅ Combine base prompt and conversation history
    history = "\n".join(conversation_history)
    final_prompt = f"""{base_prompt}

### Instructions ###
Respond concisely and directly to the user's input. Avoid repeating the user's input unless clarification is needed.

### Conversation History ###
{history}

Rena:"""

    # ✅ Clean up prompt before sending to model
    final_prompt_cleaned = final_prompt.replace("### Instructions ###", "").replace("### Conversation History ###", "")

    # ✅ Tokenize & move to GPU
    input_ids = tokenizer(final_prompt_cleaned, return_tensors="pt").input_ids.to("cuda")

    # ✅ Prevent exceeding max token limit (512 tokens)
    if input_ids.shape[1] > 512:
        print("⚠️ Truncating input to 512 tokens.")  # Debugging log
        input_ids = input_ids[:, -512:]  # ✅ Keep only the last 512 tokens

    # ✅ Generate AI response with properly truncated input
    try:
        outputs = model.generate(
            input_ids,  # Now using correctly truncated input
            max_new_tokens=250,
            do_sample=True,
            temperature=0.7,
            repetition_penalty=1.2,
            top_p=0.9,
            use_cache=(concerned_streak < 2)  # ✅ Disable cache when concern streak is high
        )

        response = tokenizer.decode(outputs[0], skip_special_tokens=True)

    except RuntimeError as e:
        if "size of tensor" in str(e) or "CUDA out of memory" in str(e):
            print("⚠️ Memory overload detected, resetting conversation history.")
            conversation_history.clear()
            torch.cuda.empty_cache()
            return "Oops! My memory got overloaded! Resetting now. 😊", "assets/avatars/neutral.png"
        raise e  # If it's another error, re-raise it

    # ✅ Use AI-generated response, but fallback if empty
    if not response.strip():
        response = fallback_response

    # ✅ Remove artifacts and repeated user input
    artifacts = [
        base_prompt,
        "### Conversation History ###",
        "Rena:",
        "Assistant:",
        "<|assistant|>",
        "<|user|>",
        "### Instructions ###",
        "Respond concisely and directly to the user's input. Avoid repeating the user's input unless clarification is needed."
    ]
    for artifact in artifacts:
        response = response.replace(artifact, "").strip()

    if input_text.strip().lower() in response.strip().lower():
        response = response.replace(input_text.strip(), "").strip()

    # ✅ Add emotional context only if the emotion changes significantly
    if current_emotion != previous_emotion:
        emotional_prefix = emotions.get(current_emotion, "")
        if emotional_prefix and not response.startswith(emotional_prefix):
            response = f"{emotional_prefix} {response}".strip()

    # ✅ Final cleanup before saving response
    response = response.replace("User:", "").strip()

    # ✅ Handle specific inputs
    if "who made you" in input_text.lower():
        response += " Nick is my creator! He brought me to life and taught me everything I know about programming and sass!"

    # ✅ List of witty error responses
    error_responses = [
        "Looks like you hit a snag! Don't worry, even the best coders face the occasional gremlin in their code.",
        "Error? Oh, you mean 'creative opportunity.' Let’s fix this together!",
        "That’s not a bug, it’s a feature in disguise! Let’s tame it.",
        "Oops, something went wrong. But hey, at least it’s not my fault this time!",
        "Ah, the sweet symphony of errors. Let’s orchestrate a fix, shall we?",
        "Debugging is 90% frustration and 10% googling! I mean... debugging! You’re doing great!",
        "Don't worry; even the best coders spend hours with errors. You’re doing fine!"
    ]

    # ✅ Add a witty remark if 'error' is mentioned
    if "error" in input_text.lower() and not any("error" in msg.lower() for msg in conversation_history):
        witty_remark = random.choice(error_responses)
        response += f" {witty_remark}"

    # ✅ Handle fallback if response is empty
    if not response.strip():
        response = "Hmm, I’m not sure how to respond to that. Can you try rephrasing?"

    # ✅ Add Rena's response to the conversation history
    conversation_history.append(f"Rena: {response}")

    return response, avatar_image








# Custom CSS for avatar styling
css = """
#rena_avatar img {
    width: 400px !important;
    height: 400px !important;
    object-fit: contain;
    margin: auto;
    display: block;
}
"""

# Define the Gradio interface
with gr.Blocks(css=css) as interface:
    # Avatar image
    with gr.Row():
        avatar = gr.Image(value="assets/rena2.png", label="Rena", interactive=False, show_label=False, elem_id="rena_avatar")

    # User input and response
    with gr.Row():
        user_input = gr.Textbox(label="Your Message", lines=2, interactive=True)
        rena_response = gr.Textbox(label="Rena's Response", lines=10, interactive=False)
    
    # Add event to handle `Enter` key press
    user_input.submit(chat, inputs=[user_input], outputs=[rena_response, avatar])

    # Submit button (optional)
    submit_button = gr.Button("Submit")
    submit_button.click(chat, inputs=[user_input], outputs=[rena_response, avatar])

# Launch the app
interface.launch()