import os
import random
import requests
from deep_translator import GoogleTranslator
from langdetect import detect
from gradio_client import Client, handle_file
import gradio as gr

# Функция для получения случайного API ключа
def get_random_api_key():
    keys = os.getenv("KEYS", "").split(",")
    if keys and keys[0]:  # Проверяем, что KEYS установлен и не пуст
        return random.choice(keys).strip()
    else:
        raise ValueError("API ключи не найдены. Пожалуйста, установите переменную окружения KEYS.")

# Ссылка на файл CSS
css_url = "https://neurixyufi-aihub.static.hf.space/style.css"

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

# Функция для обработки изображения и текста через API
def process_image_and_text(image, text):
    # Проверка и перевод текста на английский, если необходимо
    language = detect(text)
    if language != 'en':
        text = GoogleTranslator(source=language, target='en').translate(text)

    # Загрузка изображения в кеш и получение ссылки
    image_path = image.name
    image_url = handle_file(image_path)

    # Получение случайного API ключа
    api_key = get_random_api_key()

    # Инициализация клиента
    client = Client("Yuanshi/OminiControl", hf_token=api_key)

    # Вызов API
    result = client.predict(
        image=image_url,
        text=text,
        api_name="/predict"
    )

    # Возврат результата
    return result['path']

# Интерфейс Gradio
demo = gr.Interface(
    fn=process_image_and_text,
    inputs=[
        gr.Image(type="pil"),
        gr.Textbox(lines=2),
    ],
    outputs=gr.Image(type="pil", show_share_button=False),
    title="OminiControl / Subject driven generation",
    css=css,
    concurrency_limit=250
)

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