import os
import cv2
import gradio as gr
import numpy as np
import random

MAX_SEED = 999999
example_path = os.path.join(os.path.dirname(__file__), 'assets')

# Load example images
garm_list = os.listdir(os.path.join(example_path, "cloth"))
garm_list_path = [os.path.join(example_path, "cloth", garm) for garm in garm_list]
human_list = os.listdir(os.path.join(example_path, "human"))
human_list_path = [os.path.join(example_path, "human", human) for human in human_list]

def mock_tryon(person_img, garment_img, seed, randomize_seed):
    if person_img is None or garment_img is None:
        gr.Warning("Please upload both images!")
        return None, None, "Error: Empty image"
    
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    # Simple mock processing
    try:
        # Convert person to grayscale
        person_gray = cv2.cvtColor(person_img, cv2.COLOR_RGB2GRAY)
        person_gray = cv2.cvtColor(person_gray, cv2.COLOR_GRAY2RGB)
        
        # Resize garment
        garment_resized = cv2.resize(garment_img, (person_img.shape[1], person_img.shape[0]))
        
        # Blend images
        result = cv2.addWeighted(person_gray, 0.3, garment_resized, 0.7, 0)
        
        return result, seed, "Mock try-on complete"
    except Exception as e:
        gr.Error(f"Processing error: {str(e)}")
        return None, seed, f"Error: {str(e)}"

css = """
#col-left, #col-mid, #col-right { margin: 0 auto; max-width: 430px; }
#col-showcase { margin: 0 auto; max-width: 1100px; }
#button { color: blue; }
"""

with gr.Blocks(css=css) as app:
    gr.Markdown("# Virtual Try-On Demo (Mock)")
    
    with gr.Row():
        with gr.Column(elem_id="col-left"):
            gr.Markdown("### Step 1: Upload Person Image")
            person_img = gr.Image(label="Person", type="numpy")
            gr.Examples(examples=human_list_path, inputs=person_img)
        
        with gr.Column(elem_id="col-mid"):
            gr.Markdown("### Step 2: Upload Garment")
            garment_img = gr.Image(label="Garment", type="numpy")
            gr.Examples(examples=garm_list_path, inputs=garment_img)
        
        with gr.Column(elem_id="col-right"):
            gr.Markdown("### Step 3: Get Result")
            output_img = gr.Image(label="Result")
            with gr.Row():
                seed = gr.Slider(0, MAX_SEED, label="Seed", value=0)
                random_seed = gr.Checkbox(label="Random Seed", value=True)
            with gr.Row():
                seed_used = gr.Number(label="Seed Used")
                status = gr.Textbox(label="Status")
            run_btn = gr.Button("Run", variant="primary")

    run_btn.click(
        fn=mock_tryon,
        inputs=[person_img, garment_img, seed, random_seed],
        outputs=[output_img, seed_used, status]
    )

app.launch(share=True)  # Critical fix: Added share=True