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

# Global model instance
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=512,
            n_threads=os.cpu_count(),
            n_batch=128, 
            n_gpu_layers=0,  
            embedding_cache_size=1024
        )
        
        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:
            print("Загружаем модель...")
            model = load_model()
            print("Модель загружена")
        
        with gr.Progress() as progress:
            progress(0, desc="Подготовка контекста...")
            
            # Отладочная печать входных параметров
            print(f"""
            Входные параметры:
            - message: {message}
            - history length: {len(history)}
            - system_message: {system_message}
            - max_new_tokens: {max_new_tokens}
            - temperature: {temperature}
            - top_p: {top_p}
            """)
            
            # Ограничиваем историю последними 3 сообщениями
            recent_history = history[-3:] if len(history) > 3 else history
            
            context = f"{system_message}\n\n"
            for user_msg, assistant_msg in recent_history:
                context += f"User: {user_msg}\nAssistant: {assistant_msg}\n"
            context += f"User: {message}\nAssistant: "
            
            print(f"Генерируем ответ для контекста длиной {len(context)} символов")
            
            try:
                progress(0.3, desc="Генерация ответа...")
                response = model(
                    prompt=context,
                    max_tokens=max_new_tokens,
                    temperature=temperature,
                    top_p=top_p,
                    stop=["User:", "\n\n", "<|endoftext|>"],
                    echo=False
                )
                
                progress(1, desc="Готово!")
                generated_text = response['choices'][0]['text']
                print(f"Ответ сгенерирован успешно, длина: {len(generated_text)}")
                return generated_text.strip()
            except Exception as inner_e:
                print(f"Ошибка при генерации: {str(inner_e)}")
                print(f"Тип ошибки: {type(inner_e).__name__}")
                return f"Ошибка при генерации: {str(inner_e)}"
        
    except KeyboardInterrupt:
        return "Генерация прервана пользователем"
    except Exception as e:
        print(f"Основная ошибка: {str(e)}")
        print(f"Тип ошибки: {type(e).__name__}")
        return f"Произошла ошибка: {str(e)}\nТип ошибки: {type(e).__name__}"

# Создаем интерфейс с оптимизированными параметрами
demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Textbox(
            value="Ты дружелюбный и полезный ассистент. Отвечай обдуманно и по делу.", 
            label="System message"
        ),
        gr.Slider(
            minimum=1, 
            maximum=512,
            value=128,
            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="LLight Chat Model (Optimized)",
    description="Оптимизированный чат с LLight-3.2-3B",
    examples = [
        ["Привет! Как дела?", 
         "Ты дружелюбный и полезный ассистент. Отвечай обдуманно и по делу.",  # system_message
         128,    # max_new_tokens
         0.3,    # temperature
         0.95    # top_p
        ],
        ["Расскажи мне о себе",
         "Ты дружелюбный и полезный ассистент. Отвечай обдуманно и по делу.",
         128,
         0.3,
         0.95
        ],
        ["Что ты умеешь делать?",
         "Ты дружелюбный и полезный ассистент. Отвечай обдуманно и по делу.",
         128,
         0.3,
         0.95
        ]
    ],
    cache_examples=True
)

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