"""
Face Anti-Spoofing Dataset Generator

This application generates synthetic spoof images for face anti-spoofing dataset creation.
Attack types align with iBeta Level 1 and Level 2 standards:

iBeta Level 1 (Basic Attacks):
- Print Attack: Printed photos of the target face
- Display Attack: Screen replay of face images/photos
- Cut Photo Attack: Partially occluded printed photos

iBeta Level 2 (Advanced Attacks):
- Mask Attack: Paper/plastic masks
- Warped Photo Attack: Deformed/repositioned photos
- Eye Frame Attack: Photos with eye cutouts
"""

import gradio as gr
import numpy as np
from PIL import Image, ImageDraw, ImageFilter, ImageEnhance
import io
import base64
import json
from dataclasses import dataclass
from typing import List, Tuple, Optional
import random

# Attack types aligned with iBeta standards
@dataclass
class AttackType:
    """Represents a spoof attack type"""
    name: str
    level: int  # 1 or 2 (iBeta level)
    description: str
    severity: str  # low, medium, high


# Define attack types
ATTACK_TYPES = [
    AttackType("Print Attack", 1, "Printed photograph of the face", "low"),
    AttackType("Display Attack", 1, "Face shown on screen/display", "low"),
    AttackType("Cut Photo Attack", 1, "Partially cut photograph with eye holes", "medium"),
    AttackType("Paper Mask Attack", 2, "Paper-based face mask", "medium"),
    AttackType("Warped Photo Attack", 2, "Warped/deformed photograph", "high"),
    AttackType("Eye Frame Attack", 2, "Photo with eye cutouts and frame", "high"),
]


def generate_print_attack(image: Image.Image, quality: str = "high") -> Image.Image:
    """
    Simulate a printed photograph attack.
    Adds print artifacts like grain, slight blur, color shift.
    """
    img = image.copy()
    
    if quality == "high":
        # Slight blur simulating high-quality print
        img = img.filter(ImageFilter.GaussianBlur(radius=0.5))
        # Add slight noise
        np_img = np.array(img)
        noise = np.random.normal(0, 3, np_img.shape).astype(np.int16)
        np_img = np.clip(np_img.astype(np.int16) + noise, 0, 255).astype(np.uint8)
        img = Image.fromarray(np_img)
    else:
        # Lower quality print
        img = img.filter(ImageFilter.GaussianBlur(radius=1.5))
        np_img = np.array(img)
        noise = np.random.normal(0, 10, np_img.shape).astype(np.int16)
        np_img = np.clip(np_img.astype(np.int16) + noise, 0, 255).astype(np.uint8)
        img = Image.fromarray(np_img)
    
    # Slight color shift (printing ink effect)
    enhancer = ImageEnhance.Color(img)
    img = enhancer.enhance(0.9)
    
    return img


def generate_display_attack(image: Image.Image, screen_type: str = "phone") -> Image.Image:
    """
    Simulate a display replay attack.
    Adds screen artifacts like moiré patterns, reflections.
    """
    img = image.copy()
    
    # Resize to simulate different screen sizes
    if screen_type == "phone":
        img = img.resize((224, 224), Image.LANCZOS)
        img = img.resize((300, 300), Image.LANCZOS)
    else:
        img = img.resize((256, 256), Image.LANCZOS)
        img = img.resize((400, 300), Image.LANCZOS)
    
    # Add screen moiré effect
    np_img = np.array(img)
    moiré = np.zeros_like(np_img)
    for i in range(moiré.shape[0]):
        for j in range(moiré.shape[1]):
            moiré[i, j] = int(15 * np.sin(i * 0.1) * np.sin(j * 0.1))
    
    np_img = np.clip(np_img.astype(np.int16) + moiré, 0, 255).astype(np.uint8)
    img = Image.fromarray(np_img)
    
    # Add slight glow effect
    img = img.filter(ImageFilter.GaussianBlur(radius=0.5))
    
    return img


