"""
Inpainter MCP Server (FastMCP)
- Tool:
  - inpaint(image_path, mask_path, prompt) -> {"output_path": str}

Uses Stable Diffusion with ControlNet conditioning for high-quality inpainting.
SDXL: Pose ControlNet for persons, Canny ControlNet for objects.
SD 1.5: Pose ControlNet for persons, Canny ControlNet for objects.
"""
from __future__ import annotations

import warnings
warnings.filterwarnings("ignore", message="Unable to find acceptable character detection dependency")
warnings.filterwarnings("ignore", message="Importing from timm.models.layers is deprecated")
warnings.filterwarnings("ignore", message="Importing from timm.models.registry is deprecated")
warnings.filterwarnings("ignore", message="Overwriting tiny_vit")

from pathlib import Path
import sys
import os
from typing import Dict, Any
from PIL import Image, ImageFilter

try:
    import torch
    import numpy as np
    from diffusers import StableDiffusionInpaintPipeline, StableDiffusionControlNetInpaintPipeline, ControlNetModel
    from diffusers import StableDiffusionXLInpaintPipeline, StableDiffusionXLControlNetPipeline, StableDiffusionXLControlNetInpaintPipeline, AutoencoderKL, UniPCMultistepScheduler
    from diffusers.utils import load_image
    from controlnet_aux import OpenposeDetector, MidasDetector
    from color_transfer import color_transfer
    from PIL import Image
    import cv2
except ImportError:
    torch = None
    np = None
    StableDiffusionControlNetInpaintPipeline = None
    StableDiffusionXLInpaintPipeline = None
    StableDiffusionXLControlNetInpaintPipeline = None
    ControlNetModel = None
    load_image = None
    OpenposeDetector = None
    MidasDetector = None
    color_transfer = None
    Image = None
    cv2 = None


ARTIFACTS = Path("artifacts")
ARTIFACTS.mkdir(parents=True, exist_ok=True)


def _run_without_mcp() -> None:
    print(
        "Inpainter MCP server requires the 'mcp' package.\n"
        "Install with: pip install mcp\n"
        "Then run: python -m mcp.servers.inpainter_server"
    )


try:
    from mcp.server.fastmcp import FastMCP
except Exception:
    FastMCP = None  # type: ignore[misc,assignment]


if FastMCP is None:
    def main() -> None:  # noqa: D401
        _run_without_mcp()
