import os
import numpy as np
import base64
import json
import random
import cv2
from PIL import Image
import requests
import gradio as gr

# In-memory user storage for demonstration
users_db = {}

# Ensure the model path is correct
current_dir = os.path.dirname(os.path.abspath(__file__))
default_model = os.path.join(current_dir, "models/eva/Eva_0.png")

# Map of AI models with their corresponding file paths
MODEL_MAP = {
    "AI Model Rouyan_0": os.path.join(current_dir, 'models/rouyan_new/Rouyan_0.png'),
    "AI Model Eva_0": os.path.join(current_dir, 'models/eva/Eva_0.png'),
    "AI Model Eva_1": os.path.join(current_dir, 'models/eva/Eva_1.png'),
    "AI Model Simon_0": os.path.join(current_dir, 'models/simon_online/Simon_0.png'),
    "AI Model Simon_1": os.path.join(current_dir, 'models/simon_online/Simon_1.png'),
    "AI Model Xuanxuan_0": os.path.join(current_dir, 'models/xiaoxuan_online/Xuanxuan_0.png'),
    "AI Model Xuanxuan_1": os.path.join(current_dir, 'models/xiaoxuan_online/Xuanxuan_1.png'),
    "AI Model Xuanxuan_2": os.path.join(current_dir, 'models/xiaoxuan_online/Xuanxuan_2.png'),
    "AI Model Yaqi_0": os.path.join(current_dir, 'models/yaqi/Yaqi_0.png'),
    "AI Model Yaqi_1": os.path.join(current_dir, 'models/yaqi/Yaqi_1.png'),
    "AI Model Yaqi_2": os.path.join(current_dir, 'models/yaqi/Yaqi_2.png'),
    "AI Model Yaqi_3": os.path.join(current_dir, 'models/yaqi/Yaqi_3.png'),
    "AI Model Yifeng_0": os.path.join(current_dir, 'models/yifeng_online/Yifeng_0.png'),
    "AI Model Yifeng_1": os.path.join(current_dir, 'models/yifeng_online/Yifeng_1.png'),
    "AI Model Yifeng_2": os.path.join(current_dir, 'models/yifeng_online/Yifeng_2.png'),
    "AI Model Yifeng_3": os.path.join(current_dir, 'models/yifeng_online/Yifeng_3.png'),
}

# Registration function
def register(username, password, confirm_password):
    if username in users_db:
        return "Username already exists", gr.update(visible=True), gr.update(visible=False)
    if password != confirm_password:
        return "Passwords do not match", gr.update(visible=True), gr.update(visible=False)
    users_db[username] = password
    return "Registration successful! Please login.", gr.update(visible=True), gr.update(visible=False)

# Login function
def login(username, password):
    if users_db.get(username) == password:
        return gr.update(visible=False), gr.update(visible=True)
    else:
        return gr.update(visible=True), gr.update(visible=False)

# Watermark function
def add_watermark(image):
    height, width, _ = image.shape
    cv2.putText(image, 'Powered by OutfitAnyone', (int(0.3 * width), height - 20),
                cv2.FONT_HERSHEY_PLAIN, 2, (128, 128, 128), 2, cv2.LINE_AA)
    return image

