import os
import cv2
import gradio as gr
import numpy as np
import random
import base64
import requests
import json
import time
import os
import gradio as gr
import subprocess
from shutil import copyfile

# Base directory
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
DATA_DIR = os.path.join(BASE_DIR, "..", "data", "test")
RESULT_IMAGE_PATH = os.path.join(BASE_DIR, "..", "results", "test", "tryon", "user_result.jpg")

# Ensure required directories exist
os.makedirs(os.path.join(DATA_DIR, "image"), exist_ok=True)
os.makedirs(os.path.join(DATA_DIR, "cloth"), exist_ok=True)

def virtual_tryon(person_image, cloth_image):
    # Save user inputs to test data folder
    person_path = os.path.join(DATA_DIR, "image", "user_image.jpg")
    cloth_path = os.path.join(DATA_DIR, "cloth", "user_cloth.jpg")
    
    person_image.save(person_path)
    cloth_image.save(cloth_path)

    # Update test.txt with filenames
    test_list_path = os.path.join(DATA_DIR, "test.txt")
    with open(test_list_path, "w") as f:
        f.write("user_image.jpg user_cloth.jpg\n")

    # Call test.py
    result = subprocess.run(
        ["python", os.path.join(BASE_DIR, "..", "test.py")],
        cwd=os.path.join(BASE_DIR, ".."),
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
    )

    print("STDOUT:", result.stdout)
    print("STDERR:", result.stderr)

    # Check if result image exists
    if os.path.exists(RESULT_IMAGE_PATH):
        return RESULT_IMAGE_PATH
    else:
        return "❌ Error: Try-on result not found."

# Gradio UI
demo = gr.Interface(
    fn=virtual_tryon,
    inputs=[
        gr.Image(type="pil", label="Upload Person Image"),
        gr.Image(type="pil", label="Upload Clothing Image")
    ],
    outputs=gr.Image(type="filepath", label="Try-On Result"),
    title="Virtual Try-On with VITON-HD",
    description="Upload a person and a clothing image to try-on the clothes virtually!"
)

if __name__ == "__main__":
    demo.launch()