from PIL import Image
from dataclasses import dataclass
from functools import cached_property
from pathlib import Path
from typing import Optional

from .patch_detector import PatchDetector
from .patch_merger import PatchMerger


@dataclass
class EditorCtx:
    original_img: Image.Image
    edited_img: Image.Image
    detection_alg: str
    region_size: int
    stride: int
    allow_rescale: bool
    min_area_size: int
    threshold: int
    blur_radius: float
    mask_expansion: int
    edit_max_alpha: float


@dataclass
class PatchEditor:
    ctx: EditorCtx

    @cached_property
    def patch_detector(self) -> PatchDetector:
        return PatchDetector(
            original_img=self.ctx.original_img,
            edited_img=self.ctx.edited_img,
            region_size=self.ctx.region_size,
            stride=self.ctx.stride,
            allow_rescale=self.ctx.allow_rescale,
            detection_alg=self.ctx.detection_alg,
        )

    @cached_property
    def patch_merger(self) -> PatchMerger:
        return PatchMerger(
            detector=self.patch_detector,
            min_area_size=self.ctx.min_area_size,
            threshold=self.ctx.threshold,
            blur_radius=self.ctx.blur_radius,
            mask_expansion=self.ctx.mask_expansion,
            edit_max_alpha=self.ctx.edit_max_alpha,
        )

    @staticmethod
    def from_node(
        original_image: Image.Image,
        edited_image: Image.Image,
        detection_alg: str,
        region_size: int,
        stride: int,
        min_area_size: int,
        threshold: int,
        blur_radius: float,
        mask_expansion: int,
        edit_max_alpha: float,
        allow_rescale: bool,
    ) -> "PatchEditor":
        ctx = EditorCtx(
            original_img=original_image,
            edited_img=edited_image,
            detection_alg=detection_alg,
            region_size=region_size,
            stride=stride,
            allow_rescale=allow_rescale,
            min_area_size=min_area_size,
            threshold=threshold,
            blur_radius=blur_radius,
            mask_expansion=mask_expansion,
            edit_max_alpha=edit_max_alpha,
        )
        return PatchEditor(ctx)


def main() -> None:
    work_dir = Path('/tmp/work')
    original_path = work_dir / 'a.png'
    edited_path = work_dir / 'a-patch.png'
    output_path = work_dir / 'a-patched.png'

    original_img = Image.open(original_path)
    edited_img = Image.open(edited_path)

    editor = PatchEditor.from_node(
        original_image=original_img,
        edited_image=edited_img,
        detection_alg="difference",
        region_size=8,
        stride=1,
        min_area_size=4,
        threshold=10,
        blur_radius=3.0,
        mask_expansion=0,
        edit_max_alpha=1.0,
        allow_rescale=False,
    )

    merger = editor.patch_merger

    print(f"Processing {original_path} and {edited_path}...")
    merger.save_merged(output_path)
    print(f"Saved merged image to {output_path}")

    merger.save_debug(work_dir)
    print(f"Saved debug images to {work_dir}")


if __name__ == '__main__':
    main()