import os
import cv2
import numpy as np
import json
import random
from PIL import Image, ImageDraw, ImageFont
import asyncio

import requests
import base64
import gradio as gr

machine_number = 0
model = os.path.join(os.path.dirname(__file__), "models/female_model.png")

MODEL_MAP = {
    "AI Model female_model": 'models/female_model6.png',
    "AI Model female_model1": 'models/female_model1.png',
    "AI Model male_model": 'models/male_model3.png',
    "AI Model male_model2": 'models/male_model2.png',
}

# Sample clothing items from your clothes folder
SAMPLE_CLOTHES = {
    "top": [
        "clothes/shirt.jpg",
        "clothes/cardigan.jpg", 
        "clothes/jacket.jpg",
        "clothes/dress.jpg"
    ],
    "bottom": [
        "clothes/pants.jpg",
        "clothes/jeans.jpg",
        "clothes/skirt.jpg", 
        "clothes/shorts.jpg"
    ]
}

def add_waterprint(img):
    h, w, _ = img.shape
    img = cv2.putText(img, 'Powered by OutfitAnyone', (int(0.3*w), h-20), cv2.FONT_HERSHEY_PLAIN, 2, (128, 128, 128), 2, cv2.LINE_AA)
    return img

def load_sample_image(image_path):
    """Load a sample clothing image and return it as numpy array"""
    try:
        if os.path.exists(image_path):
            img = cv2.imread(image_path)
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            return img
        else:
            # Return a placeholder if image doesn't exist
            return np.zeros((400, 400, 3), dtype=np.uint8)
    except Exception as e:
        print(f"Error loading image {image_path}: {e}")
        return np.zeros((400, 400, 3), dtype=np.uint8)

def load_model_image(model_path):
    """Load a model image and return the file path"""
    try:
        if os.path.exists(model_path):
            return model_path
        else:
            print(f"Model image not found: {model_path}")
            return model  # fallback to default
    except Exception as e:
        print(f"Error loading model {model_path}: {e}")
        return model  # fallback to default

