#!/usr/bin/env python3
"""
Virtual Makeup App using MediaPipe FaceMesh
Simple approach: landmarks-based masks for lipstick, blush, foundation
"""

import cv2
import numpy as np
import gradio as gr
from PIL import Image
import os
from landmarks import detect_landmarks, normalize_landmarks
from utils import mask_skin, gamma_correction


# Color presets for lipstick
LIPSTICK_COLORS = {
    "Red": (0, 0, 255),
    "Pink": (203, 192, 255),
    "Burgundy": (75, 0, 130),
    "Orange": (0, 165, 255),
    "Nude": (180, 140, 200),
    "Wine": (0, 0, 120),
    "Coral": (100, 165, 255),
}

# Color presets for blush
BLUSH_COLORS = {
    "Coral": (66, 135, 245),
    "Pink": (192, 135, 220),
    "Peach": (152, 165, 255),
    "Rose": (100, 70, 200),
    "Berry": (140, 50, 150),
    "Apricot": (140, 140, 230),
}

# Landmark indices for different regions
UPPER_LIP = [61, 185, 40, 39, 37, 0, 267, 269, 270, 408, 415, 272, 271, 268, 12, 38, 41, 42, 191, 78, 76]
LOWER_LIP = [61, 146, 91, 181, 84, 17, 314, 405, 320, 307, 308, 324, 318, 402, 317, 14, 87, 178, 88, 95]
CHEEKS = [425, 205]  # Left and right cheek points

# Foundation shade presets (BGR adjustments)
FOUNDATION_SHADES = {
    "Fair": {"shift": (-10, -5, 0), "description": "Very light skin tones"},
    "Light": {"shift": (-5, 0, 5), "description": "Light skin tones"},
    "Medium": {"shift": (0, 5, 10), "description": "Medium skin tones"},
    "Tan": {"shift": (5, 10, 15), "description": "Tan skin tones"},
    "Deep": {"shift": (10, 15, 20), "description": "Deep skin tones"},
    "Rich": {"shift": (15, 20, 25), "description": "Rich, dark skin tones"},
}

# Coverage level presets
COVERAGE_LEVELS = {
    "Light": {
        "intensity": 0.25,
        "gamma": 1.15,
        "smoothing": 5,
        "description": "Sheer coverage for natural look",
    },
    "Medium": {
        "intensity": 0.45,
        "gamma": 1.30,
        "smoothing": 7,
        "description": "Balanced coverage for everyday wear",
    },
    "Full": {
        "intensity": 0.70,
        "gamma": 1.50,
        "smoothing": 9,
        "description": "Maximum coverage for flawless finish",
    },
}

# Blending options
BLENDING_MODES = {
    "Natural": {"blur_kernel": 11, "feather": 0.9, "description": "Soft, seamless blend"},
    "Buildable": {"blur_kernel": 7, "feather": 0.7, "description": "Layered, controlled blend"},
    "Airbrushed": {"blur_kernel": 15, "feather": 0.95, "description": "Ultra-smooth, flawless blend"},
}

def apply_lipstick(image, color_rgb, landmarks, alpha=0.4):
    """Apply lipstick to lips using landmarks."""
    if landmarks is None:
        return image
    
    h, w = image.shape[:2]
    mask = np.zeros_like(image)
    
    lip_points = normalize_landmarks(landmarks, h, w, UPPER_LIP + LOWER_LIP)
    if len(lip_points) > 0:
        lip_points = lip_points.astype(np.int32)
        cv2.fillPoly(mask, [lip_points], color_rgb)
        # Smooth edges with Gaussian blur
        mask = cv2.GaussianBlur(mask, (15, 15), 3)
    
    # Blend with original image
    output = cv2.addWeighted(image, 1.0, mask, alpha, 0)
    return output


