import os
import requests
import gradio as gr
from dotenv import load_dotenv  # Chỉ dùng khi phát triển cục bộ
import io
from PIL import Image
from mtranslate import translate  # Import thư viện mtranslate
import random  # Thư viện để tạo yếu tố ngẫu nhiên

# Tải biến môi trường từ file .env nếu có (chỉ dùng khi phát triển cục bộ)
load_dotenv()

# Lấy API Key và URL từ biến môi trường
api_key = os.getenv("HF_API_KEY")
image_api_url = os.getenv("IMAGE_API_URL")
headers = {"Authorization": f"Bearer {api_key}"}

# Hàm gọi API Hugging Face để tạo hình ảnh
def query(payload):
    response = requests.post(image_api_url, headers=headers, json=payload)
    return response.content

# Hàm dịch tiếng Việt sang tiếng Anh bằng mtranslate
def translate_to_english(text):
    return translate(text, "en", "vi")  # Dịch từ tiếng Việt sang tiếng Anh

# Hàm xử lý để gọi API và trả về hình ảnh
def generate_image(prompt):
    # Dịch prompt từ tiếng Việt sang tiếng Anh
    prompt_en = translate_to_english(prompt)
    # Thêm yếu tố ngẫu nhiên vào prompt để tạo hình ảnh khác nhau
    random_factor = random.randint(1, 1000)
    prompt_with_randomness = f"{prompt_en}. Seed: {random_factor}"
    
    # Gửi prompt đã dịch vào API tạo hình ảnh
    image_bytes = query({"inputs": prompt_with_randomness})
    # Mở hình ảnh từ byte dữ liệu
    image = Image.open(io.BytesIO(image_bytes))
    return image

# Tạo giao diện Gradio
iface = gr.Interface(
    fn=generate_image,
    inputs="text",
    outputs="image",
    title="Image Generator - Trần Như Tuấn",
    description="Nhập mô tả của bạn để tạo hình ảnh"
)

# Khởi chạy ứng dụng
iface.launch()