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

# Model yükleme
print("🔍 Model yükleniyor....")
base_model_name = "vngrs-ai/Kumru-2B-Base"
adapter_path = "ademireltaz/kumru-tr-poetry-lora"
# Tokenizer - use_fast=False ile yükle
tokenizer = AutoTokenizer.from_pretrained(
    base_model_name,
    use_fast=False,
    trust_remote_code=True
)

# Pad token ayarla
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(
    base_model_name,
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True,
    device_map="auto",
    trust_remote_code=True
)
model = PeftModel.from_pretrained(model, adapter_path)
model.eval()
print("✅ Model hazır!")

def chat_generate(message, history, temperature, max_length):
    """Chatbot formatında şiir üret"""
    
    # Şiir dışı istekleri kontrol et
    non_poetry_keywords = [
        'hesapla', 'çarp', 'böl', 'topla', 'çıkar', 'matematik',
        'kod yaz', 'program', 'script', 'function',
        'tarif', 'yemek', 'nasıl yapılır',
        'hava durumu', 'saat kaç', 'tarih',
        'çevir', 'translate', 'İngilizce',
        'anlat', 'açıkla', 'nedir', 'kimdir',
        'hikaye', 'makale', 'yazı yaz', 'kompozisyon',
        'özet', 'özetle', 'summarize'
    ]
    
    message_lower = message.lower()
    
    # Şiir anahtar kelimeleri
    poetry_keywords = ['şiir', 'mısra', 'beyit', 'kıta', 'dizeleri', 'nazım']
    
    # Şiir isteği içerip içermediğini kontrol et
    has_poetry_keyword = any(keyword in message_lower for keyword in poetry_keywords)
    has_non_poetry_keyword = any(keyword in message_lower for keyword in non_poetry_keywords)
    
    # Eğer açıkça şiir dışı bir şey isteniyorsa reddet
    if has_non_poetry_keyword and not has_poetry_keyword:
        return "Üzgünüm, ben sadece şiir üretmek için eğitilmiş bir modelim. Lütfen bana bir şiir konusu verin, sizin için şiir yazayım. 🌸"
    
    # Mesaj çok kısa ve genel ise (örn: "merhaba", "nasılsın")
    if len(message.split()) <= 3 and not has_poetry_keyword:
        # Basit sohbet tespiti
        greeting_words = ['merhaba', 'selam', 'hey', 'nasılsın', 'naber', 'günaydın', 'iyi akşamlar']
        if any(word in message_lower for word in greeting_words):
            return "Merhaba! Ben bir şiir üretme modeliyim. Bana bir konu verin, sizin için şiir yazayım. Örneğin: 'Aşk hakkında bir şiir yaz' veya 'Sonbahar temalı şiir' 🍂"
    
    # Şiir üretimi için prompt hazırla
    prompt = f"Görev: {message}\nYanıt:"
    
    inputs = tokenizer(
        prompt,
        return_tensors="pt",
        return_token_type_ids=False,
        padding=True
    ).to(model.device)
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=int(max_length),
            do_sample=True,
            temperature=float(temperature),
            top_p=0.92,
            repetition_penalty=1.1,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )
    
    full_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    if "Yanıt:" in full_text:
        poem = full_text.split("Yanıt:", 1)[1].strip()
    else:
        poem = full_text.replace(prompt, "").strip()
    
    # Üretilen içeriğin şiir olup olmadığını basit kontrol
    if len(poem.split()) < 10:  # Çok kısa ise
        return "Üzgünüm, sadece şiir üretebilirim. Lütfen bana bir şiir konusu belirtin. 📝"
    
    return poem

# Chatbot arayüzü - examples formatı düzeltildi
demo = gr.ChatInterface(
    fn=chat_generate,
    title="🎭 Kumru ile Taz Şiir Botu",
    description="""
    Türkçe şiir üretmek için tasarlanmış bir chat bot.
    
    **Nasıl kullanılır:**
    - Bir şair tarzı ve tema belirtin
    - Örnek: "Nazım Hikmet tarzında, deniz üzerine bir şiir yaz"
    
    **Desteklenen şairler:**
    - Nazım Hikmet
    - Aşık Veysel
    - Orhan Veli
    - Fazıl Hüsnü Dağlarca
    - Cahit Sıtkı Tarancı gibi üstadlar
    """,
    # ✅ Examples düzeltildi: [mesaj, temperature, max_length]
    examples=[
        ["Nazım Hikmet tarzında, deniz üzerine duygusal bir şiir yaz.", 0.85, 256],
        ["Aşık Veysel tarzında, yol temasıyla bir şiir yaz.", 0.85, 256],
        ["Orhan Veli tarzında, İstanbul sokaklarını anlatan bir şiir yaz.", 0.85, 256],
        ["Fazıl Hüsnü Dağlarca tarzında, barış üzerine bir şiir yaz.", 0.85, 256],
        ["Cahit Sıtkı Tarancı tarzında, sonbahar üzerine melankolik bir şiir yaz.", 0.85, 256],
    ],
    additional_inputs=[
        gr.Slider(0.1, 2.0, 0.85, step=0.05, label="🌡️ Temperature (yaratıcılık)"),
        gr.Slider(100, 512, 256, step=16, label="📏 Maksimum uzunluk"),
    ],
    theme=gr.themes.Soft(),
)

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