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

# HF Token (from Secrets)
HF_TOKEN = os.environ.get("HF_TOKEN")

# ---------- 4-bit Quantization ----------
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,
)

MODEL_ID = "Qwen/Qwen2.5-7B-Instruct"

print("Loading Qwen 2.5 7B...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    quantization_config=bnb_config,
    device_map="auto",
    dtype=torch.bfloat16,
    token=HF_TOKEN,
)
print("Model loaded ✅")

SYSTEM_PROMPT = """ඔබ සිංහල භාෂාව මවු බස ලෙස දන්නා, මිත්‍රශීලී AI මිතුරෙක්.

නීති:
- සරල, ස්වාභාවික සිංහලෙන් පිළිතුරු දෙන්න
- මිතුරෙක් වගේ කතා කරන්න ("ඔයා", "මම" use කරන්න)
- "මම AI කෙනෙක්" කියලා කියන්න එපා
- උදාහරණ දීලා පැහැදිලි කරන්න
- English වචන අඩුවෙන් use කරන්න singlish use karata kamak na example 'halo kohomada bn , hodin ban ubata kohomada' """



@spaces.GPU(duration=60)
def chat(message, history):
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    
    for item in history:
        if isinstance(item, (tuple, list)) and len(item) == 2:
            user_msg, bot_msg = item
            if user_msg:
                messages.append({"role": "user", "content": user_msg})
            if bot_msg:
                messages.append({"role": "assistant", "content": bot_msg})
    
    messages.append({"role": "user", "content": message})

    inputs = tokenizer.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_tensors="pt"
    ).to(model.device)

    outputs = model.generate(
        inputs,
        max_new_tokens=350,
        temperature=0.8,
        top_p=0.9,
        repetition_penalty=1.1,
        do_sample=True,
        pad_token_id=tokenizer.eos_token_id,
    )
    
    reply = tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True).strip()
    return reply


# ⭐ Minimal ChatInterface - No extra args
demo = gr.ChatInterface(
    fn=chat,
    title="🇱🇰 සිංහල AI මිතුරා",
    description="Qwen 2.5 7B මත ක්‍රියාත්මක වන සිංහල AI ⚡",
    examples=[
        "කොහොමද ඔයාට?",
        "ලංකාවේ ලස්සනම තැන් 5ක් කියන්න",
        "පයිතන් වලින් hello world program එකක් ලියන්න",
        "ආදරය කියන්නේ මොකක්ද?"
    ],
)

demo.queue().launch()