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

###############################################################################
# 1. Настройка окружения и инициализация моделей
###############################################################################

# Подставьте свой ключ или берите из окружения
GEMINI_API_KEY = "AIzaSyBoqoPX-9uzvXyxzse0gRwH8_P9xO6O3Bc"
if not GEMINI_API_KEY:
    print("Error: GEMINI_API_KEY is not set.")
    exit()

genai.configure(api_key=GEMINI_API_KEY)

# Выберите доступные модели (пример)
AVAILABLE_MODELS = [
    "gemini-2.0-flash-exp",
    "gemini-exp-1206",
    "gemini-2.0-flash-thinking-exp-1219",
]

MODELS = {}
for model_name in AVAILABLE_MODELS:
    try:
        MODELS[model_name] = genai.GenerativeModel(model_name=model_name)
    except Exception as e:
        print(f"[Предупреждение] Не удалось инициализировать модель {model_name}: {e}")

###############################################################################
# 2. Дефолтные промпты (developer role) для каждой модели
###############################################################################
# Когда пользователь переключается на модель, мы добавляем это сообщение в историю.

DEFAULT_DEVELOPER_PROMPTS = {
    "gemini-2.0-flash-exp": (
        "You are a normal model (developer role). "
        "Provide direct answers with no JSON wrapping."
    ),
    "gemini-exp-1206": (
        "You are an experimental normal model (developer role). "
        "Provide direct answers with no JSON wrapping."
    ),
    "gemini-2.0-flash-thinking-exp-1219": (
        "You are a thinking model (developer role). "
        "Please provide your final answer in the format {output: ...}. "
        "You may use internal thoughts but do not show them directly to the user."
    ),
}

###############################################################################
# 3. Функция для определения роли ассистента (assistant vs model)
###############################################################################

def _assistant_role(model_name: str) -> str:
    """
    Некоторые новые модели не принимают 'assistant', а требуют 'model'.
    """
    # Допустим "gemini-exp-1206" и "gemini-2.0-flash-thinking-exp-1219" хотят "model"
    if model_name in ["gemini-exp-1206", "gemini-2.0-flash-thinking-exp-1219"]:
        return "model"
    return "assistant"

###############################################################################
# 4. Преобразование истории из Gradio в формат Generative AI
###############################################################################

def _history_to_genai_enhanced(history, model_name):
    """
    Улучшенная версия, отличающая developer-сообщения
    (префикс "<developer>: ") от user-сообщений.
    """
    asst_role = _assistant_role(model_name)
    genai_history = []

    for user_text, bot_text in history:
        if user_text:
            if user_text.startswith("<developer>: "):
                # Считаем это developer role
                dev_content = user_text.replace("<developer>: ", "", 1)
                genai_history.append({"role": "system", "parts": dev_content})
            else:
                # Обычный пользователь
                genai_history.append({"role": "user", "parts": user_text})

        if bot_text:
            # Ответ ассистента / модель
            genai_history.append({"role": asst_role, "parts": bot_text})

    return genai_history

###############################################################################
# 5. Генераторы для стрима обычных моделей и "thinking" моделей
###############################################################################

async def _respond_stream_enh(model_name, user_message, history):
    """
    Стриминговый ответ для обычных моделей:
    - Кусочек за кусочком (partial_text).
    """
    if model_name not in MODELS:
        yield "Ошибка: модель не найдена."
        return

    model = MODELS[model_name]
    genai_history = _history_to_genai_enhanced(history, model_name)

    try:
        chat = model.start_chat(history=genai_history)
        stream = chat.send_message(user_message, stream=True)

        partial_text = ""
        async for chunk in stream:
            partial_text += (chunk.text or "")
            yield partial_text

        return
    except Exception as e:
        yield f"Ошибка при запросе к API: {e}"
        return

async def _respond_thinking_enh(model_name, user_message, history):
    """
    Для thinking-моделей:
    1) Выводим "Думаю..."
    2) После завершения — финальный ответ в формате {output: ...} + размышления.
    """
    if model_name not in MODELS:
        yield "Ошибка: модель не найдена.", ""
        return

    model = MODELS[model_name]
    genai_history = _history_to_genai_enhanced(history, model_name)

    # Сначала "Думаю..."
    yield "Думаю...", ""

    try:
        chat = model.start_chat(history=genai_history)
        response = chat.send_message(user_message, stream=False)

        thinking_process_text = ""
        final_text = ""

        if response.candidates:
            parts = response.candidates[0].content.parts
            for p in parts:
                if hasattr(p, "thought") and p.thought:
                    thinking_process_text += p.text or ""
                else:
                    final_text += p.text or ""

        # Для thinking-моделей просили итоговый ответ в {output: ...}
        final_text_formatted = f"{{output: {final_text}}}"

        yield final_text_formatted, thinking_process_text
        return

    except Exception as e:
        yield f"Ошибка при запросе к API: {e}", ""
        return

###############################################################################
# 6. Основная функция для ввода пользователя
###############################################################################

async def user_send_message(
    user_message: str,
    history: list[tuple[str, str]],
    model_name: str,
    thinking_text: str
):
    """
    Колбэк, когда пользователь отправляет запрос.
    Добавляем в history новый (user_msg, None), затем генерируем ответ.
    """
    # Пустой ввод
    if not user_message.strip():
        yield history, thinking_text
        return

    # Добавляем (user_message, None)
    history.append((user_message, None))

    # Если модель — thinking
    if "thinking" in model_name.lower():
        async for (assistant_text, thought_text) in _respond_thinking_enh(model_name, user_message, history):
            history[-1] = (user_message, assistant_text)
            yield history, thought_text
        return
    else:
        # Обычная модель
        partial_answer = ""
        async for chunk in _respond_stream_enh(model_name, user_message, history):
            partial_answer = chunk
            history[-1] = (user_message, partial_answer)
            yield history, ""
        return

