from PIL import Image, ImageFilter, ImageDraw
from dataclasses import dataclass
from functools import cached_property
from pathlib import Path
import numpy as np

from .patch_detector import PatchDetector


from .editor_ctx import EditorCtx


@dataclass
class PatchMerger:
    """
    Merges two images based on a detected patch, with various processing options.

    This class takes the heatmap from the PatchDetector and processes it to create
    a final mask. This mask is then used to composite the edited image onto the
    original image. It also provides methods for saving the merged image and
    debug artifacts.
    """

    detector: PatchDetector

    @cached_property
    def ctx(self) -> EditorCtx:
        return self.detector.ctx

    @cached_property
    def raw_mask(self) -> Image.Image:
        heatmap = self.detector.heatmap
        threshold = self.detector.detection_alg.get_threshold(self.ctx.tolerance)
        # Apply threshold
        mask = heatmap.point(lambda p: 255 if p > threshold else 0)

        # Drop smaller areas
        if self.ctx.min_area_size > 1:
            # Filter size must be odd
            filter_size = self.ctx.min_area_size | 1
            # Erode then Dilate (Opening) to remove small noise
            mask = mask.filter(ImageFilter.MinFilter(filter_size))
            mask = mask.filter(ImageFilter.MaxFilter(filter_size))

        return mask

    @cached_property
    def sized_mask(self) -> Image.Image:
        mask = self.raw_mask
        if self.ctx.mask_expansion == 0:
            return mask

        # The kernel size for a radius of N is 2*N+1
        size = abs(self.ctx.mask_expansion) * 2 + 1

        if self.ctx.mask_expansion > 0:
            # Expand (Dilation)
            return mask.filter(ImageFilter.MaxFilter(size))
        else:
            # Shrink (Erosion)
            return mask.filter(ImageFilter.MinFilter(size))

    @cached_property
    def final_mask(self) -> Image.Image:
        mask = self.sized_mask
        if self.ctx.blur_radius > 0:
            mask = mask.filter(ImageFilter.GaussianBlur(self.ctx.blur_radius))
        return mask

    @cached_property
    def merged_image(self) -> Image.Image:
        img_orig, img_edit = self.detector.images
        mask = self.final_mask

        if self.ctx.edit_max_alpha < 1.0:
            img_edit = Image.blend(img_orig, img_edit, self.ctx.edit_max_alpha)

        # The `img_edit` is already warped and pasted on a copy of the original.
        # We use composite to merge the images using the final mask.
        merged = Image.composite(img_edit, img_orig, mask)
        return merged.convert("RGBA")

    @cached_property
    def debug_overlay(self) -> Image.Image:
        # Overlay
        img1 = self.detector.images[0]
        overlay = img1.copy().convert('RGBA')

        # Create a red highlight layer
        highlight = Image.new('RGBA', overlay.size, (255, 0, 0, 128))

        # Blend highlight over overlay
        blended = Image.alpha_composite(overlay, highlight)

        # Composite using mask
        return Image.composite(blended, overlay, self.sized_mask)

    @property
    def base_name(self) -> str:
        if self.ctx.original_img_path:
            return self.ctx.original_img_path.stem
        return "patch_result"

    def save_merged(self, path: Path) -> None:
        self.merged_image.save(path, qmin=0, qmax=0)

    def save_debug(self, output_dir: Path) -> None:
        output_dir.mkdir(parents=True, exist_ok=True)

        self.detector.heatmap.save(output_dir / f"{self.base_name}-heatmap.png")
        self.raw_mask.save(output_dir / f"{self.base_name}-raw_mask.png")
        self.final_mask.save(output_dir / f"{self.base_name}-final_mask.png")
        self.debug_overlay.save(output_dir / f"{self.base_name}-overlay.png")
        self.detector.images[1].save(output_dir / f"{self.base_name}-wrapped.png")