import os
import sys
import glob

# 🔧 ФИКС CUDA ДЛЯ ZEROGPU
site_packages = os.path.join(sys.prefix, "lib", f"python{sys.version_info.major}.{sys.version_info.minor}", "site-packages")
nvidia_libs = glob.glob(os.path.join(site_packages, "nvidia", "*", "lib"))
if nvidia_libs:
    current_ld_path = os.environ.get("LD_LIBRARY_PATH", "")
    os.environ["LD_LIBRARY_PATH"] = ":".join(nvidia_libs) + (":" + current_ld_path if current_ld_path else "")

os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"

import spaces
import gradio as gr
from huggingface_hub import hf_hub_download
from llama_cpp import Llama

# Gemma 4 использует тот же визуальный энкодер и формат, что и Gemma 3
try:
    from llama_cpp.llama_chat_format import Gemma3ChatHandler as VisionHandler
except ImportError:
    from llama_cpp.llama_chat_format import Llava15ChatHandler as VisionHandler

REPO_ID = "kirilldual0879/rumeni-1-nano-GGUF"
FILENAME = "rumeni-1-nano-Q8_0.gguf"
MMPROJ_FILENAME = "rumeni-1-nano-mmproj-F16.gguf"

print("⬇️ Pre-downloading Q8 model and mmproj...")
MODEL_PATH = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
MMPROJ_PATH = hf_hub_download(repo_id=REPO_ID, filename=MMPROJ_FILENAME)
print("✅ Model and mmproj cached.")

chat_handler = VisionHandler(clip_model_path=MMPROJ_PATH, verbose=False)

def to_openai_format(msg):
    """Конвертирует формат Gradio 5 multimodal в формат, который понимает llama.cpp"""
    if isinstance(msg, dict):
        if "files" in msg and msg["files"]:
            content = []
            for f in msg["files"]:
                content.append({"type": "image_url", "image_url": {"url": f}})
            content.append({"type": "text", "text": msg.get("text", "")})
            return content
        return msg.get("text", "")
    return msg

@spaces.GPU
def respond(message, history):
    llm = Llama(
        model_path=MODEL_PATH,
        chat_handler=chat_handler,
        n_ctx=131072,              # 🚀 128K КОНТЕКСТ (ПОТОЛОК ДЛЯ GEMMA 4)
        n_gpu_layers=-1,           
        n_threads=8,
        verbose=False,
        flash_attn=True,           # Ускоряет генерацию и экономит VRAM
    )
    
    messages = [{"role": "system", "content": "You are Rumeni, a powerful multimodal AI assistant based on Gemma 4. You provide extremely detailed and comprehensive answers."}]
    
    for msg in history:
        messages.append({
            "role": msg["role"],
            "content": to_openai_format(msg["content"])
        })
        
    messages.append({
        "role": "user",
        "content": to_openai_format(message)
    })

    stream = llm.create_chat_completion(
        messages=messages,
        max_tokens=-1,             # 🚀 БЕЗ ЛИМИТОВ (Генерация до <eos> или конца контекста)
        temperature=0.7,
        top_p=0.9,
        stream=True
    )

    response = ""
    buffer = ""
    for chunk in stream:
        if not chunk["choices"]:
            continue
        delta = chunk["choices"][0].get("delta", {})
        if "content" in delta:
            buffer += delta["content"]
            if len(buffer) >= 5 or buffer.endswith(('.', '!', '?', '\n')):
                response += buffer
                buffer = ""
                yield response
                
    if buffer:
        response += buffer
        yield response

demo = gr.ChatInterface(
    fn=respond,
    type="messages",
    multimodal=True,
    title="Rumeni 1 Nano (Gemma 4 + ZeroGPU 🚀)",
    description=(
        "⚡️ **Gemma 4 Arch + Q8_0 + Vision:** Максимальная мощность Nvidia A100/L40S.\n"
        "🧠 **Контекст:** 131,072 токенов (128K) | **Лимит ответа:** ∞ (Без ограничений).\n"
        "👁️ **Vision:** Загрузите изображение через кнопку 📎!\n\n"
        "*⚠️ Генерация огромных текстов может занять время. ZeroGPU работает на пределе возможностей.*"
    ),
    concurrency_limit=1,
    theme="default"
)

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