###############################################################################
# 7. Очистка диалога
###############################################################################

def clear_all():
    """Сброс истории и размышлений."""
    return [], ""

###############################################################################
# 8. Когда меняем модель в Dropdown
###############################################################################

def on_model_change(selected_model, history, thinking_text):
    """
    При переключении модели добавляем в историю developer-сообщение,
    + добавляем дефолтный промпт этой модели (тоже developer).
    """
    new_history = history.copy()

    # Cообщаем модели, что переключились (developer role)
    new_history.append((
        "<developer>: Switched to model: " + selected_model,
        None
    ))

    # Добавляем дефолтный промпт (developer role)
    default_prompt = DEFAULT_DEVELOPER_PROMPTS.get(
        selected_model,
        "No default prompt for this model."
    )
    new_history.append((
        "<developer>: " + default_prompt,
        None
    ))

    return new_history, thinking_text

###############################################################################
# 9. Функция конвертации истории с учётом developer role
###############################################################################

def _history_to_genai_enhanced(history, model_name):
    """
    Улучшенная версия, отличающая developer-сообщения
    (префикс "<developer>: ") от user-сообщений.
    """
    asst_role = _assistant_role(model_name)
    genai_history = []

    for user_text, bot_text in history:
        if user_text:
            if user_text.startswith("<developer>: "):
                # Считаем это developer role
                dev_content = user_text.replace("<developer>: ", "", 1)
                genai_history.append({"role": "system", "parts": dev_content})
            else:
                # Обычный пользователь
                genai_history.append({"role": "user", "parts": user_text})

        if bot_text:
            # Ответ ассистента / модель
            genai_history.append({"role": asst_role, "parts": bot_text})

    return genai_history

###############################################################################
# 10. Построение интерфейса Gradio
###############################################################################

with gr.Blocks() as demo:
    gr.Markdown("## Chat с Gemini. Поддержка developer role, переключения моделей, JSON-ответа для thinking")

    with gr.Row():
        model_dropdown = gr.Dropdown(
            choices=AVAILABLE_MODELS,
            value="gemini-2.0-flash-exp",
            label="Выберите модель"
        )
        clear_button = gr.Button("Очистить чат")

    history_state = gr.State([])
    thinking_store = gr.State("")

    chatbot = gr.Chatbot(label="Диалог с Gemini")
    user_input = gr.Textbox(label="Ваш вопрос", placeholder="Введите текст...")
    thinking_output = gr.Textbox(label="Размышления", interactive=False)
    send_btn = gr.Button("Отправить")

    ################################################
    # (A) Обработка переключения модели
    ################################################
    def handle_model_change(selected_model, history, thinking):
        new_history, new_thinking = on_model_change(selected_model, history, thinking)
        return new_history, new_thinking

    # Когда пользователь меняет модель:
    model_change = model_dropdown.change(
        fn=handle_model_change,
        inputs=[model_dropdown, history_state, thinking_store],
        outputs=[history_state, thinking_store],
        queue=False
    ).then(
        # После добавления developer-сообщений в историю → обновляем чат
        fn=lambda h: h,
        inputs=[history_state],
        outputs=[chatbot],
        queue=False
    )

    ################################################
    # (B) При нажатии «Отправить»
    ################################################
    send_chain = send_btn.click(
        fn=user_send_message,
        inputs=[user_input, history_state, model_dropdown, thinking_store],
        outputs=[history_state, thinking_store],
        queue=True
    )
    send_chain.then(
        fn=lambda h: h,
        inputs=[history_state],
        outputs=[chatbot],
        queue=True
    )
    send_chain.then(
        fn=lambda t: t,
        inputs=[thinking_store],
        outputs=[thinking_output],
        queue=True
    )
    # Очистка поля ввода
    send_chain.then(
        fn=lambda: "",
        inputs=[],
        outputs=[user_input],
        queue=False
    )

    ################################################
    # (C) При нажатии Enter в textbox
    ################################################
    submit_chain = user_input.submit(
        fn=user_send_message,
        inputs=[user_input, history_state, model_dropdown, thinking_store],
        outputs=[history_state, thinking_store],
        queue=True
    )
    submit_chain.then(
        fn=lambda h: h,
        inputs=[history_state],
        outputs=[chatbot],
        queue=True
    )
    submit_chain.then(
        fn=lambda t: t,
        inputs=[thinking_store],
        outputs=[thinking_output],
        queue=True
    )
    # Очистка поля ввода
    submit_chain.then(
        fn=lambda: "",
        inputs=[],
        outputs=[user_input],
        queue=False
    )

    ################################################
    # (D) Кнопка «Очистить»
    ################################################
    clear_chain = clear_button.click(
        fn=clear_all,
        inputs=[],
        outputs=[history_state, thinking_store],
        queue=False
    )
    clear_chain.then(
        fn=lambda h: h,
        inputs=[history_state],
        outputs=[chatbot]
    )
    clear_chain.then(
        fn=lambda _: "",
        inputs=[],
        outputs=[thinking_output]
    )
    clear_chain.then(
        fn=lambda: "",
        inputs=[],
        outputs=[user_input]
    )

# Запуск
if __name__ == "__main__":
    demo.launch()