"""Before/after comparison helpers using the visual diff utilities."""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Optional

from PIL import Image

from . import baseline_manager
from .visual_diff import DiffResult, VisualDiffer


@dataclass
class BeforeAfterReport:
    baseline_path: Path
    candidate_path: Path
    diff_path: Optional[Path]
    result: DiffResult


class BeforeAfterComparator:
    """Compare generated screenshots against stored baselines."""

    def __init__(
        self,
        *,
        differ: Optional[VisualDiffer] = None,
        artifacts_dir: Path | str = "reports/visual-diffs",
    ):
        self.differ = differ or VisualDiffer()
        self.artifacts_dir = Path(artifacts_dir)

    def compare_paths(
        self,
        baseline_image: Path | str,
        candidate_image: Path | str,
        *,
        save_diff: bool = True,
    ) -> BeforeAfterReport:
        baseline_path = Path(baseline_image)
        candidate_path = Path(candidate_image)
        if not baseline_path.exists():
            raise FileNotFoundError(f"Baseline image not found: {baseline_path}")
        if not candidate_path.exists():
            raise FileNotFoundError(f"Candidate image not found: {candidate_path}")

        with (
            Image.open(baseline_path) as baseline_img,
            Image.open(candidate_path) as candidate_img,
        ):
            result = self.differ.compare(
                baseline_img, candidate_img, generate_diff=save_diff
            )

        diff_path = None
        if save_diff and not result.passed and result.diff_image is not None:
            diff_dir = self.artifacts_dir
            diff_dir.mkdir(parents=True, exist_ok=True)
            diff_path = diff_dir / f"{candidate_path.stem}_diff.png"
            result.save_diff(diff_path)

        return BeforeAfterReport(
            baseline_path=baseline_path,
            candidate_path=candidate_path,
            diff_path=diff_path,
            result=result,
        )

    def compare_stored_baseline(
        self,
        *,
        tier: str,
        category: str,
        template_name: str,
        candidate_image: Path | str,
        workspace: Path | str = "tests/baselines",
    ) -> BeforeAfterReport:
        workspace_path = Path(workspace)
        baseline_path = workspace_path / tier / category / f"{template_name}.png"
        if not baseline_path.exists():
            raise FileNotFoundError(
                f"Baseline image for {tier}/{category}/{template_name} not found at {baseline_path}"
            )
        return self.compare_paths(baseline_path, candidate_image)


def compare_against_baseline(
    tier: str,
    category: str,
    template_name: str,
    candidate_image: Path | str,
    *,
    workspace: Path | str = "tests/baselines",
    update_if_missing: bool = False,
    differ: Optional[VisualDiffer] = None,
) -> BeforeAfterReport:
    """Convenience helper tying baseline_manager into visual diffs."""

    workspace_path = Path(workspace)
    baseline_path = workspace_path / tier / category / f"{template_name}.png"
    if not baseline_path.exists():
        if update_if_missing:
            baseline_manager.store_baseline_image(
                str(candidate_image), tier, category, template_name
            )
        else:
            raise FileNotFoundError(
                f"Baseline not found for {tier}/{category}/{template_name}"
            )
    comparator = BeforeAfterComparator(differ=differ)
    return comparator.compare_paths(baseline_path, candidate_image)
