import os
from PIL import Image
import gradio as gr
from gradio_client import Client, handle_file

HF_TOKEN = os.getenv("HF_TOKEN")

client = Client(
    "franciszzj/Leffa",
    hf_token=HF_TOKEN,
)

def rename_to_image(input_path, output_dir="Uploads"):
    os.makedirs(output_dir, exist_ok=True)
    random_id = os.urandom(4).hex()
    output_filename = f"Image_{random_id}.png"
    output_path = os.path.join(output_dir, output_filename)
    try:
        with Image.open(input_path) as img:
            img.convert('RGB').save(output_path, 'PNG', quality=100)
        return output_path
    except Exception as e:
        raise ValueError(f"Error renaming: {str(e)}")

def virtual_tryon(person_path, garment_path, garment_type):
    if not person_path or not garment_path:
        return None, "Ju lutem ngarkoni të dyja imazhet."

    # Rename uploads to Image_<random>.png
    person_png = rename_to_image(person_path)
    garment_png = rename_to_image(garment_path)

    person_file = handle_file(person_png)
    garment_file = handle_file(garment_png)

    try:
        result = client.predict(
            person_file,
            garment_file,
            ref_acceleration=False,
            step=40,
            scale=2.5,  # Corrected to 2.5 (was 42 in last version)
            seed=42,
            vt_model_type="viton_hd",
            vt_garment_type=garment_type,
            vt_repaint=False,
            api_name="/leffa_predict_vt",
        )
        # Rename result to uniqueproc_<random>.png
        input_image_path = result[0]
        random_id = os.urandom(4).hex()
        output_image_path = os.path.join("Outputs", f"uniqueproc_{random_id}.png")
        os.makedirs("Outputs", exist_ok=True)
        with Image.open(input_image_path) as img:
            img.convert('RGB').save(output_image_path, 'PNG', quality=95)
        return output_image_path, ""
    except Exception as e:
        return None, f"Error processing: {str(e)}"

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("")
    error_msg = gr.Markdown("", visible=False)
    with gr.Row():
        with gr.Column():
            src = gr.Image(sources="upload", type="filepath", label="Foto e personit")
        with gr.Column():
            ref = gr.Image(sources="upload", type="filepath", label="Foto e veshjes")
        with gr.Column():
            garment_type = gr.Radio(
                choices=[("Siper", "upper_body"),
                         ("Posht", "lower_body"),
                         ("Komplet", "dresses")],
                value="upper_body",
                label="Lloji i veshjes",
            )
        with gr.Column():
            btn = gr.Button("VISHU")
            out = gr.Image(type="filepath", label="Rezultati")

    btn.click(
        virtual_tryon,
        inputs=[src, ref, garment_type],
        outputs=[out, error_msg]
    )

demo.launch(
    share=True,
    show_error=True,
    pwa=True,
)