def generate_cut_photo_attack(image: Image.Image, cut_type: str = "eyes") -> Image.Image:
    """
    Simulate a cut photo attack with eye holes cut out.
    Used to simulate attempts to bypass eye-based liveness detection.
    """
    img = image.copy()
    width, height = img.size
    
    # Create a black background
    background = Image.new('RGB', (width, height), (0, 0, 0))
    
    # Calculate eye positions (approximate)
    left_eye_x = int(width * 0.35)
    left_eye_y = int(height * 0.35)
    right_eye_x = int(width * 0.65)
    right_eye_y = int(height * 0.35)
    eye_radius = int(min(width, height) * 0.08)
    
    draw = ImageDraw.Draw(background)
    
    if cut_type == "eyes":
        # Cut out eye regions
        mask = Image.new('L', (width, height), 0)
        mask_draw = ImageDraw.Draw(mask)
        mask_draw.ellipse(
            (left_eye_x - eye_radius, left_eye_y - eye_radius,
             left_eye_x + eye_radius, left_eye_y + eye_radius),
            fill=255
        )
        mask_draw.ellipse(
            (right_eye_x - eye_radius, right_eye_y - eye_radius,
             right_eye_x + eye_radius, right_eye_y + eye_radius),
            fill=255
        )
    else:
        # Cut out larger region around eyes
        mask = Image.new('L', (width, height), 0)
        mask_draw = ImageDraw.Draw(mask)
        mask_draw.ellipse(
            (left_eye_x - eye_radius*2, left_eye_y - eye_radius*1.5,
             left_eye_x + eye_radius*2, left_eye_y + eye_radius*1.5),
            fill=255
        )
        mask_draw.ellся:
        mask_draw.ellipse(
            (right_eye_x - eye_radius*2, right_eye_y - eye_radius*1.5,
             right_eye_x + eye_radius*2, right_eye_y + eye_radius*1.5),
            fill=255
        )
    
    # Apply the cut
    img.paste(background, mask=mask)
    
    # Add slight paper texture
    np_img = np.array(img)
    noise = np.random.normal(0, 5, np_img.shape).astype(np.int16)
    np_img = np.clip(np_img.astype(np.int16) + noise, 0, 255).astype(np.uint8)
    img = Image.fromarray(np_img)
    
    return img


