import os
import gradio as gr
from gradio import ChatMessage
from typing import Iterator
import google.generativeai as genai

# Отримуємо API ключ Gemini з environment variable
GEMINI_API_KEY = os.getenv("AIzaSyAX0KwZz-3vPFFMPcxZrYkK-b4cUqd6fLQ")
genai.configure(api_key=GEMINI_API_KEY)

# Обираємо модель Gemini
model = genai.GenerativeModel("gemini-1.5-flash")  # або інша доступна модель

def форматувати_історію_чату(повідомлення: list) -> list:
    """Форматує історію чату"""
    форматована_історія = []
    for повідомлення in повідомлення:
        if not (повідомлення.get("role") == "assistant" and "metadata" in повідомлення):
            форматована_історія.append({
                "role": "user" if повідомлення.get("role") == "user" else "assistant",
                "parts": [повідомлення.get("content", "")]
            })
    return форматована_історія

def транслювати_відповідь(повідомлення_користувача: str, повідомлення: list) -> Iterator[list]:
    """Працюю"""

    if not повідомлення_користувача.strip():
        повідомлення.append(ChatMessage(role="assistant", content="Будь ласка, введіть непусте текстове повідомлення."))
        yield повідомлення
        return

    try:
        print(f"\n=== Новий запит (Текст) ===")
        print(f"Повідомлення користувача: {повідомлення_користувача}")

        історія_чату = форматувати_історію_чату(повідомлення)
        чат = model.start_chat(history=історія_чату)
        response = чат.send_message(повідомлення_користувача, stream=True)

        буфер_думок = ""
        буфер_відповіді = ""
        думки_завершені = False

        повідомлення.append(
            ChatMessage(
                role="assistant",
                content="",
                metadata={"title": "⚙️ Працюю"}
            )
        )

        for chunk in response:
            частини = chunk.candidates[0].content.parts
            поточний_шматок = частини[0].text

            if len(частини) == 2 and not думки_завершені:
                буфер_думок += поточний_шматок
                print(f"\n=== Завершена думка ===\n{буфер_думок}")

                повідомлення[-1] = ChatMessage(
                    role="assistant",
                    content=буфер_думок,
                    metadata={"title": "⚙️ Працюю"}
                )
                yield повідомлення

                буфер_відповіді = частини[1].text
                print(f"\n=== Починаємо відповідь ===\n{буфер_відповіді}")

                повідомлення.append(ChatMessage(role="assistant", content=буфер_відповіді))
                думки_завершені = True

            elif думки_завершені:
                буфер_відповіді += поточний_шматок
                print(f"\n=== Шматок відповіді ===\n{поточний_шматок}")
                повідомлення[-1] = ChatMessage(role="assistant", content=буфер_відповіді)

            else:
                буфер_думок += поточний_шматок
                print(f"\n=== Шматок думки ===\n{поточний_шматок}")
                повідомлення[-1] = ChatMessage(
                    role="assistant",
                    content=буфер_думок,
                    metadata={"title": "⚙️ Працюю"}
                )

            yield повідомлення

        print(f"\n=== Остаточна відповідь ===\n{буфер_відповіді}")

    except Exception as e:
        print(f"\n=== Помилка ===\n{str(e)}")
        повідомлення.append(ChatMessage(role="assistant", content=f"Вибачте, сталася помилка: {str(e)}"))
        yield повідомлення

def повідомлення_користувача(msg: str, history: list) -> tuple[str, list]:
    """Додає повідомлення користувача до історії чату."""
    history.append(ChatMessage(role="user", content=msg))
    return "", history


with gr.Blocks(theme=gr.themes.Soft(primary_hue="teal", secondary_hue="slate", neutral_hue="neutral")) as demo:
    gr.Markdown(" Чатбот на основі ШІ. З питань роботи додатка: gubarostislav@gmail.com ")

    chatbot = gr.Chatbot(
        type="messages",
        label="Електронний медик",
        render_markdown=True,
        scale=1,
    )

    with gr.Row(equal_height=True):
        input_box = gr.Textbox(
            lines=1,
            label="Повідомлення",
            placeholder="Введіть ваше повідомлення тут...",
            scale=4
        )

        clear_button = gr.Button("Очистити чат", scale=1)

    msg_store = gr.State("")

    input_box.submit(
        lambda msg: (msg, msg, ""),
        inputs=[input_box],
        outputs=[msg_store, input_box, input_box],
        queue=False
    ).then(
        повідомлення_користувача,
        inputs=[msg_store, chatbot],
        outputs=[input_box, chatbot],
        queue=False
    ).then(
        транслювати_відповідь,
        inputs=[msg_store, chatbot],
        outputs=chatbot
    )

    clear_button.click(
        lambda: ([], "", ""),
        outputs=[chatbot, input_box, msg_store],
        queue=False
    )

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