import os
import requests
import gradio as gr
from dotenv import load_dotenv
import io
from PIL import Image
from mtranslate import translate
import random

# Tải biến môi trường từ file .env
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}"}

# Danh sách màu sắc
color_map = {
    "Đỏ": "red",
    "Cam": "orange",
    "Vàng": "yellow",
    "Lục": "green",
    "Lam": "blue",
    "Chàm": "indigo",
    "Tím": "violet",
    "Trắng": "white",
    "Đen": "black",
    "Xám": "gray",
    "Cầu vồng": "rainbow",
}

# Danh sách kiểu tóc
hair_styles = [
    "Tóc tự nhiên", "Tóc thẳng mượt", "Tóc buộc nửa đầu", "Tóc búi cao", "Tóc bob", 
    "Tóc xoăn sóng", "Tóc tết", "Tóc pixie"
]

# 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)
    if response.status_code != 200:
        return f"Error: {response.status_code}, {response.text}"
    return response.content

# Hàm dịch tiếng Việt sang tiếng Anh
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ý tạo hình ảnh với các lựa chọn
def generate_image(prompt, view, style, hair_color, eye_color, lip_color, hair_style):
    # Dịch prompt từ tiếng Việt sang tiếng Anh
    prompt_en = translate_to_english(prompt)
    
    # Thêm các yếu tố lựa chọn vào prompt
    prompt_with_choices = f"{prompt_en}, View: {view}, Style: {style}, Hair color: {hair_color}, Eye color: {eye_color}, Lip color: {lip_color}, Hair style: {hair_style}"

    # Thêm yếu tố ngẫu nhiên vào prompt
    random_factor = random.randint(1, 1000)
    prompt_with_randomness = f"{prompt_with_choices}. Seed: {random_factor}"
    
    # Gửi prompt 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",  # Prompt
        gr.Dropdown(["Cận cảnh", "Nửa người", "Toàn thân"], label="View"),
        gr.Dropdown(["Art", "Anime", "Realistic"], label="Style"),
        gr.Dropdown(
            ["Đỏ", "Cam", "Vàng", "Lục", "Lam", "Chàm", "Tím", "Trắng", "Đen", "Xám", "Cầu vồng"],
            label="Màu tóc"
        ),
        gr.Dropdown(
            ["Đỏ", "Cam", "Vàng", "Lục", "Lam", "Chàm", "Tím", "Trắng", "Đen", "Xám", "Cầu vồng"],
            label="Màu mắt"
        ),
        gr.Dropdown(
            ["Đỏ", "Cam", "Vàng", "Lục", "Lam", "Chàm", "Tím", "Trắng", "Đen", "Xám", "Cầu vồng"],
            label="Màu môi"
        ),
        gr.Dropdown(
            hair_styles, 
            label="Kiểu tóc"
        ),
    ],
    outputs="image",
    title="Image Generator - Trần Như Tuấn",
    description="Chọn các tùy chọn và nhập mô tả để tạo hình ảnh chân dung"
)

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