import os
import hashlib
import sys
import torch
import spaces  # ZeroGPU için gerekli
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from threading import Thread

# İstediğin modeller sabit
DEFAULT_MODELS = [
    "Qwen/Qwen3.8-27B",
    "Qwen/Qwen3.5-9B"
]

TARGET_TOKEN_HASH = "8fb29448faee18b656030e8f5a8b9e9a695900f36a3b7d7ebb0d9d51e06c8569d81a55e39b481cf50546d697e7bde1715aa6badede8ddc801c739777be77f166"

# VRAM taşmasını (OOM) önlemek için modelleri istek geldikçe tek tek belleğe alıyoruz
models_cache = {}
tokenizers_cache = {}

def get_or_load_model(model_name):
    """OOM hatasını önlemek için modelleri ihtiyaç anında (tekil olarak) yükler ve önbelleğe alır."""
    if model_name in models_cache:
        print(f"⚡ Model önbellekten (cache) getirildi: {model_name}", flush=True)
        return tokenizers_cache[model_name], models_cache[model_name]
    
    print(f"📥 Model yükleniyor (VRAM yönetimi aktif): {model_name}...", flush=True)
    
    # Bellekte yer açmak için varsa eski modeli temizle
    if models_cache:
        print("🧹 Eski model VRAM'den temizleniyor...", flush=True)
        models_cache.clear()
        tokenizers_cache.clear()
        torch.cuda.empty_cache()

    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
        device_map="auto",
        trust_remote_code=True
    )
    
    tokenizers_cache[model_name] = tokenizer
    models_cache[model_name] = model
    print(f"✅ Model başarıyla yüklendi: {model_name}", flush=True)
    return tokenizer, model

@spaces.GPU(duration=120)
def api_completion(model_name, lsu_accesstoken, messages, temperature, max_tokens, top_p):
    print("🚀 API İsteği Alındı, GPU bağlamı aktifleşti.", flush=True)
    
    if not lsu_accesstoken:
        print("❌ Hata: lsu_accesstoken boş.", flush=True)
        yield "Hata: Yetkilendirme token'ı (lsu_accesstoken) eksik."
        return

    incoming_hash = hashlib.sha512(lsu_accesstoken.encode('utf-8')).hexdigest()
    
    if incoming_hash != TARGET_TOKEN_HASH:
        print("❌ Hata: Geçersiz hash eşleşmesi.", flush=True)
        yield "Hata: Geçersiz erişim token'ı. Erişim reddedildi."
        return

    if not messages:
        print("❌ Hata: messages payload boş.", flush=True)
        yield "Hata: Geçerli bir 'messages' payload gönderilmedi."
        return

    if model_name not in DEFAULT_MODELS:
        print(f"❌ Hata: Geçersiz model seçimi -> {model_name}", flush=True)
        yield f"Hata: Geçersiz model seçimi ({model_name})."
        return

    print(f"⚙️ Model işleniyor: {model_name}", flush=True)
    tokenizer, model = get_or_load_model(model_name)
    
    try:
        prompt = tokenizer.apply_chat_template(
            messages, 
            tokenize=False, 
            add_generation_prompt=True
        )
    except Exception as e:
        print(f"⚠️ Chat template hatası: {e}, fallback uygulanıyor.", flush=True)
        prompt = ""
        for m in messages:
            role = m.get("role", "user").capitalize()
            content = m.get("content", "")
            prompt += f"{role}: {content}\n"
        prompt += "Assistant:"

    inputs = tokenizer([prompt], return_tensors="pt").to(model.device)
    
    streamer = TextIteratorStreamer(
        tokenizer, 
        skip_prompt=True, 
        skip_special_tokens=True
    )
    
    generation_kwargs = dict(
        **inputs,
        streamer=streamer,
        max_new_tokens=int(max_tokens),
        temperature=float(temperature) if temperature > 0 else 0.01,
        do_sample=True if temperature > 0 else False,
        top_p=float(top_p),
        pad_token_id=tokenizer.pad_token_id
    )
    
    print("✍️ Metin üretimi (generation) başlatılıyor...", flush=True)
    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()
    
    # Think filtreleme mekanizması
    accumulated_text = ""
    thinking_done = False
    sent_length = 0

    for new_text in streamer:
        accumulated_text += new_text
        
        if not thinking_done:
            if "</think>" in accumulated_text:
                thinking_done = True
                idx = accumulated_text.find("</think>") + len("</think>")
                clean_part = accumulated_text[idx:].lstrip("\n")
                if clean_part:
                    yield clean_part
                    sent_length = len(clean_part)
        else:
            idx = accumulated_text.find("</think>") + len("</think>")
            clean_part = accumulated_text[idx:].lstrip("\n")
            if len(clean_part) > sent_length:
                new_chunk_to_send = clean_part[sent_length:]
                yield new_chunk_to_send
                sent_length = len(clean_part)

    print("✅ Yanıt akışı tamamlandı.", flush=True)

with gr.Blocks(analytics_enabled=False) as demo:
    gr.Markdown(
        """
        # 🛑 Restricted Access API Endpoint
        
        Bu Space doğrudan **API Uç Noktası (API Endpoint)** olarak çalışmaktadır. Görsel arayüz kapatılmıştır.
        
        Kullanım detayları, endpoint adresleri ve entegrasyon kılavuzu için lütfen [README.md](https://huggingface.co/spaces/aertsimon90/OBLIHN_API/blob/main/README.md) dosyasını inceleyin.
        """
    )

    api_trigger = gr.Interface(
        fn=api_completion,
        inputs=[
            gr.Dropdown(choices=DEFAULT_MODELS, value=DEFAULT_MODELS[0]),
            gr.Textbox(),
            gr.JSON(),
            gr.Slider(value=0.7),
            gr.Slider(value=1024),
            gr.Slider(value=0.85)
        ],
        outputs=gr.Textbox(),
        api_name="predict"
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)