else:
    app = FastMCP("InpainterServer")

    # Cache for pipeline and models to avoid reloading on every call
    _pipeline_cache = {}
    _detector_cache = {}

    @app.tool()
    def inpaint(image_path: str, mask_path: str, prompt: str, description: str = "") -> Dict[str, Any]:  # noqa: D401
        """Inpaint image using mask and prompt with Stable Diffusion.
        
        Args:
            image_path: Path to input image
            mask_path: Path to mask image
            prompt: Base prompt (e.g., 'person')
            description: Anonymized description to guide generation (e.g., 'person in business attire')
        """
        if torch is None or np is None or StableDiffusionControlNetInpaintPipeline is None or ControlNetModel is None or load_image is None or OpenposeDetector is None or MidasDetector is None or color_transfer is None or Image is None:
            return {"output_path": "", "error": "diffusers, controlnet_aux (with MidasDetector), color_transfer, numpy, or opencv not installed. Install with: pip install diffusers torch accelerate controlnet_aux color-transfer numpy opencv-python"}

        ipath = Path(image_path)
        mpath = Path(mask_path)
        if not ipath.exists():
            return {"output_path": "", "error": f"image not found: {image_path}"}
        if not mpath.exists():
            return {"output_path": "", "error": f"mask not found: {mask_path}"}

        # Determine if this is person or object inpainting based on prompt (needed early for pipeline selection)
        is_person = prompt.lower().strip() == "person"

        try:
            # Load images
            original_image = load_image(str(ipath))
            mask_image = load_image(str(mpath))

            # Find bounding box of the mask
            mask_array = np.array(mask_image.convert("L"))
            mask_threshold = int(os.getenv("INPAINT_MASK_THRESHOLD", "128"))
            rows = np.any(mask_array > mask_threshold, axis=1)
            cols = np.any(mask_array > mask_threshold, axis=0)
            if not np.any(rows) or not np.any(cols):
                return {"output_path": "", "error": "mask is empty"}
            ymin, ymax = np.where(rows)[0][[0, -1]]
            xmin, xmax = np.where(cols)[0][[0, -1]]
            # Expand bbox more aggressively for better context
            margin = int(os.getenv("INPAINT_BBOX_MARGIN", "20"))
            xmin = max(0, xmin - margin)
            ymin = max(0, ymin - margin)
            xmax = min(original_image.width, xmax + margin)
            ymax = min(original_image.height, ymax + margin)
            bbox = (xmin, ymin, xmax, ymax)

            # Crop images to bbox
            cropped_image = original_image.crop(bbox)
            cropped_mask = mask_image.crop(bbox)

            # Ensure mask is binary with proper threshold from YOLOv8m segmentation
            cropped_mask = cropped_mask.convert("L").point(lambda p: 255 if p > mask_threshold else 0)  # Use threshold for clean YOLO masks
            
            # Store original mask for compositing (to avoid ghosting)
            original_mask = cropped_mask.copy()
            
            # # Dilate mask slightly for inpainting to ensure complete coverage
            # # This prevents gray/unfilled areas at mask boundaries
            # if cv2 is not None:
            #     mask_np = np.array(cropped_mask)
            #     kernel = np.ones((3, 3), np.uint8)
            #     dilated_mask = cv2.dilate(mask_np, kernel, iterations=1)
            #     cropped_mask = Image.fromarray(dilated_mask)

            # Get cropped size
            crop_width, crop_height = cropped_image.size

            # Calculate upscale ratio to maximize size in target resolution
            # FADM paper tested: 512x512, 768x768, and 1024x1024 (up to 16x more pixels)
            # SDXL works best at 1024x1024, SD 1.5 at configured resolution
            target_size = int(os.getenv("INPAINT_RESOLUTION", "768"))
            if is_person and os.getenv("INPAINT_USE_SDXL", "true").lower() == "true":
                target_size = int(os.getenv("INPAINT_SDXL_RESOLUTION", "1024"))
            elif not is_person and os.getenv("INPAINT_OBJECT_USE_SDXL", "true").lower() == "true":
                target_size = int(os.getenv("INPAINT_SDXL_RESOLUTION", "1024"))
            
            ratio = min(target_size / crop_width, target_size / crop_height)
            new_w = int(crop_width * ratio)
            new_h = int(crop_height * ratio)

            # Load/cache controlnets and pipeline (only once)
            device = "cuda" if torch.cuda.is_available() else "cpu"
            dtype_key = torch.float16 if torch.cuda.is_available() else torch.float32
            
            # Load appropriate pipeline based on subject type
            if is_person:
                # Check which model to use (SDXL by default, SD 1.5 fallback)
                use_sdxl = os.getenv("INPAINT_USE_SDXL", "true").lower() == "true"
                
                if use_sdxl:
                    pipeline_key = "pipe_person_sdxl"
                    if pipeline_key not in _pipeline_cache:
                        print("[InpainterServer] Loading SDXL person inpainting pipeline (pose ControlNet)...", file=sys.stderr, flush=True)
                        # Load SDXL ControlNet for OpenPose
                        # Note: This model uses .bin format (pickle), not .safetensors
                        controlnet_pose = ControlNetModel.from_pretrained(
                            "thibaud/controlnet-openpose-sdxl-1.0", 
                            torch_dtype=dtype_key,
                            use_safetensors=False  # Model uses .bin format
                        )
                        # Load SDXL INPAINTING pipeline with ControlNet
                        import warnings as _warnings
                        with _warnings.catch_warnings():
                            _warnings.filterwarnings("ignore", message=".*torch_dtype.*deprecated.*")
                            _warnings.filterwarnings("ignore", message=".*config attributes.*were passed.*but are not expected.*")
                            pipe = StableDiffusionXLControlNetInpaintPipeline.from_pretrained(
                                "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
                                controlnet=controlnet_pose,
                                torch_dtype=dtype_key
                            )
                        pipe.to(device)
                        
                        # Use UniPCMultistepScheduler for better quality/speed trade-off
                        pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
                        
                        # Enable model CPU offload for memory efficiency
                        if hasattr(pipe, 'enable_model_cpu_offload'):
                            try:
                                pipe.enable_model_cpu_offload()
                            except Exception:
                                pass
                        
                        # Enable additional memory optimizations for SDXL
                        # VAE slicing: Reduces VRAM by ~2-3GB by processing VAE in slices
                        if hasattr(pipe, 'enable_vae_slicing'):
                            try:
                                pipe.enable_vae_slicing()
                                print("[InpainterServer] SDXL VAE slicing enabled (saves ~2-3GB VRAM)", file=sys.stderr, flush=True)
                            except Exception:
                                pass
                        
                        # Attention slicing: Reduces VRAM by ~1-2GB by processing attention in steps
                        if hasattr(pipe, 'enable_attention_slicing'):
                            try:
                                pipe.enable_attention_slicing(slice_size="auto")
                                print("[InpainterServer] SDXL attention slicing enabled (saves ~1-2GB VRAM)", file=sys.stderr, flush=True)
                            except Exception:
                                pass
                        
                        pipe.safety_checker = None  # SDXL doesn't have safety_checker by default
                        _pipeline_cache[pipeline_key] = pipe
                        print("[InpainterServer] SDXL person inpainting pipeline loaded and cached.", file=sys.stderr, flush=True)
                    else:
                        pipe = _pipeline_cache[pipeline_key]
                else:
                    pipeline_key = "pipe_person"
                    if pipeline_key not in _pipeline_cache:
                        print("[InpainterServer] Loading SD 1.5 person pipeline (pose ControlNet only)...", file=sys.stderr, flush=True)
                        # ControlNetModel uses torch_dtype, not dtype
                        controlnet_pose = ControlNetModel.from_pretrained(
                            "lllyasviel/sd-controlnet-openpose", torch_dtype=dtype_key
                        )
                        # Pipeline - suppress deprecation warning by using variant parameter
                        import warnings as _warnings
                        with _warnings.catch_warnings():
                            _warnings.filterwarnings("ignore", message=".*torch_dtype.*deprecated.*")
                            pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
                                "runwayml/stable-diffusion-v1-5", 
                                controlnet=controlnet_pose, 
                                torch_dtype=dtype_key
                            )
                        pipe.to(device)
                        
                        # Use better scheduler for quality (DPMSolverMultistep)
                        from diffusers import DPMSolverMultistepScheduler
                        pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
                        
                        # Enable xformers for memory efficiency and speed
                        if hasattr(pipe, 'enable_xformers_memory_efficient_attention'):
                            try:
                                pipe.enable_xformers_memory_efficient_attention()
                            except Exception:
                                pass
                        
                        pipe.safety_checker = lambda images, **kwargs: (images, [False] * len(images))  # Disable NSFW checker
                        _pipeline_cache[pipeline_key] = pipe
                        print("[InpainterServer] SD 1.5 person pipeline loaded and cached.", file=sys.stderr, flush=True)
                    else:
                        pipe = _pipeline_cache[pipeline_key]
            else:
                # Check which model to use for objects (SDXL by default, SD 1.5 fallback)
                use_sdxl_objects = os.getenv("INPAINT_OBJECT_USE_SDXL", "true").lower() == "true"
                
                if use_sdxl_objects:
                    pipeline_key = "pipe_object_sdxl"
                    if pipeline_key not in _pipeline_cache:
                        print("[InpainterServer] Loading SDXL object inpainting pipeline (Canny ControlNet)...", file=sys.stderr, flush=True)
                        # Load SDXL-compatible Canny ControlNet
                        # Note: Using SD 1.5 Canny ControlNet with SDXL (works with proper scaling)
                        controlnet_canny = ControlNetModel.from_pretrained(
                            "diffusers/controlnet-canny-sdxl-1.0",
                            torch_dtype=dtype_key
                        )
                        # Load SDXL INPAINTING pipeline with Canny ControlNet
                        import warnings as _warnings
                        with _warnings.catch_warnings():
                            _warnings.filterwarnings("ignore", message=".*torch_dtype.*deprecated.*")
                            _warnings.filterwarnings("ignore", message=".*config attributes.*were passed.*but are not expected.*")
                            pipe = StableDiffusionXLControlNetInpaintPipeline.from_pretrained(
                                "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
                                controlnet=controlnet_canny,
                                torch_dtype=dtype_key
                            )
                        pipe.to(device)
                        
                        # Use UniPCMultistepScheduler for better quality/speed trade-off
                        pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
                        
                        # Enable model CPU offload for memory efficiency
                        if hasattr(pipe, 'enable_model_cpu_offload'):
                            try:
                                pipe.enable_model_cpu_offload()
                            except Exception:
                                pass
                        
                        # Enable additional memory optimizations for SDXL
                        # VAE slicing: Reduces VRAM by ~2-3GB by processing VAE in slices
                        if hasattr(pipe, 'enable_vae_slicing'):
                            try:
                                pipe.enable_vae_slicing()
                                print("[InpainterServer] SDXL VAE slicing enabled (saves ~2-3GB VRAM)", file=sys.stderr, flush=True)
                            except Exception:
                                pass
                        
                        # Attention slicing: Reduces VRAM by ~1-2GB by processing attention in steps
                        if hasattr(pipe, 'enable_attention_slicing'):
                            try:
                                pipe.enable_attention_slicing(slice_size="auto")
                                print("[InpainterServer] SDXL attention slicing enabled (saves ~1-2GB VRAM)", file=sys.stderr, flush=True)
                            except Exception:
                                pass
                        
                        pipe.safety_checker = None  # SDXL doesn't have safety_checker by default
                        _pipeline_cache[pipeline_key] = pipe
                        print("[InpainterServer] SDXL object inpainting pipeline loaded and cached.", file=sys.stderr, flush=True)
                    else:
                        pipe = _pipeline_cache[pipeline_key]
                else:
                    pipeline_key = "pipe_object"
                    if pipeline_key not in _pipeline_cache:
                        print("[InpainterServer] Loading SD 1.5 object pipeline (Canny ControlNet with low threshold)...", file=sys.stderr, flush=True)
                        # Load Canny ControlNet for minimal edge-preserving object removal
                        controlnet_canny = ControlNetModel.from_pretrained(
                            "lllyasviel/sd-controlnet-canny", torch_dtype=dtype_key
                        )
                        # Pipeline with Canny ControlNet
                        import warnings as _warnings
                        with _warnings.catch_warnings():
                            _warnings.filterwarnings("ignore", message=".*torch_dtype.*deprecated.*")
                            pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
                                "runwayml/stable-diffusion-v1-5", 
                                controlnet=controlnet_canny, 
                                torch_dtype=dtype_key
                            )
                        pipe.to(device)
                        
                        # Use better scheduler for quality (DPMSolverMultistep)
                        from diffusers import DPMSolverMultistepScheduler
                        pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
                        
                        # Enable xformers for memory efficiency and speed
                        if hasattr(pipe, 'enable_xformers_memory_efficient_attention'):
                            try:
                                pipe.enable_xformers_memory_efficient_attention()
                            except Exception:
                                pass
                        
                        pipe.safety_checker = lambda images, **kwargs: (images, [False] * len(images))  # Disable NSFW checker
                        _pipeline_cache[pipeline_key] = pipe
                        print("[InpainterServer] SD 1.5 object pipeline loaded and cached.", file=sys.stderr, flush=True)
                    else:
                        pipe = _pipeline_cache[pipeline_key]

            # Load/cache detectors based on subject type
            if is_person:
                # For persons: Use pose ControlNet
                if "openpose" not in _detector_cache:
                    _detector_cache["openpose"] = OpenposeDetector.from_pretrained("lllyasviel/Annotators")
                
                openpose = _detector_cache["openpose"]
                
                # Estimate pose from cropped image for controlnet conditioning
                cropped_pose = openpose(cropped_image)

            # Pad cropped inputs to target_size with upscaling
            def pad_to_target(img, fill=0):
                resized = img.resize((new_w, new_h), Image.LANCZOS)
                new_img = Image.new(img.mode, (target_size, target_size), fill)
                offset_x = (target_size - new_w) // 2
                offset_y = (target_size - new_h) // 2
                new_img.paste(resized, (offset_x, offset_y))
                return new_img

            padded_image = pad_to_target(cropped_image, (0, 0, 0))
            padded_mask = pad_to_target(cropped_mask, 0)
            
            # Pad control images based on subject type
            if is_person:
                padded_pose = pad_to_target(cropped_pose, (0, 0, 0))
            else:
                # For objects: Use Canny edge detection with very low thresholds for minimal guidance
                if cv2 is not None:
                    cropped_np = np.array(cropped_image)
                    # Convert to grayscale for edge detection
                    gray = cv2.cvtColor(cropped_np, cv2.COLOR_RGB2GRAY)
                    # Apply Canny edge detection with very low thresholds for subtle guidance
                    canny_low = int(os.getenv("INPAINT_OBJECT_CANNY_LOW", "10"))
                    canny_high = int(os.getenv("INPAINT_OBJECT_CANNY_HIGH", "30"))
                    edges = cv2.Canny(gray, canny_low, canny_high)
                    # Convert back to RGB for ControlNet
                    cropped_canny = Image.fromarray(cv2.cvtColor(edges, cv2.COLOR_GRAY2RGB))
                else:
                    # Fallback if cv2 not available
                    cropped_canny = cropped_image.convert('L').convert('RGB')
                padded_canny = pad_to_target(cropped_canny, (0, 0, 0))
            
            # Build prompt based on subject type
            if is_person:
                # For persons: Use FADM quality prompts with anonymized description
                prompt_suffix = os.getenv("INPAINT_PERSON_PROMPT_SUFFIX", "RAW photo, 8k uhd, dslr, soft lighting, high quality, film grain, Fujifilm XT3, photorealistic, detailed")
                prompt_fallback = os.getenv("INPAINT_PERSON_PROMPT_FALLBACK", "RAW photo, subject, 8k, dslr, soft lighting, high quality, film grain, Fujifilm XT3")
                if description and description.strip():
                    subject_desc = description.strip()
                    positive_prompt = f"{subject_desc}, {prompt_suffix}" if prompt_suffix else subject_desc
                else:
                    positive_prompt = prompt_fallback
                
                # FADM negative prompt - prevents person-specific artifacts
                negative_prompt = os.getenv(
                    "INPAINT_PERSON_NEGATIVE_PROMPT",
                    "deformed iris, deformed pupils, semi-realistic, cgi, 3d, render, sketch, cartoon, drawing, anime, "
                    "text, cropped, out of frame, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, "
                    "mutilated, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, "
                    "blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, "
                    "gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, "
                    "fused fingers, too many fingers, long neck"
                )
                
                # Log prompts for person inpainting
                print(f"[InpainterServer] === Person Inpainting Prompts ===", file=sys.stderr, flush=True)
                print(f"[InpainterServer] Description: {description if description else '(none - using fallback)'}", file=sys.stderr, flush=True)
                print(f"[InpainterServer] Positive: {positive_prompt}", file=sys.stderr, flush=True)
                print(f"[InpainterServer] Negative: {negative_prompt[:100]}..." if len(negative_prompt) > 100 else f"[InpainterServer] Negative: {negative_prompt}", file=sys.stderr, flush=True)
            else:
                # For objects: Use description to guide background generation
                # The description describes what to REMOVE and WHERE it is (context for background)
                # Example: "a blue sign on a building" → generate "building facade" background
                background_fallback = os.getenv("INPAINT_OBJECT_PROMPT_FALLBACK", "natural scene, building facade, street background, realistic outdoor environment")
                
                if description and description.strip():
                    # Use description as positive prompt to guide background generation
                    # The mask tells the model WHERE, the prompt tells it WHAT context
                    positive_prompt = description.strip()
                else:
                    positive_prompt = background_fallback
                
                # Negative prompt: prevent regenerating similar objects
                negative_prompt = os.getenv("INPAINT_OBJECT_NEGATIVE_PROMPT", "")
                
                # Log prompts for object inpainting
                print(f"[InpainterServer] === Object Inpainting Prompts ===", file=sys.stderr, flush=True)
                print(f"[InpainterServer] Description (context): {description if description else '(none - using fallback)'}", file=sys.stderr, flush=True)
                print(f"[InpainterServer] Positive: {positive_prompt}", file=sys.stderr, flush=True)
                print(f"[InpainterServer] Negative: {negative_prompt if negative_prompt else ''}", file=sys.stderr, flush=True)
            
            # βd (beta_d) parameter: noise strength / anonymization strength
            # FADM recommends βd = 0.6 as balance between quality and anonymization
            # Lower values (e.g., 0.3-0.5) = less anonymization, more original preserved
            # Higher values (e.g., 0.7-0.9) = stronger anonymization, more changed
            beta_d = float(os.getenv("INPAINT_STRENGTH", "0.9"))  # Can be overridden via env
            
            # Scale inference steps with resolution for consistent quality
            # SDXL: Fewer steps needed (25-50) vs SD 1.5 (50-100) due to better architecture
            use_sdxl = is_person and os.getenv("INPAINT_USE_SDXL", "true").lower() == "true"
            if use_sdxl:
                num_steps = int(os.getenv("INPAINT_SDXL_STEPS", "50"))
            else:
                base_steps = {
                    512: int(os.getenv("INPAINT_STEPS_512", "60")),
                    768: int(os.getenv("INPAINT_STEPS_768", "75")),
                    1024: int(os.getenv("INPAINT_STEPS_1024", "100"))
                }
                num_steps = base_steps.get(target_size, int(base_steps[512] * (target_size / 512)))
            
            # Set up generator with seed variation (only if seed is explicitly set)
            # For persons: Add hash-based variation to increase clothing diversity across different masks
            # For objects: Use base seed for consistency
            inpaint_seed_env = os.getenv("INPAINT_SEED", "")
            if inpaint_seed_env:
                base_seed = int(inpaint_seed_env)
                if is_person:
                    # Add variation based on mask path to ensure different persons get different appearances
                    # while remaining deterministic (same mask = same result)
                    import hashlib
                    mask_hash = int(hashlib.md5(mask_path.encode()).hexdigest()[:8], 16)
                    seed = (base_seed + mask_hash) % (2**32)
                    print(f"[InpainterServer] Using person-specific seed: {seed} (base={base_seed}, hash={mask_hash})", file=sys.stderr, flush=True)
                else:
                    seed = base_seed
                generator = torch.Generator(device).manual_seed(seed)
            else:
                # No seed set - use random generation for non-deterministic results
                generator = torch.Generator(device)
                if is_person:
                    print(f"[InpainterServer] Using random seed for person inpainting (non-deterministic)", file=sys.stderr, flush=True)
                else:
                    print(f"[InpainterServer] Using random seed for object inpainting (non-deterministic)", file=sys.stderr, flush=True)
            
            # Call pipeline with appropriate ControlNet based on subject type
            guidance_scale = float(os.getenv("INPAINT_GUIDANCE_SCALE", "7.5"))
            eta = float(os.getenv("INPAINT_ETA", "0.0"))
            
            # Log actual parameters being used
            print(f"[InpainterServer] === ACTUAL PARAMETERS USED ===", file=sys.stderr, flush=True)
            print(f"[InpainterServer] Target size: {target_size}px", file=sys.stderr, flush=True)
            print(f"[InpainterServer] Inference steps: {num_steps}", file=sys.stderr, flush=True)
            print(f"[InpainterServer] Strength (beta_d): {beta_d}", file=sys.stderr, flush=True)
            print(f"[InpainterServer] Guidance scale: {guidance_scale}", file=sys.stderr, flush=True)
            print(f"[InpainterServer] Eta (determinism): {eta}", file=sys.stderr, flush=True)
            print(f"[InpainterServer] Is person: {is_person}", file=sys.stderr, flush=True)
            print(f"[InpainterServer] Using SDXL: {use_sdxl}", file=sys.stderr, flush=True)
            
            if is_person:
                # Person inpainting: Use pose ControlNet (SDXL or SD 1.5)
                pose_scale = float(os.getenv("INPAINT_CONTROLNET_POSE_SCALE", "0.8"))
                print(f"[InpainterServer] Pose ControlNet scale: {pose_scale}", file=sys.stderr, flush=True)
                
                if use_sdxl:
                    # SDXL INPAINTING pipeline call (with pose ControlNet)
                    result_padded = pipe(
                        prompt=positive_prompt,
                        negative_prompt=negative_prompt,
                        image=padded_image,  # Original image for inpainting context
                        mask_image=padded_mask,  # Mask defining region to inpaint
                        control_image=padded_pose,  # Pose ControlNet conditioning
                        num_inference_steps=num_steps,
                        guidance_scale=guidance_scale,
                        strength=beta_d,  # Denoising strength for inpainting
                        controlnet_conditioning_scale=pose_scale,
                        generator=generator,
                        eta=eta,
                    ).images[0]
                else:
                    # SD 1.5 inpainting pipeline call with pose ControlNet
                    result_padded = pipe(
                        prompt=positive_prompt,
                        negative_prompt=negative_prompt,
                        image=padded_image,
                        mask_image=padded_mask,
                        control_image=padded_pose,
                        num_inference_steps=num_steps,  # Scales with resolution
                        guidance_scale=guidance_scale,  # FADM uses standard 7.5
                        strength=beta_d,  # FADM: βd = 0.6 recommended for quality vs. anonymization balance
                        controlnet_conditioning_scale=pose_scale,  # pose guidance only
                        generator=generator,
                        eta=eta,  # Deterministic sampling
                    ).images[0]
            else:
                # Object inpainting: Use Canny ControlNet (SDXL or SD 1.5)
                use_sdxl_objects = os.getenv("INPAINT_OBJECT_USE_SDXL", "true").lower() == "true"
                object_scale = float(os.getenv("INPAINT_CONTROLNET_OBJECT_SCALE", "0.8"))
                print(f"[InpainterServer] Object SDXL enabled: {use_sdxl_objects}", file=sys.stderr, flush=True)
                print(f"[InpainterServer] Canny ControlNet scale: {object_scale}", file=sys.stderr, flush=True)
                
                # Use SDXL steps if enabled, otherwise use scaled SD 1.5 steps
                if use_sdxl_objects and num_steps == int(os.getenv("INPAINT_SDXL_STEPS", "25")):
                    # Already using SDXL steps for objects
                    pass
                elif use_sdxl_objects:
                    # Object is using SDXL, use SDXL steps
                    num_steps = int(os.getenv("INPAINT_SDXL_STEPS", "25"))
                
                result_padded = pipe(
                    prompt=positive_prompt,
                    negative_prompt=negative_prompt,
                    image=padded_image,
                    mask_image=padded_mask,
                    control_image=padded_canny,
                    num_inference_steps=num_steps,
                    guidance_scale=guidance_scale,
                    strength=beta_d,
                    controlnet_conditioning_scale=object_scale,
                    generator=generator,
                    eta=eta,
                ).images[0]

            # Crop and resize result back to cropped size
            offset_x = (target_size - new_w) // 2
            offset_y = (target_size - new_h) // 2
            result_resized = result_padded.crop((offset_x, offset_y, offset_x + new_w, offset_y + new_h))
            result_cropped = result_resized.resize((crop_width, crop_height), Image.LANCZOS)

            # Apply advanced color matching using LAB color space
            # CRITICAL: Only compute statistics from MASKED regions (person/object area)
            # This prevents background colors from contaminating person harmonization
            # Strategy: Skip color matching for both SDXL persons and objects to prevent blue tint
            
            # Check if we should skip color matching for SDXL (persons or objects)
            use_sdxl_objects = os.getenv("INPAINT_OBJECT_USE_SDXL", "true").lower() == "true"
            skip_color_matching_person = is_person and use_sdxl and os.getenv("INPAINT_SDXL_SKIP_COLOR_MATCH", "false").lower() == "true"
            skip_color_matching_object = not is_person and use_sdxl_objects and os.getenv("INPAINT_SDXL_SKIP_COLOR_MATCH", "false").lower() == "true"
            skip_color_matching = skip_color_matching_person or skip_color_matching_object
            
            if skip_color_matching:
                if is_person:
                    print(f"[InpainterServer] SDXL person color matching SKIPPED (preserving native SDXL colors)", file=sys.stderr, flush=True)
                else:
                    print(f"[InpainterServer] SDXL object color matching SKIPPED (preserving native SDXL colors)", file=sys.stderr, flush=True)
            
            if cv2 is not None and not skip_color_matching:
                result_np = np.array(result_cropped)
                cropped_np = np.array(cropped_image)
                
                # Convert to LAB color space for perceptually uniform matching
                result_lab = cv2.cvtColor(result_np, cv2.COLOR_RGB2LAB).astype(np.float32)
                cropped_lab = cv2.cvtColor(cropped_np, cv2.COLOR_RGB2LAB).astype(np.float32)
                
                # Create mask array for computing statistics only from inpainted region
                mask_np = np.array(cropped_mask)
                mask_binary = (mask_np > mask_threshold).astype(bool)
                
                # Only compute statistics from MASKED pixels (not background)
                # This prevents gray backgrounds from making faces look gray
                if np.any(mask_binary):
                    # Match luminance (L channel) - reduced for persons to preserve brightness diversity
                    # For persons: Partial luminance matching to preserve generated clothing brightness
                    if is_person:
                        l_match_strength = float(os.getenv("INPAINT_PERSON_LUMINANCE_MATCH", "0.5"))  # 50% luminance matching
                    else:
                        l_match_strength = 1.0  # Full luminance matching for objects
                    
                    # Compute statistics ONLY from masked region
                    result_l_mean = result_lab[:,:,0][mask_binary].mean()
                    result_l_std = result_lab[:,:,0][mask_binary].std()
                    cropped_l_mean = cropped_lab[:,:,0][mask_binary].mean()
                    cropped_l_std = cropped_lab[:,:,0][mask_binary].std()
                    
                    if result_l_std > 0:
                        matched_l = (result_lab[:,:,0] - result_l_mean) * (cropped_l_std / result_l_std) + cropped_l_mean
                        # Blend matched with original based on strength
                        result_lab[:,:,0] = result_lab[:,:,0] * (1 - l_match_strength) + matched_l * l_match_strength
                    
                    # Match color channels (A and B) - weight depends on subject type
                    # Persons: MINIMAL matching to preserve generated clothing diversity
                    # Objects: Stronger matching for natural background reconstruction
                    lab_strong = float(os.getenv("INPAINT_COLOR_LAB_STRONG", "0.7"))
                    lab_subtle = float(os.getenv("INPAINT_COLOR_LAB_SUBTLE", "0.3"))
                    
                    # For persons: Use minimal chrominance matching to preserve anonymization
                    if is_person:
                        # Use configurable minimal values to preserve generated clothing colors
                        lab_strong = float(os.getenv("INPAINT_PERSON_COLOR_STRONG", "0.05"))  # 5% color transfer
                        lab_subtle = float(os.getenv("INPAINT_PERSON_COLOR_SUBTLE", "0.02"))  # 2% original color influence
                        print(f"[InpainterServer] Person color matching (masked-only): L={l_match_strength:.2f}, strong={lab_strong:.2f}, subtle={lab_subtle:.2f}", file=sys.stderr, flush=True)
                    
                    for i in [1, 2]:
                        # Compute statistics ONLY from masked region
                        result_mean = result_lab[:,:,i][mask_binary].mean()
                        result_std = result_lab[:,:,i][mask_binary].std()
                        cropped_mean = cropped_lab[:,:,i][mask_binary].mean()
                        cropped_std = cropped_lab[:,:,i][mask_binary].std()
                        
                        if result_std > 0:
                            result_lab[:,:,i] = (result_lab[:,:,i] - result_mean) * (cropped_std / result_std) * lab_strong + cropped_mean * lab_subtle + result_mean * (1.0 - lab_strong - lab_subtle)
                
                result_lab[:,:,0] = np.clip(result_lab[:,:,0], 0, 255)
                result_lab[:,:,1] = np.clip(result_lab[:,:,1], 0, 255)
                result_lab[:,:,2] = np.clip(result_lab[:,:,2], 0, 255)
                result_np = cv2.cvtColor(result_lab.astype(np.uint8), cv2.COLOR_LAB2RGB)
                result_cropped = Image.fromarray(result_np)
            else:
                # Fallback: simple brightness transfer
                result_hsv = result_cropped.convert('HSV')
                cropped_hsv = cropped_image.convert('HSV')
                result_hsv_np = np.array(result_hsv)
                cropped_hsv_np = np.array(cropped_hsv)
                mean_v_result = np.mean(result_hsv_np[:, :, 2])
                mean_v_cropped = np.mean(cropped_hsv_np[:, :, 2])
                if mean_v_result > 0:
                    scale = mean_v_cropped / mean_v_result
                    result_hsv_np[:, :, 2] = np.clip(result_hsv_np[:, :, 2] * scale, 0, 255).astype(np.uint8)
                result_cropped = Image.fromarray(result_hsv_np, 'HSV').convert('RGB')

            # Create tight edge-preserving blend to eliminate ghosting
            # Strategy: Slightly ERODE the mask to stay well within YOLOv8m boundaries,
            # then apply minimal feathering only at the very edge
            
            erosion_iterations = int(os.getenv("INPAINT_BLEND_EROSION", "1"))
            feather_iterations = int(os.getenv("INPAINT_BLEND_FEATHER", "1"))
            blur_sigma = float(os.getenv("INPAINT_BLEND_BLUR_SIGMA", "0.5"))
            
            if cv2 is not None:
                # Use original_mask as base (no dilation = no ghost)
                orig_mask_np = np.array(original_mask)
                
                # Erode by 1-2 pixels to ensure we stay inside the true detection boundary
                # This eliminates any ghosting from the inpainting bleeding outside
                kernel_erode = np.ones((3, 3), np.uint8)
                eroded_mask = cv2.erode(orig_mask_np, kernel_erode, iterations=erosion_iterations)
                
                # Create a very narrow 1-2 pixel feather at the edge for seamless blend
                # Dilate eroded mask slightly to create thin transition zone
                kernel_feather = np.ones((3, 3), np.uint8)
                feathered_outer = cv2.dilate(eroded_mask, kernel_feather, iterations=feather_iterations)
                
                # Apply minimal blur ONLY to the feathered edge (1-2 pixels)
                feathered_mask = cv2.GaussianBlur(feathered_outer, (3, 3), blur_sigma)
                
                soft_mask = Image.fromarray(feathered_mask)
            else:
                # Fallback: erode then minimal blur (PIL only)
                soft_mask = original_mask.filter(ImageFilter.MinFilter(size=3))
                soft_mask = soft_mask.filter(ImageFilter.GaussianBlur(radius=1))

            # Composite inpainted result with original cropped image using tight mask
            composited = Image.composite(result_cropped, cropped_image, soft_mask)

            # Paste back into original
            original_image.paste(composited, bbox)
            result_image = original_image

            # Save in same directory as input (preserves subdirectory structure)
            # Use {stem}_anon pattern: output.jpeg -> output_anon.jpeg
            out_path = ipath.parent / f"{ipath.stem}_anon{ipath.suffix}"
            out_path.parent.mkdir(parents=True, exist_ok=True)
            result_image.save(out_path)
            
            # Return ControlNet information for logging
            controlnets_used = []
            if is_person:
                controlnets_used = ["openpose"]  # Both SDXL and SD 1.5 use pose ControlNet
            else:
                controlnets_used = ["canny"]  # Objects use Canny ControlNet with low threshold
            
            return {
                "output_path": str(out_path),
                "controlnets": controlnets_used,
                "is_person": is_person
            }
        except Exception as e:
            return {"output_path": "", "error": f"inpainting failed: {e}"}

    @app.tool()
    def batch_inpaint(image_path: str, instances: list) -> dict:  # noqa: D401
        """Batch inpaint multiple persons in a single call (keeps pipeline loaded).
        
        Args:
            image_path: Path to input image
            instances: List of dicts, each with 'mask_path' and 'description' keys
            
        Returns: {"results": [{"output_path": str}], "final_output": str, "processed_count": int}
        
        Note: Returns ONLY the final output path in results (not intermediate steps).
        """
        if not instances:
            return {"results": [], "final_output": "", "error": "no instances provided"}
        
        import shutil
        import tempfile
        
        # Use working copy to avoid escalating "_anon_anon_anon..." filenames
        ipath = Path(image_path)
        # Create temp file in same directory as input to preserve subdirectory structure
        working_copy = Path(tempfile.mktemp(suffix=ipath.suffix, dir=ipath.parent))
        shutil.copy2(image_path, working_copy)
        
        current_image = str(working_copy)
        
        for i, instance in enumerate(instances):
            mask_path = instance.get('mask_path')
            description = instance.get('description', '')
            
            if not mask_path:
                continue
            
            # Inpaint this instance (pipeline stays loaded)
            result = inpaint(
                image_path=current_image,
                mask_path=mask_path,
                prompt="person",
                description=description
            )
            
            # Chain: use output as input for next instance
            if result.get("output_path"):
                # Move result to working copy location to avoid filename escalation
                result_path = Path(result["output_path"])
                if result_path.exists() and result_path != working_copy:
                    shutil.move(str(result_path), str(working_copy))
                    current_image = str(working_copy)
        
        # Final output: copy working copy to proper output name (single _anon suffix)
        # Use {stem}_anon pattern: 0000000932.jpeg -> 0000000932_anon.jpeg
        final_output_path = ipath.parent / f"{ipath.stem}_anon{ipath.suffix}"
        final_output_path.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(working_copy, final_output_path)
        
        # Cleanup working copy
        try:
            working_copy.unlink()
        except Exception:
            pass
        
        # Return ControlNet information (persons always use openpose)
        return {
            "results": [{"output_path": str(final_output_path), "controlnets": ["openpose"], "is_person": True}],
            "final_output": str(final_output_path),
            "processed_count": len(instances)
        }

    @app.tool()
    def blur(image_path: str, mask_path: str, radius: float | None = None) -> Dict[str, str]:  # noqa: D401
        """Blur masked regions in the image using Gaussian blur.

        Args:
            image_path: Path to input image.
            mask_path: Path to mask image (white=blur, black=keep original). Will be resized to image size if needed.
            radius: Gaussian blur radius (pixels). Defaults to env INPAINT_BLUR_RADIUS or 8.

        Returns: {"output_path": str} or {"output_path": "", "error": str}
        """
        if radius is None:
            radius = float(os.getenv("INPAINT_BLUR_RADIUS", "8"))
        ipath = Path(image_path)
        mpath = Path(mask_path)
        if not ipath.exists():
            return {"output_path": "", "error": f"image not found: {image_path}"}
        if not mpath.exists():
            return {"output_path": "", "error": f"mask not found: {mask_path}"}

        try:
            img = Image.open(ipath).convert("RGB")
            mask = Image.open(mpath).convert("L")
        except Exception as e:
            return {"output_path": "", "error": f"failed to load image or mask: {e}"}

        # Ensure mask matches image size
        if mask.size != img.size:
            mask = mask.resize(img.size, Image.NEAREST)

        # Normalize mask to binary-ish (threshold at 128) to avoid semi-transparent bleed if detector outputs grayscale
        try:
            mask = mask.point(lambda p: 255 if p >= 128 else 0)
        except Exception:
            pass

        try:
            blurred = img.filter(ImageFilter.GaussianBlur(float(radius)))
            composited = Image.composite(blurred, img, mask)
            # Save in same directory as input (overwrites to update canonical output)
            out_path = ipath.parent / ipath.name
            out_path.parent.mkdir(parents=True, exist_ok=True)
            composited.save(out_path)
            return {"output_path": str(out_path)}
        except Exception as e:
            return {"output_path": "", "error": f"blur operation failed: {e}"}

    def main() -> None:  # noqa: D401
        print("[InpainterServer] Starting (transport=stdio)…", file=sys.stderr, flush=True)
        try:
            app.run(transport="stdio")
        except BrokenPipeError:
            # Client closed connection during shutdown - this is expected behavior
            pass
        except KeyboardInterrupt:
            # Graceful shutdown on Ctrl+C
            pass
        except BaseException as e:
            # Check if it's an ExceptionGroup containing BrokenPipeError
            if hasattr(e, '__cause__') or hasattr(e, '__context__'):
                cause = getattr(e, '__cause__', None) or getattr(e, '__context__', None)
                if isinstance(cause, BrokenPipeError):
                    pass  # Expected shutdown
                    return
            # Check for ExceptionGroup with BrokenPipeError
            if e.__class__.__name__ in ('ExceptionGroup', 'BaseExceptionGroup'):
                try:
                    exceptions = getattr(e, 'exceptions', [])
                    if any(isinstance(ex, BrokenPipeError) or 
                           'BrokenPipeError' in str(type(ex)) or
                           'Broken pipe' in str(ex) for ex in exceptions):
                        pass  # Expected shutdown
                        return
                except Exception:
                    pass
            print(f"[InpainterServer] Unexpected error: {e}", file=sys.stderr, flush=True)
            raise


if __name__ == "__main__":
    main()
