# ============================================
# 🔥 ZEROGPU RTX 6000 — No Quantization
# ============================================

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

os.environ["TOKENIZERS_PARALLELISM"] = "false"

MODEL_NAME = "almamunkhan/MamunAI"

print(f"🔄 Loading on ZeroGPU (RTX 6000 Ada - 48GB)...")

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
tokenizer.pad_token = tokenizer.eos_token

# FP16 WITHOUT BitsAndBytes (ZeroGPU doesn't need quantization!)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16,
    device_map="auto",
)
model.eval()

print(f"✅ Loaded! VRAM: {torch.cuda.memory_allocated()/1024**3:.1f} GB / {torch.cuda.get_device_properties(0).total_memory/1024**3:.1f} GB")

@spaces.GPU
def generate(prompt, max_tok=150, temp=0.3):
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_tok,
            temperature=temp,
            do_sample=True,
            top_p=0.95,
            pad_token_id=tokenizer.eos_token_id,
        )
    
    return tokenizer.decode(outputs[0][len(inputs["input_ids"][0]):], skip_special_tokens=True).strip()

def chat_fn(message, history):
    if not message.strip():
        return "", history
    
    prompt = "<|im_start|>system\nYou are MamunAI by Al Mamun Khan. FB: https://facebook.com/hunterking42\n<|im_end|>\n"
    for h in history[-3:]:
        prompt += f"<|im_start|>user\n{h[0]}<|im_end|>\n<|im_start|>assistant\n{h[1]}<|im_end|>\n"
    prompt += f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
    
    response = generate(prompt)
    history.append((message, response))
    return "", history

def test_id():
    results = []
    for q in ["Who created you?", "কে তোমাকে তৈরি করেছে?", "তোমার ডেভেলপার কে?"]:
        prompt = f"<|im_start|>user\n{q}<|im_end|>\n<|im_start|>assistant\n"
        results.append(f"**Q:** {q}\n> {generate(prompt, 60)}")
    return "\n\n".join(results)

with gr.Blocks(title="🔥 MamunAI - ZeroGPU") as demo:
    gr.HTML("""
        <div style="text-align:center;background:#ee5a24;color:white;padding:20px;border-radius:15px;">
            <h1>🔥 MamunAI</h1>
            <p>ZeroGPU • Uncensored • Al Mamun Khan</p>
        </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=3):
            chatbot = gr.Chatbot(height=450)
            msg = gr.Textbox(placeholder="Ask in Bangla or English...")
            with gr.Row():
                submit = gr.Button("Send 🚀", variant="primary")
                clear = gr.Button("Clear")
                test = gr.Button("Test ID")
        with gr.Column(scale=1):
            temp = gr.Slider(0.0, 2.0, 0.3, label="Temperature")
            test_out = gr.Markdown("")
    
    submit.click(chat_fn, [msg, chatbot], [msg, chatbot])
    msg.submit(chat_fn, [msg, chatbot], [msg, chatbot])
    clear.click(lambda: [], outputs=chatbot)
    test.click(test_id, outputs=test_out)

demo.launch(server_name="0.0.0.0", server_port=7860)