import gradio as gr
from gradio_client import Client, handle_file
import os
import random
from PIL import Image
from io import BytesIO
import requests

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

def swap_face(source_image, target_image, do_face_enhancer):
    try:
        api_key = get_random_api_key()
        client = Client("tuan2308/face-swap", hf_token=api_key)

        temp_dir = "temp"
        os.makedirs(temp_dir, exist_ok=True)

        source_filename = os.path.join(temp_dir, f"source_{random.randint(1000, 9999)}.png")
        target_filename = os.path.join(temp_dir, f"target_{random.randint(1000, 9999)}.png")

        source_image.save(source_filename)
        target_image.save(target_filename)

        result = client.predict(
            source_file=handle_file(source_filename),
            target_file=handle_file(target_filename),
            doFaceEnhancer=do_face_enhancer,
            api_name="/predict",
        )

        if isinstance(result, str) and os.path.exists(result):
            output_image = Image.open(result)
        elif isinstance(result, bytes):
            output_image = Image.open(BytesIO(result))
        else:
            raise ValueError(f"Неожиданный тип результата API: {type(result)}")

        return output_image

    except Exception as e:
        print(f"Ошибка при вызове API: {e}")
        import traceback
        traceback.print_exc()
        raise gr.Error(f"Произошла ошибка при обработке изображений: {e}")


# Ссылка на файл 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 не удалась



with gr.Blocks(css=css) as demo:
    gr.Markdown("# Замена лиц")  # Заголовок
    with gr.Row():
        with gr.Column():
            target_image_input = gr.Image(label="Основа", type="pil")
            source_image_input = gr.Image(label="Лицо", type="pil")
            do_face_enhancer_input = gr.Checkbox(label="Улучшение лица", info="Включить улучшение качества лица? (Дольше)")
            submit_btn = gr.Button("Заменить лицо", variant="primary")
        with gr.Column():
            output_image = gr.Image(label="Результат")

    submit_btn.click(
        fn=swap_face,
        inputs=[source_image_input, target_image_input, do_face_enhancer_input],
        outputs=output_image,
    )


demo.launch()