import argparse
from pathlib import Path

import cv2

from cleaning import (
    clean_display_image_and_mask,
    cleaned_image_name,
    read_image_rgb_and_preview,
    write_image_rgb,
)


def main(
    image_paths: list[str],
    output_dir: str | Path = "out",
) -> None:
    output_dir = Path(output_dir)
    cleaned_dir = output_dir / "cleaned"
    masks_dir = output_dir / "masks"
    cleaned_dir.mkdir(parents=True, exist_ok=True)
    masks_dir.mkdir(parents=True, exist_ok=True)

    for image_path_string in image_paths:
        image_path = Path(image_path_string)
        image_rgb, display_rgb = read_image_rgb_and_preview(image_path)
        cleaned_rgb, debris_mask = clean_display_image_and_mask(
            image_rgb,
            display_rgb,
        )

        cv2.imwrite(str(masks_dir / f"{image_path.stem}_mask.png"), debris_mask)
        cleaned_path = cleaned_dir / cleaned_image_name(image_path)
        write_image_rgb(cleaned_path, cleaned_rgb)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Clean debris from IHC images.")
    parser.add_argument(
        "images",
        nargs="+",
        help="Paths to one or more input images.",
    )
    parser.add_argument(
        "--output-dir",
        default="out",
        help="Directory for cleaned images and masks (default: out).",
    )
    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()
    main(args.images, args.output_dir)