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

# --------------------------------------------------
# Model configuration
# --------------------------------------------------

REPO = "vermarjun/chad"
SUBFOLDER = "transformers/dpo450"

print("Loading Chad tokenizer...")

tokenizer = AutoTokenizer.from_pretrained(
    REPO,
    subfolder=SUBFOLDER,
)

print("Loading Chad model...")

model = AutoModelForCausalLM.from_pretrained(
    REPO,
    subfolder=SUBFOLDER,
    torch_dtype=torch.float32,
)

model.eval()

print("Chad model loaded successfully!")


# --------------------------------------------------
# Chat response function
# --------------------------------------------------

@spaces.GPU
def respond(message, history):
    if not message or not message.strip():
        return "Bhai, kuch message toh likho!"

    # Build conversation prompt
    prompt = ""

    for item in history or []:
        if not isinstance(item, dict):
            continue

        role = item.get("role")
        content = item.get("content", "")

        # Ignore non-text messages
        if not isinstance(content, str):
            continue

        if role == "user":
            prompt += f"<|user|>{content}<|assistant|>"
        elif role == "assistant":
            prompt += f"{content}<|endoftext|>"

    prompt += f"<|user|>{message.strip()}<|assistant|>"

    # Tokenize input
    inputs = tokenizer(
        prompt,
        return_tensors="pt",
        add_special_tokens=False,
    )

    # Generate response on the GPU allocated by ZeroGPU
    with torch.no_grad():
        output = model.generate(
            **inputs,
            max_new_tokens=120,
            do_sample=True,
            temperature=0.8,
            top_p=0.9,
            repetition_penalty=1.3,
            eos_token_id=0,
        )

    # Decode only newly generated tokens
    input_length = inputs["input_ids"].shape[1]
    new_tokens = output[0][input_length:]

    answer = tokenizer.decode(
        new_tokens,
        skip_special_tokens=True,
    ).strip()

    return answer or "Bhai, ek baar phir try karo."


# --------------------------------------------------
# Gradio interface
# --------------------------------------------------

demo = gr.ChatInterface(
    fn=respond,
    title="⚡ Chad AI",
    description="A Hinglish conversational AI chatbot",
    textbox=gr.Textbox(
        placeholder="Bol bhai, kya scene hai?",
        container=False,
    ),
)

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