import cv2
import numpy as np
import gradio as gr
from inference_sdk import InferenceHTTPClient
from PIL import Image

# Initialize API client
CLIENT = InferenceHTTPClient(
    api_url="https://outline.roboflow.com",
    api_key="yu2FwTnkRok26cNiMo8e"
)

# Predefined color options
color_options = {
    "🔴 Red": (0, 0, 255),
    "🔵 Blue": (255, 0, 0),
    "🟢 Green": (0, 255, 0),
    "💖 Pink": (203, 192, 255),
    "💜 Purple": (128, 0, 128),
    "⚫ Black": (0, 0, 0),
    "⚪ White": (255, 255, 255)
}

# Function to apply nail paint
def apply_nail_paint(image, selected_color):
    image = np.array(image)
    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)

    result = CLIENT.infer(image, model_id="seg_nail_test/1")

    mask = np.zeros(image.shape[:2], dtype=np.uint8)
    for pred in result['predictions']:
        if 'points' in pred:
            points = np.array([[p['x'], p['y']] for p in pred['points']], dtype=np.int32)
            cv2.fillPoly(mask, [points], 255)

    mask = cv2.GaussianBlur(mask, (7, 7), 5)
    mask = cv2.erode(mask, None, iterations=1)
    mask = cv2.dilate(mask, None, iterations=2)
    mask_3d = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR)

    color_bgr = color_options[selected_color]
    color_nails = np.full_like(image, color_bgr, dtype=np.uint8)
    colored_nails = cv2.addWeighted(image, 0.4, color_nails, 0.6, 0)
    final_nails = np.where(mask_3d == 255, colored_nails, image)

    glossy_nails = cv2.addWeighted(final_nails, 0.9, mask_3d, 0.1, 0)
    glossy_nails = cv2.cvtColor(glossy_nails, cv2.COLOR_BGR2RGB)

    return Image.fromarray(glossy_nails)

# Function to apply nail texture
def apply_nail_texture(image, texture):
    image = np.array(image)
    texture = np.array(texture)

    result = CLIENT.infer(image, model_id="seg_nail_test/1")
    
    mask = np.zeros(image.shape[:2], dtype=np.uint8)
    for pred in result['predictions']:
        if 'points' in pred:
            points = np.array([[p['x'], p['y']] for p in pred['points']], dtype=np.int32)
            cv2.fillPoly(mask, [points], 255)

    mask = cv2.GaussianBlur(mask, (7, 7), 5)
    mask = cv2.erode(mask, None, iterations=1)
    mask = cv2.dilate(mask, None, iterations=2)
    mask_3d = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR)

    texture_resized = cv2.resize(texture, (image.shape[1], image.shape[0]))
    textured_nails = np.where(mask_3d == 255, texture_resized, image)
    # textured_nails = cv2.cvtColor(textured_nails, cv2.COLOR_BGR2RGB)

    return Image.fromarray(textured_nails)

# Function to handle dropdown selection
def process_image(mode, image, selected_color=None, texture=None):
    if mode == "Apply Paint":
        return apply_nail_paint(image, selected_color)
    elif mode == "Apply Texture":
        return apply_nail_texture(image, texture)

# Function to clear all inputs/outputs
def clear_inputs():
    return None, None, None, None

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## Nail Paint & Texture Application")
    
    mode = gr.Dropdown(
        ["Apply Paint", "Apply Texture"],
        label="Select Mode",
        value="Apply Paint"
    )

    with gr.Row():
        image_input = gr.Image(type="pil", label="Upload Hand Image")
        color_input = gr.Radio(choices=list(color_options.keys()), label="Select Nail Color", visible=True)
        texture_input = gr.Image(type="pil", label="Upload Texture Image", visible=False)

    output_image = gr.Image(type="pil", label="Output Image")

    def update_inputs(selected_mode):
        if selected_mode == "Apply Paint":
            return gr.update(visible=True), gr.update(visible=False)
        else:
            return gr.update(visible=False), gr.update(visible=True)

    mode.change(update_inputs, inputs=[mode], outputs=[color_input, texture_input])
    
    with gr.Row():
        process_btn = gr.Button("Process")
        clear_btn = gr.Button("Clear")

    process_btn.click(process_image, inputs=[mode, image_input, color_input, texture_input], outputs=output_image)
    clear_btn.click(clear_inputs, outputs=[image_input, color_input, texture_input, output_image])

demo.launch()