def get_tryon_result(model_name, garment1, garment2, seed=1234):
    # model_name = "AI Model " + model_name.split("\\")[-1].split(".")[0] # windows
    model_name = "AI Model " + model_name.split("/")[-1].split(".")[0] # linux
    print(model_name)

    encoded_garment1 = cv2.imencode('.jpg', garment1)[1].tobytes()
    encoded_garment1 = base64.b64encode(encoded_garment1).decode('utf-8')

    if garment2 is not None:
        encoded_garment2 = cv2.imencode('.jpg', garment2)[1].tobytes()
        encoded_garment2 = base64.b64encode(encoded_garment2).decode('utf-8')
    else:
        encoded_garment2 = ''

    # Fix for the missing environment variable
    try:
        url = os.environ['OA_IP_ADDRESS']
    except KeyError:
        print("Error: OA_IP_ADDRESS environment variable not set!")
        # Return a dummy image or handle the error appropriately
        dummy_img = np.zeros((512, 512, 3), dtype=np.uint8)
        dummy_img = cv2.putText(dummy_img, 'Error: OA_IP_ADDRESS not set', (50, 256), 
                               cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
        return add_waterprint(dummy_img)
    
    headers = {'Content-Type': 'application/json'}
    seed = random.randint(0, 1222222222)
    data = {
        "garment1": encoded_garment1,
        "garment2": encoded_garment2,
        "model_name": model_name,
        "seed": seed
    }
    
    try:
        response = requests.post(url, headers=headers, data=json.dumps(data))
        print("response code", response.status_code)
        if response.status_code == 200:
            result = response.json()
            result = base64.b64decode(result['images'][0])
            result_np = np.frombuffer(result, np.uint8)
            result_img = cv2.imdecode(result_np, cv2.IMREAD_UNCHANGED)
        else:
            print('server error!')
            # Return error image
            result_img = np.zeros((512, 512, 3), dtype=np.uint8)
            result_img = cv2.putText(result_img, f'Server Error: {response.status_code}', 
                                   (50, 256), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
    except Exception as e:
        print(f"Request error: {e}")
        result_img = np.zeros((512, 512, 3), dtype=np.uint8)
        result_img = cv2.putText(result_img, 'Connection Error', (50, 256), 
                               cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)

    final_img = add_waterprint(result_img)
    return final_img


with gr.Blocks(css = ".output-image, .input-image, .image-preview {height: 400px !important} ") as demo:
    gr.HTML(
        """
        <div style="display: flex; justify-content: center; align-items: center; text-align: center;">
        <a href="https://github.com/HumanAIGC/OutfitAnyone" style="margin-right: 20px; text-decoration: none; display: flex; align-items: center;">
        </a>
        <div>
            <h1 >Outfit Anyone: Ultra-high quality virtual try-on for Any Clothing and Any Person</h1>
            <h4 >v0.9</h4>
            <h5 style="margin: 0;">OutfitAnyone plus version is now online with any model and any cloth: https://www.outfitanyone.life/</h5>
            <div style="display: flex; justify-content: center; align-items: center; text-align: center;>
                <a href="https://arxiv.org/abs/2407.16224"><img src="https://img.shields.io/badge/Arxiv-2407.16224-red"></a>
                <a href='https://humanaigc.github.io/outfit-anyone/'><img src='https://img.shields.io/badge/Project_Page-OutfitAnyone-green' alt='Project Page'></a>
                <a href='https://github.com/HumanAIGC/OutfitAnyone'><img src='https://img.shields.io/badge/Github-Repo-blue'></a>
            </div>
        </div>
        </div>
        """)
    with gr.Row():
        with gr.Column():
            # Model selection buttons with preview images
            gr.HTML("<h3>Select AI Model:</h3>")
            with gr.Row():
                with gr.Column():
                    model_preview1 = gr.Image(value=MODEL_MAP["AI Model female_model"], 
                                            label="Female Model", 
                                            height=150, width=100, 
                                            interactive=False, show_label=True)
                    model_btn1 = gr.Button("Select Female Model", variant="secondary")
                
                with gr.Column():
                    model_preview2 = gr.Image(value=MODEL_MAP["AI Model female_model1"], 
                                            label="Female Model 1", 
                                            height=150, width=100, 
                                            interactive=False, show_label=True)
                    model_btn2 = gr.Button("Select Female Model 1", variant="secondary")
                
                with gr.Column():
                    model_preview3 = gr.Image(value=MODEL_MAP["AI Model male_model"], 
                                            label="Male Model", 
                                            height=150, width=100, 
                                            interactive=False, show_label=True)
                    model_btn3 = gr.Button("Select Male Model", variant="secondary")
                
                with gr.Column():
                    model_preview4 = gr.Image(value=MODEL_MAP["AI Model male_model2"], 
                                            label="Male Model 2", 
                                            height=150, width=100, 
                                            interactive=False, show_label=True)
                    model_btn4 = gr.Button("Select Male Model 2", variant="secondary")
            
            init_image = gr.Image(sources='clipboard', type="filepath", label="Selected Model", value=model)
            
        with gr.Column():
            gr.HTML(
                """
                <div style="display: flex; justify-content: center; align-items: center; text-align: center;">
                <div>
                    <h3>Models are fixed and cannot be uploaded or modified; we only support users uploading their own garments.</h3>
                    <h4 style="margin: 0;">For a one-piece dress or coat, you only need to upload the image to the 'top garment' section and leave the 'lower garment' section empty.</h4>
                </div>
                </div>
                """)
            
            # Sample clothing buttons section with preview images
            gr.HTML("<h3>Quick Select Sample Clothes:</h3>")
            
            gr.HTML("<h4>Top Garments:</h4>")
            with gr.Row():
                with gr.Column():
                    top_preview1 = gr.Image(value=SAMPLE_CLOTHES["top"][0], 
                                          label="Shirt", 
                                          height=150, width=100, 
                                          interactive=False, show_label=True)
                    top_btn1 = gr.Button("Select Shirt", variant="secondary")
                
                with gr.Column():
                    top_preview2 = gr.Image(value=SAMPLE_CLOTHES["top"][1], 
                                          label="Cardigan", 
                                          height=150, width=100, 
                                          interactive=False, show_label=True)
                    top_btn2 = gr.Button("Select Cardigan", variant="secondary")
                
                with gr.Column():
                    top_preview3 = gr.Image(value=SAMPLE_CLOTHES["top"][2], 
                                          label="Jacket", 
                                          height=150, width=100, 
                                          interactive=False, show_label=True)
                    top_btn3 = gr.Button("Select Jacket", variant="secondary")
                
                with gr.Column():
                    top_preview4 = gr.Image(value=SAMPLE_CLOTHES["top"][3], 
                                          label="Dress", 
                                          height=150, width=100, 
                                          interactive=False, show_label=True)
                    top_btn4 = gr.Button("Select Dress", variant="secondary")
            
            gr.HTML("<h4>Bottom Garments:</h4>")
            with gr.Row():
                with gr.Column():
                    bottom_preview1 = gr.Image(value=SAMPLE_CLOTHES["bottom"][0], 
                                             label="Pants", 
                                             height=150, width=100, 
                                             interactive=False, show_label=True)
                    bottom_btn1 = gr.Button("Select Pants", variant="secondary")
                
                with gr.Column():
                    bottom_preview2 = gr.Image(value=SAMPLE_CLOTHES["bottom"][1], 
                                             label="Jeans", 
                                             height=150, width=100, 
                                             interactive=False, show_label=True)
                    bottom_btn2 = gr.Button("Select Jeans", variant="secondary")
                
                with gr.Column():
                    bottom_preview3 = gr.Image(value=SAMPLE_CLOTHES["bottom"][2], 
                                             label="Skirt", 
                                             height=150, width=100, 
                                             interactive=False, show_label=True)
                    bottom_btn3 = gr.Button("Select Skirt", variant="secondary")
                
                with gr.Column():
                    bottom_preview4 = gr.Image(value=SAMPLE_CLOTHES["bottom"][3], 
                                             label="Shorts", 
                                             height=150, width=100, 
                                             interactive=False, show_label=True)
                    bottom_btn4 = gr.Button("Select Shorts", variant="secondary")

            with gr.Row():
                garment_top = gr.Image(sources='upload', type="numpy", label="top garment")
                garment_down = gr.Image(sources='upload', type="numpy", label="lower garment")

            run_button = gr.Button(value="Run")
            
        with gr.Column():
            gallery = gr.Image()

    # Connect model buttons to load model images
    model_btn1.click(
        lambda: load_model_image(MODEL_MAP["AI Model female_model"]),
        outputs=[init_image]
    )
    model_btn2.click(
        lambda: load_model_image(MODEL_MAP["AI Model female_model1"]),
        outputs=[init_image]
    )
    model_btn3.click(
        lambda: load_model_image(MODEL_MAP["AI Model male_model"]),
        outputs=[init_image]
    )
    model_btn4.click(
        lambda: load_model_image(MODEL_MAP["AI Model male_model2"]),
        outputs=[init_image]
    )

    # Connect buttons to load sample images
    top_btn1.click(
        lambda: load_sample_image(SAMPLE_CLOTHES["top"][0]),
        outputs=[garment_top]
    )
    top_btn2.click(
        lambda: load_sample_image(SAMPLE_CLOTHES["top"][1]), 
        outputs=[garment_top]
    )
    top_btn3.click(
        lambda: load_sample_image(SAMPLE_CLOTHES["top"][2]),
        outputs=[garment_top]
    )
    top_btn4.click(
        lambda: load_sample_image(SAMPLE_CLOTHES["top"][3]),
        outputs=[garment_top]
    )
    
    bottom_btn1.click(
        lambda: load_sample_image(SAMPLE_CLOTHES["bottom"][0]),
        outputs=[garment_down]
    )
    bottom_btn2.click(
        lambda: load_sample_image(SAMPLE_CLOTHES["bottom"][1]),
        outputs=[garment_down] 
    )
    bottom_btn3.click(
        lambda: load_sample_image(SAMPLE_CLOTHES["bottom"][2]),
        outputs=[garment_down]
    )
    bottom_btn4.click(
        lambda: load_sample_image(SAMPLE_CLOTHES["bottom"][3]),
        outputs=[garment_down]
    )

    run_button.click(fn=get_tryon_result, 
                     inputs=[
                            init_image,
                            garment_top,
                            garment_down,
                            ], 
                     outputs=[gallery],
                     concurrency_limit=2)

if __name__ == "__main__":
    ip = requests.get('http://ifconfig.me/ip', timeout=1).text.strip()
    print("ip address alibaba", ip)
    demo.queue(max_size=10)
    demo.launch()