def apply_blush(image, color_rgb, landmarks, intensity=0.3, radius=40):
    """Apply blush to cheeks with intensity control (0-1)."""
    if landmarks is None:
        return image
    
    h, w = image.shape[:2]
    mask = np.zeros_like(image)
    
    cheek_points = normalize_landmarks(landmarks, h, w, CHEEKS)
    
    for point in cheek_points:
        x, y = int(point[0]), int(point[1])
        # Create circular blush with gradient
        y_min = max(0, y - radius)
        y_max = min(h, y + radius + 1)
        x_min = max(0, x - radius)
        x_max = min(w, x + radius + 1)
        
        yy, xx = np.ogrid[y_min:y_max, x_min:x_max]
        dist = np.sqrt((yy - y) ** 2 + (xx - x) ** 2)
        
        # Smooth gradient with cosine falloff
        gradient = np.zeros_like(dist, dtype=np.float32)
        valid = dist <= radius
        gradient[valid] = (1.0 + np.cos(np.pi * dist[valid] / radius)) / 2.0
        
        # Apply color with gradient
        for c in range(3):
            color_val = color_rgb[c]
            mask[y_min:y_max, x_min:x_max, c] = np.maximum(
                mask[y_min:y_max, x_min:x_max, c],
                (color_val * gradient).astype(np.uint8)
            )
    
    # Blur for smooth blending
    blur_rad = max(3, radius // 3)
    if blur_rad % 2 == 0:
        blur_rad += 1
    mask = cv2.GaussianBlur(mask, (blur_rad, blur_rad), blur_rad // 2)
    
    # Apply with intensity control
    alpha = intensity * 0.5  # Scale down so it's subtle
    output = cv2.addWeighted(image, 1.0, mask, alpha, 0)
    return output


def auto_detect_shade(image):
    """Analyze skin tone to suggest the best foundation shade."""
    if image is None:
        return "Medium"

    skin_mask_binary = mask_skin(image)
    if skin_mask_binary.ndim == 3:
        skin_mask_binary = skin_mask_binary[:, :, 0]

    skin_pixels = image[skin_mask_binary > 0]
    if len(skin_pixels) == 0:
        return "Medium"

    avg_b, avg_g, avg_r = np.mean(skin_pixels, axis=0)
    luminance = 0.299 * avg_r + 0.587 * avg_g + 0.114 * avg_b

    if luminance < 100:
        return "Rich"
    elif luminance < 130:
        return "Deep"
    elif luminance < 160:
        return "Tan"
    elif luminance < 190:
        return "Medium"
    elif luminance < 220:
        return "Light"
    else:
        return "Fair"


def apply_foundation_advanced(image, shade_name="Medium", coverage="Medium", blending="Natural", warmth_adjust=0):
    """Apply foundation with advanced controls."""
    if image is None:
        return image

    skin_mask_binary = mask_skin(image)
    if skin_mask_binary.ndim == 3:
        skin_mask_binary = skin_mask_binary[:, :, 0]

    shade_info = FOUNDATION_SHADES.get(shade_name, FOUNDATION_SHADES["Medium"])
    coverage_info = COVERAGE_LEVELS.get(coverage, COVERAGE_LEVELS["Medium"])
    blend_info = BLENDING_MODES.get(blending, BLENDING_MODES["Natural"])

    intensity = coverage_info["intensity"]
    gamma_val = coverage_info["gamma"]
    smoothing = coverage_info["smoothing"]
    blur_kernel = blend_info["blur_kernel"]
    feather = blend_info["feather"]

    corrected = gamma_correction(image, gamma_val, coefficient=1)
    corrected = corrected.astype(np.float32)
    b_shift, g_shift, r_shift = shade_info["shift"]
    corrected[:, :, 0] = np.clip(corrected[:, :, 0] + b_shift, 0, 255)
    corrected[:, :, 1] = np.clip(corrected[:, :, 1] + g_shift, 0, 255)
    corrected[:, :, 2] = np.clip(corrected[:, :, 2] + r_shift, 0, 255)

    if warmth_adjust != 0:
        warmth_factor = 1.0 + (warmth_adjust / 100.0)
        corrected[:, :, 2] = np.clip(corrected[:, :, 2] * warmth_factor, 0, 255)

    corrected = corrected.astype(np.uint8)

    if smoothing > 0:
        corrected = cv2.bilateralFilter(corrected, smoothing, 75, 75)

    skin_mask_float = skin_mask_binary.astype(np.float32)
    if blur_kernel % 2 == 0:
        blur_kernel += 1
    skin_mask_float = cv2.GaussianBlur(skin_mask_float, (blur_kernel, blur_kernel), 0)
    skin_mask_float = np.power(skin_mask_float, feather)
    skin_mask_float = np.expand_dims(skin_mask_float, axis=-1)
    skin_mask_float = np.repeat(skin_mask_float, 3, axis=2)

    output = image.astype(np.float32)
    corrected = corrected.astype(np.float32)
    output = output * (1.0 - skin_mask_float * intensity) + corrected * (skin_mask_float * intensity)

    return output.astype(np.uint8)


def process_foundation(image, shade, warmth):
    """Process image to apply foundation."""
    if image is None:
        return None, "No image provided", ""

    if isinstance(image, Image.Image):
        img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
    else:
        img = image.copy()

    coverage = "Full"
    blending = "Natural"
    output = apply_foundation_advanced(img, shade, coverage, blending, warmth)
    output_rgb = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)

    coverage_info = COVERAGE_LEVELS.get(coverage, COVERAGE_LEVELS["Medium"])
    status = (
        f"Applied {shade} foundation with {coverage} coverage "
        f"({coverage_info['description']}) using {blending} blending"
    )

    return output_rgb, status, shade


def process_image(image, apply_lipstick_flag, lipstick_color,
                  apply_blush_flag, blush_color, blush_intensity,
                  apply_foundation_flag, foundation_shade, warmth):
    """Process image with selected makeup features."""
    
    if image is None:
        return image, "No image provided"
    
    # Convert PIL to OpenCV format
    if isinstance(image, Image.Image):
        img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
    else:
        img = image.copy()
    
    # Detect landmarks
    landmarks = detect_landmarks(img)
    if landmarks is None:
        return cv2.cvtColor(img, cv2.COLOR_BGR2RGB), "No face detected"
    
    output = img.copy()
    applied_features = []
    
    # Apply selected features
    if apply_lipstick_flag:
        color = LIPSTICK_COLORS.get(lipstick_color, LIPSTICK_COLORS["Red"])
        output = apply_lipstick(output, color, landmarks, alpha=0.4)
        applied_features.append("Lipstick")
    
    if apply_blush_flag:
        color = BLUSH_COLORS.get(blush_color, BLUSH_COLORS["Pink"])
        # Blush intensity is percentage (0-100), convert to 0-1
        intensity = blush_intensity / 100.0
        output = apply_blush(output, color, landmarks, intensity=intensity)
        applied_features.append(f"Blush ({blush_intensity}%)")
    
    if apply_foundation_flag:
        foundation_image, foundation_status, used_shade = process_foundation(
            output, foundation_shade, warmth
        )
        output = cv2.cvtColor(foundation_image, cv2.COLOR_RGB2BGR)
        applied_features.append(f"Foundation ({used_shade}, Full, Natural)")
    
    # Convert back to RGB for display
    output_rgb = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
    
    status = f"Applied: {', '.join(applied_features) if applied_features else 'None'}"
    return output_rgb, status

# Create Gradio interface
with gr.Blocks(title="Virtual Makeup App") as demo:
    gr.Markdown("# 💄 Virtual Makeup App")
    gr.Markdown("Upload a photo and apply virtual makeup with customizable intensity")
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(label="Upload Image", type="pil")
            
            # Lipstick controls
            lipstick_check = gr.Checkbox(label="Apply Lipstick", value=True)
            lipstick_color = gr.Dropdown(
                choices=list(LIPSTICK_COLORS.keys()),
                value="Red",
                label="Lipstick Color"
            )
            
            # Blush controls
            blush_check = gr.Checkbox(label="Apply Blush", value=True)
            blush_color = gr.Dropdown(
                choices=list(BLUSH_COLORS.keys()),
                value="Pink",
                label="Blush Color"
            )
            blush_intensity = gr.Slider(
                minimum=0, maximum=100, value=50, step=5,
                label="Blush Intensity (%)"
            )
            
            # Foundation controls
            foundation_check = gr.Checkbox(label="Apply Foundation", value=True)
            gr.Markdown("### 🎨 Shade Selection")
            foundation_shade = gr.Dropdown(
                choices=list(FOUNDATION_SHADES.keys()),
                value="Medium",
                label="Foundation Shade",
                interactive=True,
            )

            shade_desc = gr.Markdown(
                f"*{FOUNDATION_SHADES['Medium']['description']}*",
                visible=True
            )

            gr.Markdown("### 🌡️ Warmth Adjustment")
            warmth_slider = gr.Slider(
                minimum=-20,
                maximum=20,
                value=0,
                step=2,
                label="Warmth (Cool ← → Warm)",
                interactive=True,
            )
            
            apply_btn = gr.Button("✨ Apply Makeup", variant="primary")
        
        with gr.Column():
            image_output = gr.Image(label="Result", type="pil")
            status_text = gr.Textbox(label="Status", interactive=False)
    
    def update_shade_desc(shade):
        desc = FOUNDATION_SHADES.get(shade, FOUNDATION_SHADES["Medium"])["description"]
        return f"*{desc}*"

    shade_dropdown = foundation_shade
    shade_dropdown.change(
        fn=update_shade_desc,
        inputs=[shade_dropdown],
        outputs=[shade_desc]
    )

    # Link button to processing function
    apply_btn.click(
        fn=process_image,
        inputs=[
            image_input,
            lipstick_check, lipstick_color,
            blush_check, blush_color, blush_intensity,
            foundation_check, foundation_shade, warmth_slider
        ],
        outputs=[image_output, status_text],
        api_name="apply_makeup"
    )


if __name__ == "__main__":
    # Hugging Face Spaces launch configuration
    port = int(os.environ.get("PORT", 7860))
    
    demo.launch(server_name="0.0.0.0", server_port=port, share=True)