import gradio as gr
import os
import random
from PIL import Image
from io import BytesIO
import base64
import requests

# Получаем ключи API из секретной переменной
KEYS = os.getenv("KEYS", "").split(",")

# Создаем функцию для получения случайного ключа API
def get_random_api_key():
    if not KEYS:
        raise ValueError("KEYS environment variable is not set or empty.")
    return random.choice(KEYS)

def swap_face_api(source_img, target_img, doFaceEnhancer):
    try:
        api_key = get_random_api_key()
        api_url = "https://tuan2308-face-swap-predict.hf.space/api/predict/"

        # Преобразуем изображения в base64
        source_bytes = BytesIO()
        source_img.save(source_bytes, format="JPEG")
        source_b64 = source_bytes.getvalue()

        target_bytes = BytesIO()
        target_img.save(target_bytes, format="JPEG")
        target_b64 = target_bytes.getvalue()

        # Создаем полезные данные для POST-запроса
        payload = {
            "source_file": ("source.jpg", source_b64, "image/jpeg"),
            "target_file": ("target.jpg", target_b64, "image/jpeg"),
            "doFaceEnhancer": doFaceEnhancer
        }

        # Создаем заголовки с API ключом
        headers = {
            "Content-Type": "multipart/form-data",
            "Authorization": f"Bearer {api_key}"
        }

        # Отправляем POST-запрос
        response = requests.post(api_url, files=payload)
        response.raise_for_status()

        # Получаем результат
        result = response.json()
        result_data = result["data"][0]  # Убедитесь, что индекс корректный
        result_b64 = base64.b64decode(result_data)

        # Преобразуем результат в изображение PIL
        output_image = Image.open(BytesIO(result_b64))
        return output_image

    except requests.exceptions.RequestException as e:
        print(f"Ошибка API: {e}")
        return None
    except Exception as e:
        print(f"Ошибка: {e}")
        return None

# Создаем интерфейс Gradio
iface = gr.Interface(
    fn=swap_face_api,
    inputs=[
        gr.Image(type="pil", label="Source Image"),
        gr.Image(type="pil", label="Target Image"),
        gr.Checkbox(label="Face Enhancer?")
    ],
    outputs=gr.Image(type="pil", label="Output Image"),
    title="Face Swap via API"
)

iface.launch(server_name="0.0.0.0", server_port=7860)  # Запускаем на всех интерфейсах