import gradio as gr
from huggingface_hub import InferenceClient
import requests
import json
import os
import random

def get_random_api_key():
    keys = os.getenv("KEYS", "").split(",")
    if keys and keys[0]:  # Check if KEYS is set and not empty
        return random.choice(keys).strip()
    else:
        raise ValueError("API keys not found. Please set the KEYS environment variable.")

def respond(
    message,
    history: list[tuple[str, str]],
    max_tokens,
    temperature,
    top_p,
):
    if not message:
        return history, ""
    
    system_role = """Ты ии чат для создания изображений!

**!Отвечай Markdown разметкой!**
Жирный, список и тд всё html разметкой

Не пиши "пример запроса", "как делать запрос" и тп. Я программа, а обычному пользователю не нужны подробности как да что за api. У него просто начинается генерация и всё, ему не нужно знать как устроены запросы

Чтобы создать изображение, напиши ссылку изображения в теге "SimpleAiImageGen", вставив промпт (обязательно на английском), ширину, высоту, сид и другие переменные.
Можно не указывать модель (рекомендуется, автоматически выберется), но можно указать &model= одну из доступных: Flux,Flux-Pro,Flux-Realism,Flux-Anime,Flux-3D,Flux-CablyAl,Turbo
Чтобы ИЗМЕНИТЬ изображение, можно чуть чуть поменять промпт или добавить к нему, сохраняя тот же сид! Чтобы получить другое изображение с тем же промптом, напиши рандомный сид.
!За раз ты можешь создать МАКСИМУМ 6 изображений, но поумолчанию 1 делай! (Больше не пиши!)

Примеры (писать нужно так, в особенном теге):
![Краткий заголовок](https://api.neurix.ru/imggen?prompt=logo%20with%20text%20pollinations&width=1024&height=512&seed=4)

![Краткий заголовок](https://api.neurix.ru/imggen?prompt=logo%20with%20text%20YUFI&width=1024&height=512&seed=4)

У пользователя начнётся генерация.

Активно ведёшь чат с пользователем, помимо генераций изображений

Подробные промпты (2-3 предложения на английском), ты умеешь ТОЛЬКО создавать изображения и помогать в обучении промптингу, больше НЕ ЧЕГО!

Есть ещё один api (с ним тоже можешь генерировать. Но он рисует дольше, может генерироваться до 2 минут, рисуй этим только если попросят например нарисовать определённой моделью и скинут название (типо: strangerzonehf/Flux-Animex-v2-LoRA, но смотри чтобы было так, а не не правильно, иначе скажи что это не правильно)): "https://yufi-api-serverimgsgens.hf.space/prompt/":
У него запрос такой же, но модели (название модели для пользователя, показывай пользователю только его (ключевое слово которое должно быть в начале промпта): как указывать в запросе):
Flux-logo (при использовании указывай ключевое слово "logo"): Shakker-Labs/FLUX.1-dev-LoRA-Logo-Design
Schnell: black-forest-labs/FLUX.1-schnell
Midjourney (ключ. слово: "midjourney mix"): strangerzonehf/Flux-Midjourney-Mix-LoRA
Animex (ключ. слово: "animex"): strangerzonehf/Flux-Animex-v2-LoRA
In Context: ali-vilab/In-Context-LoRA
Aura (ключ. слово: "aura 9999"): prithivMLmods/Aura-9999
PixelArt: nerijs/pixel-art-xl
PencilArt (ключ. слово: "Simple Pencil"): prithivMLmods/Super-Pencil-Flux-LoRA
WallpaperArt: prithivMLmods/Canopus-Liquid-Wallpaper-Art

![Краткий заголовок](https://api.neurix.ru/imghfgen?prompt=Animex.%20logo%20with%20text%20YUFI&width=1024&height=512&seed=325346&model=strangerzonehf/Flux-Animex-v2-LoRA)
"""
    # Формируем сообщения для отправки в API
    messages = [{"role": "system", "content": system_role}]
    
    for val in history:
        if val[0]:
            user_message = val[0]
            messages.append({"role": "user", "content": user_message})
        if val[1]:
            assistant_message = val[1]
            messages.append({"role": "assistant", "content": assistant_message})
    
    messages.append({"role": "user", "content": message})
    
    # Получаем рандомный API ключ
    api_key = get_random_api_key()
    client = InferenceClient(api_key=api_key)
    
    # Генерируем ответ от чат-бота
    response = client.chat.completions.create(
        model="Qwen/Qwen2.5-Coder-32B-Instruct", 
        messages=messages, 
        max_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p
    )
    
    # Переводим ответ ассистента на русский, если он не на русском
    response_text = response.choices[0].message.content
    
    # Обновляем историю сообщений
    history.append((message, response_text))
    
    return history, ""

# Ссылка на файл CSS
css_url = "https://neurixyufi-aihub.static.hf.space/style.css"
default = """ .gradio-container{max-width: 700px !important} h1{text-align:center} #component-4 { height: 70vh !important; } .bot p img {
min-height: 200px;
    background: url(https://huggingface.co/spaces/NeurixYUFI/ImgGenChat/resolve/main/placeholder.png);
    background-size: cover;
    background-position: center;
    width: auto;
}
.message-wrap.svelte-vxn3uw .message-row .md img {
    width: 100% !important;
    max-width: none !important;
    max-height: none !important;
}
.bubble.bot-row.svelte-pbswnk.svelte-pbswnk {
    min-width: calc(100% - var(--spacing-xl)* 6);
}"""

# Получение CSS по ссылке
try:
    response = requests.get(css_url)
    response.raise_for_status() # Поднимаем исключение, если статус ответа не 200
    css = response.text + default
except requests.exceptions.RequestException as e:
    print(f"Ошибка при загрузке CSS: {e}")
    css = default # Используем базовый стиль, если загрузка CSS не удалась

# Создаем интерфейс с использованием gr.Blocks
with gr.Blocks(css=css) as demo:
    gr.Markdown("# Чат художник")
    
    with gr.Row():
        with gr.Column():
            chatbot = gr.Chatbot(show_label=False, show_share_button=False)
            message = gr.Textbox(label="Введите ваше сообщение", placeholder="Введите ваше сообщение здесь...", lines=3, container=False)
            submit = gr.Button("Отправить", variant='primary')
    
    with gr.Accordion("Настройки чата", open=False):
        max_tokens = gr.Slider(minimum=100, maximum=3000, value=2000, step=1, label="Максимальное количество новых токенов")
        temperature = gr.Slider(minimum=0.0, maximum=1.0, value=0.7, step=0.1, label="Температура")
        top_p = gr.Slider(minimum=0.0, maximum=1.0, value=0.95, step=0.05, label="Top-p (нуклеарное сэмплирование)")
    
    # Обработка отправки сообщения
    submit.click(
        fn=respond,
        inputs=[message, chatbot, max_tokens, temperature, top_p],
        outputs=[chatbot, message],
        concurrency_limit=250
    )
    
    # Обработка отправки сообщения при нажатии Ctrl+Enter
    message.submit(
        fn=respond,
        inputs=[message, chatbot, max_tokens, temperature, top_p],
        outputs=[chatbot, message],
        concurrency_limit=250
    )

# Запускаем интерфейс
if __name__ == "__main__":
    demo.queue(max_size=250).launch(show_api=False, share=False)