import os
import random
from typing import List, Tuple, Union
from gradio_client import Client
import gradio as gr

# Получаем список ключей из переменной окружения
api_keys = os.getenv("KEYS", "").split(",")
if not api_keys:
    raise ValueError("API keys are not set in the environment variable KEYS")

# Выбираем случайный ключ для каждого запроса
def get_random_api_key():
    return random.choice(api_keys)

# Функция для вызова API /add_text
def add_text(_input: dict, _chatbot: List[Tuple[dict, dict]]) -> Tuple[dict, List[Tuple[dict, dict]]]:
    api_key = get_random_api_key()
    client = Client("Qwen/Qwen2.5-Turbo-1M-Demo", hf_token=api_key)
    result = client.predict(
        _input=_input,
        _chatbot=_chatbot,
        api_name="/add_text"
    )
    return result

# Функция для вызова API /agent_run
def agent_run(_chatbot: List[Tuple[dict, dict]]) -> List[Tuple[dict, dict]]:
    api_key = get_random_api_key()
    client = Client("Qwen/Qwen2.5-Turbo-1M-Demo", hf_token=api_key)
    result = client.predict(
        _chatbot=_chatbot,
        api_name="/agent_run"
    )
    return result

# Функция для вызова API /flushed
def flushed() -> dict:
    api_key = get_random_api_key()
    client = Client("Qwen/Qwen2.5-Turbo-1M-Demo", hf_token=api_key)
    result = client.predict(
        api_name="/flushed"
    )
    return result

# Основная функция приложения
def app_gui():
    def chat(message, chat_history):
        # Добавляем текст в чат
        _input = {"files": [], "text": message}
        _chatbot = chat_history if chat_history else []
        response, _chatbot = add_text(_input, _chatbot)
        return "", _chatbot

    def clear_chat():
        # Очищаем чат
        _chatbot = flushed()
        return _chatbot

    # Создаем интерфейс Gradio
    with gr.Blocks() as demo:
        chatbot = gr.Chatbot(label="Chatbot")
        with gr.Row():
            msg = gr.Textbox(label="Message")
            btn = gr.Button("Send")

        btn.click(chat, [msg, chatbot], [msg, chatbot])
        msg.submit(chat, [msg, chatbot], [msg, chatbot])
        clear = gr.Button("Clear")
        clear.click(clear_chat, None, chatbot)

    demo.launch()

if __name__ == '__main__':
    app_gui()