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

# Так надо
model = None

def load_model():
    global model
    try:
        print("Начинаем загрузку модели из Hub...")
        
        model_path = hf_hub_download(
            repo_id="AugustLight/LLight-3.2-3B-Instruct",
            filename="Llight.Q8_0.gguf",
            repo_type="model"
        )
        
        print(f"Модель загружена в: {model_path}")
        
        model = Llama(
            model_path=model_path,
            n_ctx=2048,        # Размер контекста
            n_threads=4,       # Количество потоков
            n_batch=512        # Размер батча
        )
        
        print("Модель успешно инициализирована!")
        return model
        
    except Exception as e:
        print(f"Подробная ошибка при загрузке модели: {str(e)}")
        raise e

def respond(message, history, system_message, max_new_tokens, temperature, top_p):
    try:
        global model
        if model is None:
            model = load_model()
            
        context = f"{system_message}\n\n"
        for user_msg, assistant_msg in history:
            context += f"User: {user_msg}\nAssistant: {assistant_msg}\n"
        context += f"User: {message}\nAssistant: "

        print(f"Генерируем ответ для контекста длиной {len(context)} символов")
        
        response = model(
            prompt=context,
            max_tokens=max_new_tokens,
            temperature=temperature,
            top_p=top_p,
            stop=["User:", "\n\n", "<|endoftext|>"],
            echo=False  # Не возвращать промпт в ответе
        )
        
        generated_text = response['choices'][0]['text']
        print(f"Ответ сгенерирован успешно, длина: {len(generated_text)}")
        return generated_text.strip()
        
    except Exception as e:
        error_msg = f"Произошла ошибка: {str(e)}"
        print(error_msg)
        return error_msg


demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Textbox(
            value="Ты дружелюбный и полезный ассистент. Отвечай обдуманно и по делу.", 
            label="System message"
        ),
        gr.Slider(
            minimum=1, 
            maximum=2048, 
            value=512, 
            step=1, 
            label="Max new tokens"
        ),
        gr.Slider(
            minimum=0.1, 
            maximum=2.0, 
            value=0.3, 
            step=0.1, 
            label="Temperature"
        ),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05, 
            label="Top-p (nucleus sampling)"
        ),
    ],
    title="GGUF Chat Model",
    description="Чат с GGUF моделью (LLight-3.2-3B-Instruct)",
    examples=[
        ["Привет! Как дела?"],
        ["Расскажи мне о себе"],
        ["Что ты умеешь делать?"]
    ],
    cache_examples=False
)

# Запускаем приложение
if __name__ == "__main__":
    try:
        print("Инициализация приложения...")
        model = load_model()
        print("Модель загружена успешно при старте")
    except Exception as e:
        print(f"Ошибка при инициализации: {str(e)}")
    
    demo.launch()