import os
import random
import requests
import tempfile
from deep_translator import GoogleTranslator
from langdetect import detect
from gradio_client import Client, handle_file
import gradio as gr
from PIL import Image
from io import BytesIO

# Функция для получения случайного 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, timeout=10)  # Добавлено timeout
    response.raise_for_status()
    css = response.text + " h1{text-align:center}"
except requests.exceptions.RequestException as e:
    print(f"Ошибка при загрузке CSS: {e}")
    css = " h1{text-align:center}"  # Используем базовый стиль

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

        # Сохранение изображения во временный файл
        with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as temp_file:
            image.save(temp_file.name)
            image_path = temp_file.name

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

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

        # Инициализация клиента (обработка ошибок)
        client = Client("Yuanshi/OminiControl", hf_token=api_key)
        result = client.predict(image=image_url, text=text, api_name="/predict")

        if result is None:
            return None
        elif isinstance(result, Image.Image):
            return result
        elif isinstance(result, str) and result.startswith("/tmp"):
            try:
                return Image.open(result)
            except FileNotFoundError:
                print(f"Файл не найден: {result}")
                return None
            except Exception as e:
                print(f"Ошибка при открытии файла: {e}, путь: {result}")
                return None
        else:
            print(f"Неизвестный тип результата: {type(result)}, данные: {result}")
            return None
    except Exception as e:
        print(f"Ошибка при обработке: {e}")
        return None

# Функция для загрузки изображения из URL
def load_image_from_url(url):
    try:
        response = requests.get(url, timeout=10)
        response.raise_for_status()
        image = Image.open(BytesIO(response.content))
        return image
    except Exception as e:
        print(f"Ошибка при загрузке изображения из URL {url}: {e}")
        return None

# Функция для получения примеров
def get_samples():
    sample_list = [
        {
            "image": "https://huggingface.co/spaces/NeurixYUFI/Omini/resolve/main/img/pear.png",
            "text": "Луна, данный предмет лежит на поверхности луны — груша. На заднем плане Земля возвышается на переднем плане, ультра реализм",
        },
        {
            "image": "https://huggingface.co/spaces/NeurixYUFI/Omini/resolve/main/img/oranges.jpg",
            "text": "Очень крупный план этого предмета. Он размещен на деревянном столе. Фон — темная комната, включен телевизор, на экране кулинарное шоу. На экране надпись: «Neurix»",
        },
        {
            "image": "https://huggingface.co/spaces/NeurixYUFI/Omini/resolve/main/img/rc_car.jpg",
            "text": "Кинематографичный кадр. На Луне эта вещь едет по лунной поверхности. На заднем плане видна гигантская Земля",
        },
    ]

    samples_with_images = []
    for sample in sample_list:
        image_url = sample["image"]
        image = load_image_from_url(image_url)
        if image:
            samples_with_images.append([image, sample["text"]])

    return samples_with_images

# Интерфейс Gradio с использованием gr.Blocks
with gr.Blocks(css=css) as demo:
    gr.Markdown("# Замена фона")
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil", label="Входное изображение")
            text_input = gr.Textbox(lines=2, label="Описание")
            submit_button = gr.Button("Обработать", variant="primary")  # Добавили variant='primary'
        with gr.Column():
            output_image = gr.Image(type="pil", label="Результат", show_share_button=False)

    submit_button.click(
        fn=process_image_and_text,
        inputs=[image_input, text_input],
        outputs=output_image,
        concurrency_limit=250
    )

    # Добавление примеров
    examples = get_samples()
    if examples:
        with gr.Row():
            gr.Examples(label="Примеры", examples=examples, fn=process_image_and_text, inputs=[image_input, text_input], outputs=output_image)

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