import gradio as gr
import google.generativeai as genai
import os

# Функция для инициализации Gemini
def init_gemini(api_key):
    try:
        genai.configure(api_key=api_key)
        model = genai.GenerativeModel('gemini-pro')
        return model, None
    except Exception as e:
        return None, f"Ошибка инициализации Gemini: {str(e)}"

# Основная функция для чата
def chat_with_gemini(message, history, api_key):
    try:
        model, error = init_gemini(api_key)
        if error:
            return error
        
        # Форматируем историю
        chat_history = []
        for human, ai in history:
            chat_history.append({"role": "user", "parts": [human]})
            chat_history.append({"role": "model", "parts": [ai]})
        
        # Создаем чат
        chat = model.start_chat(history=chat_history)
        response = chat.send_message(message)
        return response.text
        
    except Exception as e:
        return f"Ошибка: {str(e)}"

# Функция проверки API ключа
def check_api_key(api_key):
    if not api_key or api_key.strip() == "":
        return gr.update(visible=True), "Введите ваш API-ключ"
    
    try:
        model, error = init_gemini(api_key)
        if error:
            return gr.update(visible=True), f"Ошибка: {error}"
        
        # Тестовый запрос
        response = model.generate_content("Ответь 'Готов к работе!'")
        return gr.update(visible=False), "Ключ проверен! Можете начинать общение."
        
    except Exception as e:
        return gr.update(visible=True), f"Ключ не прошел проверку: {str(e)}"

# Интерфейс
with gr.Blocks(theme=gr.themes.Soft(), title="Мой ИИ-Ассистент (Gemini)") as demo:
    gr.Markdown("## 🚀 Мой Личный ИИ-Ассистент на основе Google Gemini")
    gr.Markdown("Получите бесплатный API ключ на [ai.google.dev](https://ai.google.dev/)")
    
    with gr.Row():
        api_key = gr.Textbox(
            label="Google Gemini API Key",
            placeholder="Введите ваш API-ключ здесь...",
            type="password",
            lines=1
        )
        check_btn = gr.Button("Проверить ключ", variant="primary")
    
    key_status = gr.Markdown("Статус: Ожидание ввода ключа")
    
    with gr.Column(visible=False) as chat_column:
        chatbot = gr.Chatbot(height=400)
        msg = gr.Textbox(label="Ваше сообщение", placeholder="Напишите что-нибудь...")
        
        with gr.Row():
            submit_btn = gr.Button("Отправить", variant="primary")
            clear_btn = gr.Button("Очистить")

    # Обработчики событий
    check_btn.click(
        fn=check_api_key,
        inputs=[api_key],
        outputs=[chat_column, key_status]
    )
    
    def respond(message, chat_history, api_key):
        if not api_key:
            return "", chat_history + [[message, "Введите API-ключ"]]
        
        bot_message = chat_with_gemini(message, chat_history, api_key)
        chat_history.append((message, bot_message))
        return "", chat_history
    
    msg.submit(respond, [msg, chatbot, api_key], [msg, chatbot])
    submit_btn.click(respond, [msg, chatbot, api_key], [msg, chatbot])
    clear_btn.click(lambda: None, None, chatbot, queue=False)

if __name__ == "__main__":
    demo.launch()