# Try-on result processing
def get_tryon_result(model_name, top_garment, bottom_garment=None):
    model_key = "AI Model " + model_name.split("/")[-1].split(".")[0]
    print(f"Selected model: {model_key}")

    encoded_top = base64.b64encode(cv2.imencode('.jpg', top_garment)[1].tobytes()).decode('utf-8')
    
    # Check if bottom_garment is not None and is not an empty array
    encoded_bottom = ""
    if bottom_garment is not None and bottom_garment.size > 0:
        encoded_bottom = base64.b64encode(cv2.imencode('.jpg', bottom_garment)[1].tobytes()).decode('utf-8')

    server_url = os.environ.get('OA_IP_ADDRESS', 'http://localhost:5000')
    headers = {'Content-Type': 'application/json'}
    payload = {
        "garment1": encoded_top,
        "garment2": encoded_bottom,
        "model_name": model_key,
        "seed": random.randint(0, 99999999)
    }

    try:
        response = requests.post(server_url, headers=headers, data=json.dumps(payload))

        if response.status_code == 200:
            result = response.json()
            result_img = cv2.imdecode(np.frombuffer(base64.b64decode(result['images'][0]), np.uint8), cv2.IMREAD_UNCHANGED)
            final_img = add_watermark(result_img)
            return final_img
        else:
            print(f"Error: Server responded with status code {response.status_code}")
            return None
    except requests.exceptions.ConnectionError as e:
        print(f"ConnectionError: Could not connect to server at {server_url}. Please ensure the server is running.")
        print(f"Details: {e}")
        return None
    except requests.exceptions.RequestException as e:
        print(f"RequestException: An error occurred during the request.")
        print(f"Details: {e}")
        return None

# Gradio Interface
with gr.Blocks() as demo:
    gr.HTML("<script>console.log('Custom JavaScript');</script>")
    login_block = gr.Column(visible=True)
    register_block = gr.Column(visible=False)
    main_app = gr.Column(visible=False)
    
    # Registration Page
    with register_block:
        gr.HTML("<h2>Register for Outfit Anyone</h2>")
        username = gr.Textbox(label="Username")
        password = gr.Textbox(label="Password", type="password", elem_id="reg_pass")
        confirm_password = gr.Textbox(label="Confirm Password", type="password", elem_id="reg_conf_pass")
        register_button = gr.Button("Register")
        back_to_login_button = gr.Button("Back to Login")
        register_message = gr.Textbox(label="", visible=False)

    # Login Page
    with login_block:
        gr.HTML("<h2>Login to Outfit Anyone</h2>")
        username_login = gr.Textbox(label="Username")
        password_login = gr.Textbox(label="Password", type="password", elem_id="login_pass")
        login_button = gr.Button("Login")
        go_to_register_button = gr.Button("Register")
        login_message = gr.Textbox(label="", visible=False)
    
    # Main App Interface
    with main_app:
        gr.HTML("""
            <div style="text-align: center;">
                <h1>Outfit Anyone: Virtual Try-On</h1>
                <h4>v1.0</h4>
                <p>Upload your garments and choose a model to see the virtual try-on.</p>
            </div>
        """)
        with gr.Row():
            with gr.Column():
                model_selector = gr.Image(sources='upload', type="filepath", label="Model", value=default_model)
                example_models = gr.Examples(inputs=model_selector, examples=[MODEL_MAP['AI Model Rouyan_0'], MODEL_MAP['AI Model Eva_0']], examples_per_page=4)
            with gr.Column():
                gr.HTML("<h3>Select Garments for Virtual Try-On</h3>")
                top_garment_input = gr.Image(sources='upload', type="numpy", label="Top Garment")
                bottom_garment_input = gr.Image(sources='upload', type="numpy", label="Bottom Garment (Optional)")
                generate_button = gr.Button(value="Generate Outfit")
            with gr.Column():
                result_display = gr.Image()

                generate_button.click(fn=get_tryon_result, inputs=[model_selector, top_garment_input, bottom_garment_input], outputs=[result_display])

    # Button Click Actions
    register_button.click(fn=register, inputs=[username, password, confirm_password], outputs=[register_message, register_block, login_block])
    back_to_login_button.click(lambda: (gr.update(visible=False), gr.update(visible=True)), None, [register_block, login_block])
    go_to_register_button.click(lambda: (gr.update(visible=True), gr.update(visible=False)), None, [register_block, login_block])
    login_button.click(fn=login, inputs=[username_login, password_login], outputs=[login_block, main_app])
    
if __name__ == "__main__":
    demo.queue(max_size=10)
    demo.launch(server_name="0.0.0.0", server_port=8080)