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

from .editor_ctx import EditorCtx
from .patch_detector import PatchDetector
from .patch_merger import PatchMerger


@dataclass
class PatchEditor:
    ctx: EditorCtx

    @cached_property
    def patch_detector(self) -> PatchDetector:
        return PatchDetector(self.ctx)

    @cached_property
    def patch_merger(self) -> PatchMerger:
        return PatchMerger(self.patch_detector)

    @staticmethod
    def from_node(
        original_image: Image.Image,
        edited_image: Image.Image,
        detection_alg: str,
        region_size: int,
        stride: int,
        min_area_size: int,
        tolerance: float,
        blur_radius: float,
        mask_expansion: int,
        edit_max_alpha: float,
        allow_rescale: bool,
        original_image_path: Optional[Path] = None,
        edited_image_path: Optional[Path] = None,
    ) -> "PatchEditor":
        ctx = EditorCtx(
            original_img=original_image,
            edited_img=edited_image,
            original_img_path=original_image_path,
            edited_img_path=edited_image_path,
            detection_alg=detection_alg,
            region_size=region_size,
            stride=stride,
            allow_rescale=allow_rescale,
            min_area_size=min_area_size,
            tolerance=tolerance,
            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,
        original_image_path=original_path,
        detection_alg="difference",
        region_size=8,
        stride=1,
        min_area_size=4,
        tolerance=0.5,
        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()