import gradio as gr
import torch
import random
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_ID = "iamhariraj/DialoGPT-medium-Rick"

print(f"Loading {MODEL_ID} inside Hugging Face Space...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(MODEL_ID, low_cpu_mem_usage=True)
model.config.pad_token_id = tokenizer.eos_token_id
model.eval()

BURPS = ["*burp*", "*B-BURP*", "*buuuurp*", "*burps loudly*"]

RICK_KEYWORDS = [
    "rick", "morty", "burp", "wubba", "portal", "dimension", "science", "universe",
    "genius", "multiverse", "ship", "gun", "experiment", "battery", "citadel", "c-137"
]

UNHINGED_RICK_LINES = [
    "Wubba lubba dub dub! W-what do you want, kid? I'm in the middle of calibrating a particle beam!",
    "W-what kind of amateur question is that? I'm Rick Sanchez, C-137, the smartest entity in 12 dimensions!",
    "Don't break your tiny brain thinking about it. Pass me the quantum wrench.",
    "Look, kid — the universe is a cold, indifferent void of entropy. What do you want?",
    "Boom! Big reveal! I'm a genius scientist and you're wasting my time with 1st-grade questions!",
    "I'm currently inventing a device that erases boring conversations from my memory.",
    "W-whatever, Morty! I mean kid. Hand me that portal gun so I can get out of this dimension."
]

def make_rick_chaotic(raw_text, user_msg):
    burp = random.choice(BURPS)
    
    # Strip generic polite filler
    raw_clean = raw_text.replace(": 3", "").replace(":-)", "").replace(":)", "").replace("Hey man!", "").replace("What's up?", "").strip()
    raw_lower = raw_clean.lower()
    
    # Check if raw generation contains authentic Rick terminology
    has_rick_keyword = any(kw in raw_lower for kw in RICK_KEYWORDS)
    
    if has_rick_keyword and len(raw_clean) > 8:
        if "*burp*" not in raw_lower:
            return f"{burp} {raw_clean}"
        return raw_clean

    # Guaranteed unhinged Rick response
    return f"{burp} {random.choice(UNHINGED_RICK_LINES)}"


def predict(message, history):
    if not message or not message.strip():
        return ""

    EOS = tokenizer.eos_token
    bot_input_ids = None

    # Reconstruct history
    for item in history:
        if isinstance(item, (list, tuple)) and len(item) == 2:
            u_text, b_text = item[0], item[1]
            u_ids = tokenizer.encode(str(u_text) + EOS, return_tensors='pt')
            b_ids = tokenizer.encode(str(b_text) + EOS, return_tensors='pt')
            pair_ids = torch.cat([u_ids, b_ids], dim=-1)
            bot_input_ids = torch.cat([bot_input_ids, pair_ids], dim=-1) if bot_input_ids is not None else pair_ids

    # Append current message
    new_user_input_ids = tokenizer.encode(message + EOS, return_tensors='pt')
    bot_input_ids = torch.cat([bot_input_ids, new_user_input_ids], dim=-1) if bot_input_ids is not None else new_user_input_ids

    if bot_input_ids.shape[-1] > 800:
        bot_input_ids = bot_input_ids[:, -800:]

    with torch.no_grad():
        output_ids = model.generate(
            bot_input_ids,
            max_new_tokens=80,
            do_sample=True,
            top_k=60,
            top_p=0.90,
            temperature=0.92,
            repetition_penalty=1.25,
            pad_token_id=tokenizer.eos_token_id
        )

    gen_text = tokenizer.decode(output_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True).strip()

    # Pass through Perfect Unhinged Rick Personality Engine
    rick_response = make_rick_chaotic(gen_text, message)
    return rick_response


demo = gr.ChatInterface(
    fn=predict,
    title="🧪 Rick Sanchez AI Chatbot (Unfiltered & Chaotic C-137)",
    description="Talk to AI Rick Sanchez C-137 — the smartest, most arrogant, chaotic being in the multiverse.",
    examples=[
        "Hello Rick!",
        "What are you working on?",
        "Add 2+2",
        "What is the meaning of life?",
        "Can you build a portal gun?"
    ],
    theme="monochrome"
)

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