def generate_paper_mask_attack(image: Image.Image, mask_style: str = "flat") -> Image.Image:
    """
    Simulate a paper-based mask attack.
    Creates a simplified face shape on paper.
    """
    img = image.copy()
    width, height = img.size
    
    # Create a face-shaped mask
    mask = Image.new('L', (width, height), 0)
    draw = ImageDraw.Draw(mask)
    
    # Draw oval face shape
    face_center_x = width // 2
    face_center_y = height // 2
    face_width = int(width * 0.7)
    face_height = int(height * 0.8)
    
    draw.ellipse(
        (face_center_x - face_width//2, face_center_y - face_height//2,
         face_center_x + face_width//2, face_center_y + face_height//2),
        fill=255
    )
    
    # Create RGB mask for pasting
    mask_rgb = Image.merge('RGB', [mask, mask, mask])
    
    # Apply the mask
    img = Image.composite(img, Image.new('RGB', img.size, (128, 128, 128)), mask)
    
    if mask_style == "curled":
        # Add curling effect at edges
        img = img.filter(ImageFilter.GaussianBlur(radius=2))
    else:
        # Flat paper effect
        img = img.filter(ImageFilter.GaussianBlur(radius=0.5))
    
    # Add paper texture
    np_img = np.array(img)
    paper_noise = np.random.normal(0, 8, np_img.shape).astype(np.int16)
    np_img = np.clip(np_img.astype(np.int16) + paper_noise, 0, 255).astype(np.uint8)
    img = Image.fromarray(np_img)
    
    return img


def generate_warped_photo_attack(image: Image.Image, warp_type: str = "moderate") -> Image.Image:
    """
    Simulate a warped/deformed photo attack.
    Creates non-rigid deformations in the face image.
    """
    img = image.copy()
    width, height = img.size
    
    # Apply different warping based on type
    if warp_type == "slight":
        # Very subtle warping
        coeffs = [(1.02, 0.01, -0.01), (0.01, 1.01, -0.02), (0, 0, 1)]
    elif warp_type == "moderate":
        # Moderate warping
        coeffs = [(1.05, 0.02, -0.03), (0.02, 1.03, -0.02), (0, 0, 1)]
    else:  # severe
        # More pronounced warping
        coeffs = [(1.08, 0.03, -0.05), (0.03, 1.06, -0.03), (0, 0, 1)]
    
    # Apply affine transformation
    img = img.transform(
        (width, height),
        Image.AFFINE,
        coeffs[:2],
        Image.BICUBIC
    )
    
    # Add slight blur
    img = img.filter(ImageFilter.GaussianBlur(radius=0.5))
    
    return img


def generate_eye_frame_attack(image: Image.Image, frame_type: str = "plastic") -> Image.Image:
    """
    Simulate an eye frame attack with cutouts for eyes.
    Includes a physical frame around the photo.
    """
    img = image.copy()
    width, height = img.size
    
    # Create image with frame border
    border_size = int(min(width, height) * 0.1)
    new_width = width + 2 * border_size
    new_height = height + 2 * border_size
    
    # Create new canvas with frame
    if frame_type == "plastic":
        frame_color = (30, 30, 30)  # Dark plastic frame
    else:
        frame_color = (200, 180, 140)  # Wood frame
    
    canvas = Image.new('RGB', (new_width, new_height), frame_color)
    
    # Paste original image in center
    canvas.paste(img, (border_size, border_size))
    
    # Add frame border details
    draw = ImageDraw.Draw(canvas)
    
    # Draw inner border
    inner_border = int(border_size * 0.2)
    draw.rectangle(
        [border_size - inner_border, border_size - inner_border,
         new_width - border_size + inner_border, new_height - border_size + inner_border],
        outline=(200, 200, 200) if frame_type == "plastic" else (150, 130, 90),
        width=3
    )
    
    # Calculate eye positions in the pasted image
    left_eye_x = border_size + int(width * 0.35)
    left_eye_y = border_size + int(height * 0.35)
    right_eye_x = border_size + int(width * 0.65)
    right_eye_y = border_size + int(height * 0.35)
    eye_radius = int(min(width, height) * 0.06)
    
    # Cut out eye holes
    mask = Image.new('L', canvas.size, 0)
    mask_draw = ImageDraw.Draw(mask)
    mask_draw.ellipse(
        (left_eye_x - eye_radius, left_eye_y - eye_radius,
         left_eye_x + eye_radius, left_eye_y + eye_radius),
        fill=255
    )
    mask_draw.ellipse(
        (right_eye_x - eye_radius, right_eye_y - eye_radius,
         right_eye_x + eye_radius, right_eye_y + eye_radius),
        fill=255
    )
    
    # Apply the cutouts
    np_canvas = np.array(canvas)
    np_mask = np.array(mask)
    
    # Darken the cutout regions (simulating background behind frame)
    np_canvas = np.where(np_mask[:, :, np.newaxis] == 255, np_canvas * 0.3, np_canvas)
    
    result = Image.fromarray(np_canvas.astype(np.uint8))
    
    # Add frame texture
    np_result = np.array(result)
    frame_texture = np.random.normal(0, 3, np_result.shape).astype(np.int16)
    np_result = np.clip(np_result.astype(np.int16) + frame_texture, 0, 255).astype(np.uint8)
    result = Image.fromarray(np_result)
    
    return result


def generate_spoof_image(
    reference_image: Image.Image,
    attack_type: str,
    quality_variant: str = "standard"
) -> Tuple[Image.Image, dict]:
    """
    Generate a spoof image based on the selected attack type.
    
    Args:
        reference_image: The input face image to generate spoof from
        attack_type: Type of spoof attack
        quality_variant: Quality variation of the attack
    
    Returns:
        Tuple of (spoof_image, metadata_dict)
    """
    img = reference_image.copy()
    
    # Convert to RGB if needed
    if img.mode != 'RGB':
        img = img.convert('RGB')
    
    metadata = {
        "attack_type": attack_type,
        "quality_variant": quality_variant,
        "ibeta_level": None,
        "spoof_indicators": []
    }
    
    if attack_type == "Print Attack":
        quality = "high" if quality_variant == "high_quality" else "low"
        result = generate_print_attack(img, quality)
        metadata["ibeta_level"] = 1
        metadata["spoof_indicators"] = [
            "print_texture_artifact",
            "moiré_pattern_possible",
            "flat_surface_indicator"
        ]
    
    elif attack_type == "Display Attack":
        screen = "phone" if quality_variant == "mobile" else "monitor"
        result = generate_display_attack(img, screen)
        metadata["ibeta_level"] = 1
        metadata["spoof_indicators"] = [
            "screen_reflection",
            "moiré_pattern",
            "backlight_artifact"
        ]
    
    elif attack_type == "Cut Photo Attack":
        cut_type = "eyes" if quality_variant == "standard" else "large"
        result = generate_cut_photo_attack(img, cut_type)
        metadata["ibeta_level"] = 1
        metadata["spoof_indicators"] = [
            "photo_cut_marks",
            "inconsistent_occlusion",
            "background_discontinuity"
        ]
    
    elif attack_type == "Paper Mask Attack":
        mask_style = "curled" if quality_variant == "worn" else "flat"
        result = generate_paper_mask_attack(img, mask_style)
        metadata["ibeta_level"] = 2
        metadata["spoof_indicators"] = [
            "mask_edge_artifact",
            "flat_surface_texture",
            "inconsistent_skin_texture"
        ]
    
    elif attack_type == "Warped Photo Attack":
        warp_type = "slight" if quality_variant == "minimal" else "moderate"
        result = generate_warped_photo_attack(img, warp_type)
        metadata["ibeta_level"] = 2
        metadata["spoof_indicators"] = [
            "geometric_distortion",
            "inconsistent_perspective",
            "non_rigid_deformation"
        ]
    
    elif attack_type == "Eye Frame Attack":
        frame_type = "plastic" if quality_variant == "standard" else "wooden"
        result = generate_eye_frame_attack(img, frame_type)
        metadata["ibeta_level"] = 2
        metadata["spoof_indicators"] = [
            "frame_artifact",
            "eye_cutout_marks",
            "inconsistent_depth"
        ]
    
    else:
        result = img.copy()
        metadata["error"] = "Unknown attack type"
    
    return result, metadata


def create_dataset_preview(
    reference_image: Image.Image,
    selected_attacks: List[str],
    generate_all: bool = False
) -> Tuple[Image.Image, str, dict]:
    """
    Create a preview of generated spoof images.
    
    Returns:
        Preview image, attack info summary, and dataset metadata
    """
    if reference_image is None:
        return None, "Please upload a reference image first.", {}
    
    if not generate_all and not selected_attacks:
        return None, "Please select at least one attack type.", {}
    
    attacks_to_generate = ATTACK_TYPES if generate_all else [
        at for at in ATTACK_TYPES if at.name in selected_attacks
    ]
    
    # Create a grid preview
    cols = min(len(attacks_to_generate), 3)
    rows = (len(attacks_to_generate) + cols - 1) // cols
    
    img = reference_image.copy()
    if img.mode != 'RGB':
        img = img.convert('RGB')
    
    # Resize for consistent preview
    preview_size = (200, 200)
    img = img.resize(preview_size, Image.LANCZOS)
    
    # Calculate grid dimensions
    cell_width = preview_size[0] + 20
    cell_height = preview_size[1] + 40
    
    grid_width = cols * cell_width
    grid_height = rows * cell_height + 60
    
    # Create preview canvas
    preview = Image.new('RGB', (grid_width, grid_height), (245, 245, 245))
    draw = ImageDraw.Draw(preview)
    
    # Title
    from PIL import ImageFont
    try:
        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 14)
    except:
        font = ImageFont.load_default()
    
    draw.text((10, 10), "Generated Spoof Samples Preview", fill=(50, 50, 50), font=font)
    
    # Generate and place each attack
    dataset_metadata = {
        "total_samples": len(attacks_to_generate),
        "ibeta_level_1_count": sum(1 for a in attacks_to_generate if a.level == 1),
        "ibeta_level_2_count": sum(1 for a in attacks_to_generate if a.level == 2),
        "attacks": []
    }
    
    for idx, attack in enumerate(attacks_to_generate):
        col = idx % cols
        row = idx // cols
        
        x = col * cell_width + 10
        y = row * cell_height + 40
        
        # Generate spoof image
        spoof_img, metadata = generate_spoof_image(
            reference_image,
            attack.name,
            "standard"
        )
        
        spoof_img = spoof_img.resize(preview_size, Image.LANCZOS)
        
        # Paste into preview
        preview.paste(spoof_img, (x, y))
        
        # Draw attack name
        name_y = y + preview_size[1] + 5
        level_text = f"[L{attack.level}]"
        draw.text((x, name_y), level_text, fill=(100, 100, 100), font=font)
        draw.text((x + 30, name_y), attack.name[:15], fill=(50, 50, 50), font=font)
        
        dataset_metadata["attacks"].append({
            "attack_type": attack.name,
            "ibeta_level": attack.level,
            "severity": attack.severity,
            "description": attack.description,
            "metadata": metadata
        })
    
    # Create summary
    summary = (
        f"Dataset Preview Generated:\n"
        f"• Total Samples: {dataset_metadata['total_samples']}\n"
        f"• iBeta Level 1: {dataset_metadata['ibeta_level_1_count']} attacks\n"
        f"• iBeta Level 2: {dataset_metadata['ibeta_level_2_count']} attacks\n"
        f"• Attacks: {', '.join(a.name for a in attacks_to_generate)}"
    )
    
    return preview, summary, dataset_metadata


