import gradio as gr
import torch
import os
import re
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import PeftModel

# --- 1. Модельді Жүктеу Параметрлері ---
TUNED_MODEL_PATH = "./qwen_kz_books_tuned" 
BASE_MODEL_ID = "Qwen/Qwen1.5-0.5B-Chat" 

MODEL_KZ_NAME = "Қазақ Кітаптары AI (Оқытылған)"
MODEL_BASE_NAME = "Qwen 1.5-0.5B Chat (Базалық)"

# --- 2. Модельді Жүктеу Функциясы ---
def load_model():
    """Оқытылған немесе базалық Qwen 0.5B моделін CPU-да жүктейді"""
    print("🔥 Модельді жүктеу басталды. CPU режимі...")
    
    try:
        if os.path.exists(TUNED_MODEL_PATH):
            print(f"✅ Оқытылған модель табылды: {TUNED_MODEL_PATH}")
            tokenizer = AutoTokenizer.from_pretrained(TUNED_MODEL_PATH)
            
            base_model = AutoModelForCausalLM.from_pretrained(
                BASE_MODEL_ID,
                torch_dtype=torch.float32, 
                device_map="cpu", 
                trust_remote_code=True
            )
            model = PeftModel.from_pretrained(base_model, TUNED_MODEL_PATH)
            model_name = MODEL_KZ_NAME
        else:
            print(f"⚠️ Оқытылған модель табылмады. Базалық {BASE_MODEL_ID} жүктелуде.")
            tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID)
            model = AutoModelForCausalLM.from_pretrained(
                BASE_MODEL_ID,
                torch_dtype=torch.float32,
                device_map="cpu",
                trust_remote_code=True
            )
            model_name = MODEL_BASE_NAME
            
    except Exception as e:
        print(f"❌ Модельді жүктеу кезінде қате: {e}")
        return None, None, f"Қате: {e}"

    print(f"✅ Модель сәтті жүктелді: {model_name}")
    return model, tokenizer, model_name

# Глобальды айнымалыларды жүктеу
model, tokenizer, model_status = load_model()

# --- 3. Генерация Процесі (AI Қозғалтқышы) ---
def generate_content(prompt):
    """Модельге сұрақ жібереді және нәтижесін алады."""
    if model is None:
        return "Модель жүктелмеді. Қателерді консольден тексеріңіз.", "Қате"

    # Prompt-ты Qwen-нің чат үлгісіне бейімдеу
    formatted_prompt = f"Сіз қазақ әдебиеті бойынша білімді чатботсыз. Жауапты толық және Қазақша беріңіз. Сұрақ: {prompt}"
    
    # Qwen моделінің чат үлгісін қолдану
    messages = [
        {"role": "user", "content": formatted_prompt}
    ]
    input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    
    # Токенизация
    input_ids = tokenizer(input_text, return_tensors="pt", padding=True).to(model.device)
    
    # Генерациялау параметрлері
    outputs = model.generate(
        **input_ids,
        max_new_tokens=512,
        temperature=0.7,
        do_sample=True,
        top_k=50,
        top_p=0.95,
        pad_token_id=tokenizer.eos_token_id
    )
    
    # Нәтижені декодтау
    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # Qwen үлгісіндегі prompt-ты кесіп тастау
    response_start = generated_text.rfind(messages[-1]['content']) + len(messages[-1]['content'])
    generated_text = generated_text[response_start:].strip()

    # Таза мәтінді қайтару
    result_text = f"## 📚 Қазақша Жауап\n\n{generated_text}"
    
    return result_text, f"Генерация сәтті аяқталды. Қолданылған модель: {model_status}"

# --- 4. Gradio Интерфейсі ---
with gr.Blocks(title="Қазақ Кітаптары Чатботы (Qwen AI)") as demo:
    gr.HTML(f"""
        <h1 style='text-align: center; color: #8B4513;'>📖 ҚАЗАҚ КІТАПТАРЫ ТУРАЛЫ ЧАТБОТ 🇰🇿</h1>
        <p style='text-align: center;'>Бұл AI тек сен оқытқан кітаптар бойынша (Абай жолы, Көшпенділер, Қан мен Тер, т.б.) әңгімелесе алады.</p>
        <p style='text-align: center; font-weight: bold;'>Модель Статусы: <span id='model-status'>{model_status}</span></p>
    """)

    with gr.Row():
        with gr.Column(scale=1):
            
            input_prompt = gr.Textbox(
                label="AI-мен әңгімелесуді бастаңыз",
                placeholder="Мысалы: 'Абай жолындағы басты кейіпкерлерді атаңыз' немесе 'Мұхтар Әуезовтің бұл романының тарихи маңызы қандай?'",
                lines=5
            )
            
            generate_btn = gr.Button("💬 Сұрақ Жіберу")
            
        with gr.Column(scale=2):
            output_text = gr.Markdown(
                label="AI Жауабы"
            )
            
            status_message = gr.Textbox(
                label="Процесс Статусы",
                interactive=False
            )

    generate_btn.click(
        fn=generate_content,
        inputs=[input_prompt],
        outputs=[output_text, status_message]
    )

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