import gradio as gr
from huggingface_hub import InferenceClient
from typing import List, Tuple

# Список 10 моделей (добавь свои)
MODELS = [
    "microsoft/Phi-3-mini-4k-instruct",
    "Qwen/Qwen2.5-0.5B-Instruct",
    "Qwen/Qwen2.5-1.5B-Instruct",
    "google/gemma-2-2b-it",
    "HuggingFaceH4/zephyr-7b-beta",
    "mistralai/Mistral-7B-Instruct-v0.3",
    "unsloth/Mistral-Nemo-Instruct-2407-12B-GGUF",
    "microsoft/Phi-3.5-mini-instruct",
    "NousResearch/Hermes-2-Theta-Llama-3.1-8B",
    "cognitivecomputations/dolphin-2.9-llama3-8b"
]

def convert_history(history: List[Tuple[str, str]]) -> List[dict]:
    """Конвертирует Gradio history [[user_msg, bot_msg]] в OpenAI messages."""
    messages = []
    for user_msg, bot_msg in history:
        messages.append({"role": "user", "content": user_msg})
        if bot_msg:  # Если ответ есть
            messages.append({"role": "assistant", "content": bot_msg})
    return messages

def chat_response(message: str, history: List[Tuple[str, str]], model_id: str, system_prompt: str):
    """Главная функция: отправляет чат в HF Inference API."""
    try:
        client = InferenceClient(model=model_id)
        
        # Строим messages
        messages = []
        if system_prompt.strip():  # Если system не пустой
            messages.append({"role": "system", "content": system_prompt})
        
        messages.extend(convert_history(history))  # История чата
        messages.append({"role": "user", "content": message})  # Текущее сообщение
        
        # Генерируем ответ (max_new_tokens=512 для скорости)
        response = client.chat_completion(
            messages=messages,
            max_tokens=512,  # Лимит токенов ответа (токены ~ слова/4)
            temperature=0.7,  # Креативность (0=детерминировано, 1=случайно)
            stream=False  # Без стриминга для простоты
        )
        
        return response.choices[0].message.content
    
    except Exception as e:
        # Обработка ошибок (rate limit, модель не найдена)
        return f"Ошибка: {str(e)}. Проверь модель или подожди (API лимит)."

# UI: ChatInterface с доп. полями
with gr.Blocks(title="Тест Чат-Ботов HF") as demo:
    gr.Markdown("# Тестер ИИ-моделей HF\nВыбери модель, system prompt (опционально), чатай!")
    
    # Дополнительные inputs (выше чата)
    model_dropdown = gr.Dropdown(
        choices=MODELS,
        value=MODELS[0],  # По умолчанию Phi-3
        label="Модель HF (выбери или замени на свою)",
        interactive=True
    )
    system_input = gr.Textbox(
        label="System Prompt (системное сообщение, опционально)",
        placeholder="Пример: Ты полезный ассистент. Отвечай кратко.",
        lines=2
    )
    
    # Чат
    chat = gr.ChatInterface(
        fn=chat_response,
        additional_inputs=[model_dropdown, system_input],
        title="Чат с моделью",
        description="Тестируй промпты, jailbreak, код. Смени модель — чат обновится.",
        examples=None,  # Добавь свои примеры позже
        cache_examples=False,
        retry_btn="🔄 Повторить",
        undo_btn="↶ Назад",
        clear_btn="🗑️ Очистить"
    )

if __name__ == "__main__":
    demo.queue(max_size=10).launch(share=True, debug=True)  # share=True для публичной ссылки