inpainting_templates.py

import logging
from dataclasses import dataclass, field
from typing import Dict, List, Optional

logger = logging.getLogger(__name__)


@dataclass
class InpaintingTemplate:
    """Data class representing an inpainting template."""

    key: str
    name: str
    category: str
    icon: str
    description: str

    # Prompt templates
    prompt_template: str
    negative_prompt: str

    # Recommended parameters
    controlnet_conditioning_scale: float = 0.7
    feather_radius: int = 8
    guidance_scale: float = 7.5
    num_inference_steps: int = 25

    # Conditioning type preference
    preferred_conditioning: str = "canny"  # "canny" or "depth"

    # Tips for users
    usage_tips: List[str] = field(default_factory=list)


class InpaintingTemplateManager:
    """
    Manages inpainting templates for various use cases.

    Provides categorized presets optimized for different inpainting scenarios
    including object replacement, removal, style transfer, and enhancement.

    Attributes:
        TEMPLATES: Dictionary of all available templates
        CATEGORIES: List of category names in display order

    Example:
        >>> manager = InpaintingTemplateManager()
        >>> template = manager.get_template("object_replacement")
        >>> print(template.prompt_template)
    """

    TEMPLATES: Dict[str, InpaintingTemplate] = {
        # Object Replacement Category 
        "object_replacement": InpaintingTemplate(
            key="object_replacement",
            name="Object Replacement",
            category="Replacement",
            icon="🔄",
            description="Replace selected objects with new content while preserving context",
            prompt_template="{content}, seamlessly integrated into scene, matching lighting and perspective, realistic placement",
            negative_prompt=(
                "inconsistent lighting, wrong perspective, mismatched colors, "
                "visible seams, floating objects, unrealistic placement, original object, "
                "poorly integrated, disconnected from scene, keeping original"
            ),
            controlnet_conditioning_scale=0.48,
            feather_radius=14,
            guidance_scale=10.0,
            num_inference_steps=38,
            preferred_conditioning="canny",
            usage_tips=[
                "Draw mask PRECISELY around the object to replace (include small margin)",
                "Be specific: Instead of 'plant', use 'large green potted fern with detailed leaves'",
                "📝 Example prompts:",
                "  • 'elegant white ceramic teacup with delicate gold rim and floral pattern'",
                "  • 'modern silver laptop computer with sleek metallic finish'",
                "  • 'vintage wooden desk lamp with warm brass details'"
            ]
        ),

        "face_swap": InpaintingTemplate(
            key="face_swap",
            name="Face Enhancement",
            category="Replacement",
            icon="👤",
            description="Enhance or modify facial features while maintaining identity cues",
            prompt_template="{content}, natural skin texture, proper facial proportions, realistic lighting, detailed facial features",
            negative_prompt=(
                "deformed face, asymmetric features, unnatural skin, "
                "plastic appearance, wrong eye direction, blurry features, "
                "artificial smoothing, uncanny valley, distorted proportions"
            ),
            controlnet_conditioning_scale=0.88,
            feather_radius=6,
            guidance_scale=8.5,
            num_inference_steps=35,
            preferred_conditioning="canny",
            usage_tips=[
                "Draw mask CAREFULLY around face outline (avoid hair and neck)",
                "High conditioning preserves facial structure - good for subtle enhancements",
                "📝 Example prompts:",
                "  • 'warm friendly smile with bright eyes and natural expression'",
                "  • 'professional headshot with confident neutral expression'",
                "  • 'gentle smile with soft natural lighting on face'"
            ]
        ),

        "clothing_change": InpaintingTemplate(
            key="clothing_change",
            name="Clothing Change",
            category="Replacement",
            icon="👕",
            description="Change clothing color, pattern, style, and material (moderate changes)",
            prompt_template="transform to {content}, vivid fabric texture, natural folds and wrinkles, correct fit, proper color saturation",
            negative_prompt=(
                "wrong body proportions, floating fabric, unrealistic wrinkles, "
                "mismatched lighting, visible edges, original clothing style, "
                "keeping same color, faded colors, unchanged appearance, partial change"
            ),
            controlnet_conditioning_scale=0.42,
            feather_radius=14,
            guidance_scale=10.5,
            num_inference_steps=38,
            preferred_conditioning="depth",
            usage_tips=[
                "Best for: Similar color changes (black→navy, white→cream) and pattern changes",
                "For extreme changes (black↔white), use 'Dramatic Color Change' template instead",
                "📝 Example prompts:",
                "  • 'deep navy blue formal blazer with fine texture and professional fit'",
                "  • 'bright red polo shirt with clean collar and soft fabric'",
                "  • 'charcoal gray sweater with ribbed knit texture'"
            ]
        ),

        "dramatic_color_change": InpaintingTemplate(
            key="dramatic_color_change",
            name="Dramatic Color Change",
            category="Replacement",
            icon="🎨",
            description="Extreme color transformations (dark↔light, black↔white)",
            prompt_template="completely transform to {content}, vivid color saturation, high contrast, sharp color definition, clear distinct color",
            negative_prompt=(
                "original color, keeping same shade, partial change, color bleeding, "
                "faded colors, mixed tones, subtle change, gradual transition, "
                "original appearance, unchanged, dark remnants, light patches"
            ),
            controlnet_conditioning_scale=0.38,
            feather_radius=16,
            guidance_scale=12.0,
            num_inference_steps=42,
            preferred_conditioning="depth",
            usage_tips=[
                "Optimized for: black→white, white→black, dark→light extreme transformations",
                "Use VERY vivid descriptors: 'PURE WHITE', 'JET BLACK', 'BRIGHT crimson red'",
                "📝 Example prompts:",
                "  • 'pure white dress shirt with crisp texture and sharp collar'",
                "  • 'jet black leather jacket with smooth matte finish'",
                "  • 'bright crimson red blazer with vivid saturated color'"
            ]
        ),

        "clothing_addition": InpaintingTemplate(
            key="clothing_addition",
            name="Add Accessories",
            category="Replacement",
            icon="👔",
            description="Add ties, pockets, buttons, or accessories to clothing (Advanced)",
            prompt_template="{content}, clearly visible, highly detailed accessory, seamlessly integrated into clothing, proper placement and perspective",
            negative_prompt=(
                "missing details, incomplete, floating objects, disconnected, "
                "unrealistic placement, wrong perspective, blurry, poorly integrated, "
                "invisible, faint, unclear, hidden, absent, not visible"
            ),
            controlnet_conditioning_scale=0.25,
            feather_radius=12,
            guidance_scale=14.0,
            num_inference_steps=50,
            preferred_conditioning="depth",
            usage_tips=[
                "⚡ Advanced feature: For adding neckties, pocket squares, badges, or buttons",
                "Draw mask from NECK to CHEST (vertical strip) for ties, not just collar area",
                "📝 Example prompts:",
                "  • 'burgundy silk necktie with diagonal stripes and Windsor knot, hanging down from collar to chest'",
                "  • 'white pocket square with neat fold, visible in breast pocket'",
                "  • 'silver lapel pin with detailed engraving on left collar'",
                "⚠️ TIP: For ties, mask should cover the entire length where tie should appear"
            ]
        ),

        # Object Removal Category 
        "object_removal": InpaintingTemplate(
            key="object_removal",
            name="Object Removal",
            category="Removal",
            icon="🗑️",
            description="Remove unwanted objects and fill with matching background",
            prompt_template="clean background, seamless continuation, {content}",
            negative_prompt=(
                "visible patches, color mismatch, texture inconsistency, "
                "ghost artifacts, blur spots, repeated patterns, visible seams"
            ),
            controlnet_conditioning_scale=0.48,
            feather_radius=14,
            guidance_scale=8.0,
            num_inference_steps=30,
            preferred_conditioning="canny",
            usage_tips=[
                "Draw mask slightly BEYOND object edges for better blending",
                "Describe what background SHOULD look like (e.g., 'grass lawn', 'wooden floor')",
                "📝 Example prompts:",
                "  • 'clean grass lawn with natural green color'",
                "  • 'smooth wooden floor with consistent grain pattern'",
                "  • 'plain white wall with even texture'"
            ]
        ),

        "watermark_removal": InpaintingTemplate(
            key="watermark_removal",
            name="Watermark Removal",
            category="Removal",
            icon="💧",
            description="Remove watermarks and text overlays",
            prompt_template="clean image, no text, seamless background, {content}",
            negative_prompt=(
                "text, watermark, logo, signature, letters, numbers, visible artifacts, "
                "color inconsistency, blur, remnants, ghost text"
            ),
            controlnet_conditioning_scale=0.45,
            feather_radius=12,
            guidance_scale=8.5,
            num_inference_steps=30,
            preferred_conditioning="canny",
            usage_tips=[
                "Draw mask covering ALL watermark/text areas precisely",
                "Describe what SHOULD be there instead (e.g., 'sky', 'fabric texture')",
                "📝 Example prompts:",
                "  • 'clean blue sky with smooth gradient'",
                "  • 'natural skin texture without marks'",
                "  • 'smooth fabric surface with consistent color'"
            ]
        ),

        "blemish_removal": InpaintingTemplate(
            key="blemish_removal",
            name="Blemish Removal",
            category="Removal",
            icon="✨",
            description="Remove skin blemishes, scratches, or small imperfections",
            prompt_template="clean smooth surface, natural texture, {content}",
            negative_prompt=(
                "artificial smoothing, plastic texture, visible editing, "
                "color patches, unnatural appearance, over-processed"
            ),
            controlnet_conditioning_scale=0.6,
            feather_radius=6,
            guidance_scale=6.5,
            num_inference_steps=20,
            preferred_conditioning="canny",
            usage_tips=[
                "Draw small precise masks for EACH blemish/imperfection",
                "Lower guidance (6.5) preserves natural skin texture",
                "📝 Example prompts:",
                "  • 'natural clean skin with smooth texture'",
                "  • 'smooth surface without scratches or marks'",
                "  • 'clear skin with natural pores and texture'"
            ]
        ),

        # Style Transfer Category 
        "style_artistic": InpaintingTemplate(
            key="style_artistic",
            name="Artistic Style",
            category="Style",
            icon="🎨",
            description="Apply artistic style to selected region",
            prompt_template="{content}, distinctive artistic style, strong painterly effect, creative interpretation, visible brushstrokes",
            negative_prompt=(
                "photorealistic, plain, boring, low contrast, unchanged, "
                "inconsistent style, harsh transitions, original appearance, realistic photo"
            ),
            controlnet_conditioning_scale=0.52,
            feather_radius=12,
            guidance_scale=11.5,
            num_inference_steps=38,
            preferred_conditioning="canny",
            usage_tips=[
                "Works best on larger areas (faces, clothing, backgrounds) for visible transformation",
                "Be VERY specific about art style for best results",
                "📝 Example prompts:",
                "  • 'impressionist oil painting with visible thick brushstrokes and vibrant colors'",
                "  • 'watercolor painting with soft edges and delicate color washes'",
                "  • 'Van Gogh style with swirling brushstrokes and bold color contrasts'"
            ]
        ),

        "style_vintage": InpaintingTemplate(
            key="style_vintage",
            name="Vintage Look",
            category="Style",
            icon="📻",
            description="Apply vintage or retro aesthetic to selected area",
            prompt_template="{content}, strong vintage aesthetic, warm sepia tones, film grain texture, nostalgic atmosphere",
            negative_prompt=(
                "modern, digital, cold colors, harsh contrast, "
                "oversaturated, neon colors, contemporary look, clean digital, crisp"
            ),
            controlnet_conditioning_scale=0.55,
            feather_radius=14,
            guidance_scale=10.5,
            num_inference_steps=35,
            preferred_conditioning="canny",
            usage_tips=[
                "Works best on medium to large regions for visible aesthetic change",
                "Specify era and style for best results",
                "📝 Example prompts:",
                "  • '1920s sepia photograph with faded brown tones and soft grain'",
                "  • '1970s vintage photo with warm orange tones and slight film grain'",
                "  • '1950s Kodachrome with saturated warm colors and nostalgic feel'"
            ]
        ),

        "style_anime": InpaintingTemplate(
            key="style_anime",
            name="Anime Style",
            category="Style",
            icon="🎌",
            description="Transform selected region to anime/illustration style",
            prompt_template="{content}, anime illustration style, clean sharp lines, vibrant saturated colors, cel-shaded with flat colors",
            negative_prompt=(
                "photorealistic, blurry lines, muddy colors, realistic photo, "
                "3D render, uncanny valley, western cartoon, gradient shading, photographic"
            ),
            controlnet_conditioning_scale=0.48,
            feather_radius=10,
            guidance_scale=12.5,
            num_inference_steps=40,
            preferred_conditioning="canny",
            usage_tips=[
                "⚡ DRAMATIC transformation - best for portraits and characters",
                "Expect significant stylistic changes from realistic to anime",
                "📝 Example prompts:",
                "  • 'modern anime style with large expressive eyes and vibrant colors'",
                "  • 'Studio Ghibli style with soft features and warm color palette'",
                "  • 'manga style with clean black lines and cel-shaded coloring'"
            ]
        ),

        # Detail Enhancement Category 
        "detail_enhance": InpaintingTemplate(
            key="detail_enhance",
            name="Detail Enhancement",
            category="Enhancement",
            icon="🔍",
            description="Add fine details and textures to selected area",
            prompt_template="{content}, highly detailed, intricate textures, fine details, sharp focus",
            negative_prompt=(
                "blurry, smooth, low detail, soft focus, "
                "oversimplified, lacking texture"
            ),
            controlnet_conditioning_scale=0.85,
            feather_radius=4,
            guidance_scale=8.0,
            num_inference_steps=30,
            preferred_conditioning="depth",
            usage_tips=[
                "High conditioning (0.85) preserves overall structure while adding detail",
                "Best for adding fine details to existing objects",
                "📝 Example prompts:",
                "  • 'highly detailed fabric with visible weave and fine threads'",
                "  • 'intricate wood grain with natural knots and detailed texture'",
                "  • 'sharp facial features with fine skin pores and detail'"
            ]
        ),

        "texture_add": InpaintingTemplate(
            key="texture_add",
            name="Texture Addition",
            category="Enhancement",
            icon="🧱",
            description="Add or enhance surface textures",
            prompt_template="{content} texture, realistic surface detail, natural material appearance",
            negative_prompt=(
                "flat, smooth, unrealistic, plastic, "
                "wrong material, inconsistent texture"
            ),
            controlnet_conditioning_scale=0.8,
            feather_radius=5,
            guidance_scale=7.5,
            num_inference_steps=25,
            preferred_conditioning="depth",
            usage_tips=[
                "Specify material type clearly for best results",
                "Depth conditioning preserves 3D form while changing texture",
                "📝 Example prompts:",
                "  • 'rough wood texture with natural grain and knots'",
                "  • 'soft cotton fabric with gentle weave pattern'",
                "  • 'smooth marble surface with subtle veining'"
            ]
        ),

        "lighting_fix": InpaintingTemplate(
            key="lighting_fix",
            name="Lighting Correction",
            category="Enhancement",
            icon="💡",
            description="Correct or enhance lighting in selected area",
            prompt_template="{content}, proper lighting, natural shadows, balanced exposure",
            negative_prompt=(
                "harsh shadows, overexposed, underexposed, "
                "flat lighting, unnatural highlights"
            ),
            controlnet_conditioning_scale=0.65,
            feather_radius=15,
            guidance_scale=7.0,
            num_inference_steps=25,
            preferred_conditioning="depth",
            usage_tips=[
                "Use large feather (15px) for smooth lighting transitions",
                "Best for fixing uneven lighting or adding natural light",
                "📝 Example prompts:",
                "  • 'soft natural lighting from window, gentle shadows'",
                "  • 'balanced exposure with warm golden hour light'",
                "  • 'even studio lighting with soft diffused shadows'"
            ]
        ),

        # Background Category 
        "background_extend": InpaintingTemplate(
            key="background_extend",
            name="Background Extension",
            category="Background",
            icon="📐",
            description="Extend image background seamlessly",
            prompt_template="seamless background extension, {content}, consistent style and lighting",
            negative_prompt=(
                "visible seams, style mismatch, lighting inconsistency, "
                "repeated elements, unnatural continuation, abrupt changes"
            ),
            controlnet_conditioning_scale=0.55,
            feather_radius=20,
            guidance_scale=8.0,
            num_inference_steps=32,
            preferred_conditioning="canny",
            usage_tips=[
                "Draw mask on area to extend (edges of image)",
                "Large feather (20px) ensures smooth blending with existing background",
                "📝 Example prompts:",
                "  • 'continue the wooden floor with same grain pattern'",
                "  • 'extend blue sky with matching clouds and lighting'",
                "  • 'seamless continuation of brick wall texture'"
            ]
        ),

        "background_replace": InpaintingTemplate(
            key="background_replace",
            name="Background Replacement",
            category="Background",
            icon="🖼️",
            description="Replace background while keeping subject intact",
            prompt_template="{content}, professional background scene, seamless integration with subject, matching lighting and atmosphere",
            negative_prompt=(
                "floating subject, inconsistent lighting, disconnected, "
                "wrong perspective, visible edges, color mismatch, original background, "
                "poor integration, obvious composite"
            ),
            controlnet_conditioning_scale=0.60,
            feather_radius=12,
            guidance_scale=9.5,
            num_inference_steps=35,
            preferred_conditioning="depth",
            usage_tips=[
                "Draw mask around ENTIRE background (leave subject unmasked with small margin)",
                "Include lighting description to match subject for natural results",
                "📝 Example prompts:",
                "  • 'professional photography studio with white backdrop and soft lighting'",
                "  • 'modern minimalist office with white walls and bright natural lighting'",
                "  • 'sunny beach with blue ocean and golden hour lighting'"
            ]
        ),
    }

    # Category display order
    CATEGORIES = ["Replacement", "Removal", "Style", "Enhancement", "Background"]

    def __init__(self):
        """Initialize the InpaintingTemplateManager."""
        logger.info(f"InpaintingTemplateManager initialized with {len(self.TEMPLATES)} templates")

    def get_all_templates(self) -> Dict[str, InpaintingTemplate]:
        """
        Get all available templates.

        Returns
        -------
        dict
            Dictionary of all templates keyed by template key
        """
        return self.TEMPLATES

    def get_template(self, key: str) -> Optional[InpaintingTemplate]:
        """
        Get a specific template by key.

        Parameters
        ----------
        key : str
            Template identifier

        Returns
        -------
        InpaintingTemplate or None
            Template if found, None otherwise
        """
        return self.TEMPLATES.get(key)

    def get_templates_by_category(self, category: str) -> List[InpaintingTemplate]:
        """
        Get all templates in a specific category.

        Parameters
        ----------
        category : str
            Category name

        Returns
        -------
        list
            List of templates in the category
        """
        return [t for t in self.TEMPLATES.values() if t.category == category]

    def get_categories(self) -> List[str]:
        """
        Get list of all categories in display order.

        Returns
        -------
        list
            Category names
        """
        return self.CATEGORIES

    def get_template_choices_sorted(self) -> List[str]:
        """
        Get template choices formatted for Gradio dropdown.

        Returns list of display strings sorted by category then A-Z.
        Format: "icon Name"

        Returns
        -------
        list
            Formatted display strings for dropdown
        """
        display_list = []

        for category in self.CATEGORIES:
            templates = self.get_templates_by_category(category)
            for template in sorted(templates, key=lambda t: t.name):
                display_name = f"{template.icon} {template.name}"
                display_list.append(display_name)

        return display_list

    def get_template_key_from_display(self, display_name: str) -> Optional[str]:
        """
        Get template key from display name.

        Parameters
        ----------
        display_name : str
            Display string like "🔄 Object Replacement"

        Returns
        -------
        str or None
            Template key if found
        """
        if not display_name:
            return None

        for key, template in self.TEMPLATES.items():
            if f"{template.icon} {template.name}" == display_name:
                return key
        return None

    def get_parameters_for_template(self, key: str) -> Dict[str, any]:
        """
        Get recommended parameters for a template.

        Parameters
        ----------
        key : str
            Template key

        Returns
        -------
        dict
            Dictionary of parameter names and values
        """
        template = self.get_template(key)
        if not template:
            return {}

        return {
            "controlnet_conditioning_scale": template.controlnet_conditioning_scale,
            "feather_radius": template.feather_radius,
            "guidance_scale": template.guidance_scale,
            "num_inference_steps": template.num_inference_steps,
            "preferred_conditioning": template.preferred_conditioning
        }

    def build_prompt(self, key: str, content: str) -> str:
        """
        Build complete prompt from template and user content.

        Parameters
        ----------
        key : str
            Template key
        content : str
            User-provided content description

        Returns
        -------
        str
            Formatted prompt with content inserted
        """
        template = self.get_template(key)
        if not template:
            return content

        return template.prompt_template.format(content=content)

    def get_negative_prompt(self, key: str) -> str:
        """
        Get negative prompt for a template.

        Parameters
        ----------
        key : str
            Template key

        Returns
        -------
        str
            Negative prompt string
        """
        template = self.get_template(key)
        if not template:
            return ""
        return template.negative_prompt

    def get_usage_tips(self, key: str) -> List[str]:
        """
        Get usage tips for a template.

        Parameters
        ----------
        key : str
            Template key

        Returns
        -------
        list
            List of tip strings
        """
        template = self.get_template(key)
        if not template:
            return []
        return template.usage_tips

    def build_gallery_html(self) -> str:
        """
        Build HTML for template gallery display.

        Returns
        -------
        str
            HTML string for Gradio display
        """
        html_parts = ['<div class="inpainting-gallery">']

        for category in self.CATEGORIES:
            templates = self.get_templates_by_category(category)
            if not templates:
                continue

            html_parts.append(f'''
            <div class="inpainting-category">
                <h4 class="inpainting-category-title">{category}</h4>
                <div class="inpainting-grid">
            ''')

            for template in sorted(templates, key=lambda t: t.name):
                html_parts.append(f'''
                <div class="inpainting-card" data-template="{template.key}">
                    <span class="inpainting-icon">{template.icon}</span>
                    <span class="inpainting-name">{template.name}</span>
                    <span class="inpainting-desc">{template.description[:50]}...</span>
                </div>
                ''')

            html_parts.append('</div></div>')

        html_parts.append('</div>')
        return ''.join(html_parts)