def export_dataset_metadata(metadata: dict) -> str:
    """Export dataset metadata as JSON string."""
    return json.dumps(metadata, indent=2)


# Custom theme for the app
def create_custom_theme():
    """Create a custom theme for the anti-spoofing dataset generator."""
    return gr.themes.Soft(
        primary_hue="red",
        secondary_hue="orange",
        neutral_hue="slate",
        font=gr.themes.GoogleFont("Inter"),
        text_size="lg",
        spacing_size="lg",
        radius_size="md"
    ).set(
        button_primary_background_fill="*primary_600",
        button_primary_background_fill_hover="*primary_700",
        block_title_text_weight="600",
        body_text_weight="500",
    )


# Gradio 6 App
with gr.Blocks() as demo:
    # Custom CSS for styling
    custom_css = """
    .spoof-header {
        background: linear-gradient(135deg, #1a1a2e 0%, #16213e 100%);
        padding: 20px;
        border-radius: 12px;
        margin-bottom: 20px;
    }
    .attack-card {
        background: white;
        border-radius: 10px;
        padding: 15px;
        box-shadow: 0 2px 8px rgba(0,0,0,0.1);
        margin: 10px 0;
    }
    .level-badge {
        background: linear-gradient(135deg, #f97316 0%, #ef4444 100%);
        color: white;
        padding: 4px 12px;
        border-radius: 20px;
        font-size: 12px;
        font-weight: bold;
    }
    .level-badge-1 {
        background: linear-gradient(135deg, #22c55e 0%, #16a34a 100%);
    }
    .level-badge-2 {
        background: linear-gradient(135deg, #f97316 0%, #ef4444 100%);
    }
    .info-box {
        background: #fef3c7;
        border-left: 4px solid #f59e0b;
        padding: 12px;
        border-radius: 4px;
        margin: 10px 0;
    }
    .metadata-box {
        background: #f1f5f9;
        border: 1px solid #e2e8f0;
        padding: 15px;
        border-radius: 8px;
        font-family: monospace;
        font-size: 12px;
        max-height: 300px;
        overflow-y: auto;
    }
    """
    
    # Header with branding
    with gr.Group(elem_classes=["spoof-header"]):
        gr.Markdown(
            """
            # 🔒 Face Anti-Spoofing Dataset Generator
            
            Generate synthetic spoof images for face anti-spoofing dataset creation.
            Aligned with **iBeta Level 1 & 2** standards.
            
            *Built with [anycoder](https://huggingface.co/spaces/akhaliq/anycoder)*
            """
        )
    
    # Info boxes about iBeta levels
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown(
                """
                ### 📋 iBeta Level 1 (Basic)
                - **Print Attack**: Printed photos
                - **Display Attack**: Screen replays
                - **Cut Photo Attack**: Partially occluded
                """,
                elem_classes=["attack-card"]
            )
        with gr.Column(scale=1):
            gr.Markdown(
                """
                ### ⚠️ iBeta Level 2 (Advanced)
                - **Paper Mask Attack**: Paper masks
                - **Warped Photo Attack**: Deformed photos
                - **Eye Frame Attack**: Eye cutout frames
                """,
                elem_classes=["attack-card"]
            )
    
    # Main input section
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 📤 Reference Image")
            reference_image = gr.Image(
                label="Upload Live Face Image",
                type="pil",
                sources=["upload"],
                height=300
            )
            
            gr.Markdown("### 🎯 Attack Selection")
            attack_checkbox = gr.CheckboxGroup(
                choices=[at.name for at in ATTACK_TYPES],
                value=[ATTACK_TYPES[0].name],  # Default to first attack
                label="Select Attack Types",
                info="Choose which spoof attacks to generate"
            )
            
            generate_all_checkbox = gr.Checkbox(
                value=False,
                label="Generate All Attack Types",
                info="Generate samples for all available attacks"
            )
            
            generate_btn = gr.Button(
                "Generate Spoof Samples",
                variant="primary",
                size="lg"
            )
        
        with gr.Column(scale=1):
            gr.Markdown("### 📊 Preview & Output")
            preview_output = gr.Image(
                label="Generated Spoof Samples",
                type="pil",
                height=400
            )
            
            summary_output = gr.Textbox(
                label="Generation Summary",
                lines=6,
                interactive=False
            )
    
    # Metadata section
    with gr.Accordion("📄 Dataset Metadata (JSON)", open=True):
        metadata_output = gr.Code(
            label="Metadata",
            language="json",
            elem_classes=["metadata-box"]
        )
    
    # Attack details section
    with gr.Accordion("ℹ️ Attack Type Details", open=False):
        gr.Markdown(
            """
            | Attack Type | iBeta Level | Severity | Description |
            |-------------|-------------|----------|-------------|
            | Print Attack | 1 | Low | Printed photograph with typical print artifacts |
            | Display Attack | 1 | Low | Face displayed on screen with moiré patterns |
            | Cut Photo Attack | 1 | Medium | Printed photo with eye cutouts |
            | Paper Mask Attack | 2 | Medium | Flat paper-based face mask |
            | Warped Photo Attack | 2 | High | Deformed photograph with geometric distortion |
            | Eye Frame Attack | 2 | High | Photo with eye cutouts and physical frame |
            
            ### Spoof Indicators
            Each generated sample includes metadata with expected spoof indicators for training:
            - Texture artifacts (print, paper)
            - Moiré patterns (display)
            - Geometric distortions (warped)
            - Occlusion patterns (cut, frame)
            - Surface inconsistencies (mask)
            """
        )
    
    # Event handlers
    def handle_generate(ref_img, attacks, generate_all):
        preview, summary, metadata = create_dataset_preview(
            ref_img,
            attacks,
            generate_all
        )
        metadata_json = export_dataset_metadata(metadata)
        return preview, summary, metadata_json
    
    generate_btn.click(
        fn=handle_generate,
        inputs=[reference_image, attack_checkbox, generate_all_checkbox],
        outputs=[preview_output, summary_output, metadata_output]
    )
    
    # Update when "Generate All" changes
    generate_all_checkbox.change(
        fn=lambda x: gr.CheckboxGroup(interactive=not x),
        inputs=generate_all_checkbox,
        outputs=attack_checkbox
    )
    
    # Live preview on image change
    reference_image.change(
        fn=handle_generate,
        inputs=[reference_image, attack_checkbox, generate_all_checkbox],
        outputs=[preview_output, summary_output, metadata_output]
    )

# Launch with Gradio 6 theme and configuration
demo.launch(
    theme=create_custom_theme(),
    css=custom_css if 'custom_css' in dir() else None,
    footer_links=[
        {"label": "Built with anycoder", "url": "https://huggingface.co/spaces/akhaliq/anycoder"},
        {"label": "Documentation", "url": "https://gradio.app"},
    ],
    title="Face Anti-Spoofing Dataset Generator",
    description="Generate synthetic spoof images aligned with iBeta Level 1 & 2 standards",
)