#!/usr/bin/env -S uv run --script
# /// script
# dependencies = ["Pillow"]
# ///

import argparse
import collections
import datetime as dt
import logging
import os
import re
import shutil
import subprocess
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable

from PIL import Image, UnidentifiedImageError


logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s")


DATE_PATTERNS = (
    re.compile(r"^(?P<year>\d{4})-(?P<month>\d{2})-(?P<day>\d{2})"),
    re.compile(r"^(?P<year>\d{4})\.(?P<month>\d{2})\.(?P<day>\d{2})"),
    re.compile(r"^(?P<year>\d{4})(?P<month>\d{2})(?P<day>\d{2})"),
)

EXIF_DATETIME_TAGS = (36867, 36868, 306)


@dataclass
class ImageEntry:
    path: Path
    group_key: str
    filename_date: dt.datetime | None
    exif_date: dt.datetime | None = None


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Update file timestamps for images grouped by filename pattern",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "directory",
        nargs="?",
        default=Path.cwd(),
        type=Path,
        help="Root directory to scan for image files",
    )
    parser.add_argument(
        "--recursive",
        action="store_true",
        help="Recurse into subdirectories",
    )
    parser.add_argument(
        "--extensions",
        default=".jpg,.jpeg,.png,.tif,.tiff,.heic,.heif",
        help="Comma-separated list of extensions to include",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Show planned changes without touching the filesystem",
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="Enable debug logging",
    )
    return parser.parse_args()


def configure_logging(verbose: bool) -> None:
    if verbose:
        logging.getLogger().setLevel(logging.DEBUG)


def collect_image_paths(directory: Path, extensions: set[str], recursive: bool) -> Iterable[Path]:
    iterator = directory.rglob("*") if recursive else directory.glob("*")
    for candidate in iterator:
        if candidate.is_file() and candidate.suffix.lower() in extensions:
            yield candidate


def parse_group_and_date(path: Path) -> tuple[str, dt.datetime | None] | None:
    name = path.stem
    if " -- " not in name:
        logging.debug("Skipping %s: missing group delimiter", path)
        return None
    try:
        group_part, rest = name.rsplit(" -- ", 1)
    except ValueError:
        logging.debug("Skipping %s: unable to split group and tail", path)
        return None
    if not group_part:
        logging.debug("Skipping %s: empty group part", path)
        return None
    date_candidate = extract_date_from_tail(rest)
    if " -- " not in group_part and "__" in rest:
        album_part, _, _ = rest.partition("__")
        album_part = album_part.strip()
        if album_part:
            group_part = f"{group_part} -- {album_part}"
    return group_part, date_candidate


def extract_date_from_tail(text: str) -> dt.datetime | None:
    for pattern in DATE_PATTERNS:
        match = pattern.match(text)
        if match:
            try:
                return dt.datetime(
                    year=int(match.group("year")),
                    month=int(match.group("month")),
                    day=int(match.group("day")),
                )
            except ValueError:
                return None
    return None


def read_exif_datetime(path: Path) -> dt.datetime | None:
    try:
        with Image.open(path) as image:
            exif = image.getexif()
    except (UnidentifiedImageError, OSError) as error:
        logging.debug("Could not read EXIF from %s: %s", path, error)
        return None
    if not exif:
        return None
    for tag in EXIF_DATETIME_TAGS:
        value = exif.get(tag)
        if not value:
            continue
        if isinstance(value, bytes):
            value = value.decode(errors="ignore")
        if not isinstance(value, str):
            continue
        value = value.strip()
        for fmt in ("%Y:%m:%d %H:%M:%S", "%Y-%m-%d %H:%M:%S"):
            try:
                return dt.datetime.strptime(value, fmt)
            except ValueError:
                continue
    return None


def most_common_datetime(candidates: list[dt.datetime]) -> dt.datetime | None:
    if not candidates:
        return None
    counter = collections.Counter(candidates)
    most_common = counter.most_common()
    top_count = most_common[0][1]
    top_values = [item for item, count in most_common if count == top_count]
    return min(top_values)


def determine_group_date(entries: list[ImageEntry]) -> dt.datetime | None:
    name_dates = [entry.filename_date for entry in entries if entry.filename_date]
    if name_dates:
        logging.debug("Using filename-derived date for group %s", entries[0].group_key)
        return most_common_datetime(name_dates)
    exif_dates: list[dt.datetime] = []
    for entry in entries:
        if entry.exif_date is None:
            entry.exif_date = read_exif_datetime(path=entry.path)
        if entry.exif_date:
            exif_dates.append(entry.exif_date)
    if exif_dates:
        logging.debug("Using EXIF-derived date for group %s", entries[0].group_key)
    return most_common_datetime(exif_dates)


def apply_timestamp(path: Path, target: dt.datetime, dry_run: bool) -> None:
    timestamp = target.timestamp()
    if dry_run:
        logging.debug("DRY-RUN %s -> %s", path, target.isoformat(sep=" "))
        return
    os.utime(path, times=(timestamp, timestamp))


def chunked(items: list[str], size: int) -> Iterable[list[str]]:
    for index in range(0, len(items), size):
        yield items[index : index + size]


def apply_setfile_batch(entries: list[ImageEntry], target: dt.datetime, dry_run: bool) -> None:
    if sys.platform != "darwin":
        return
    setfile = shutil.which("SetFile")
    if not setfile:
        return
    formatted = target.strftime("%m/%d/%Y %H:%M:%S")
    paths = [str(entry.path) for entry in entries]
    for batch in chunked(items=paths, size=64):
        if dry_run:
            logging.debug("DRY-RUN SetFile %s files -> %s", len(batch), formatted)
            continue
        try:
            subprocess.run([setfile, "-d", formatted, *batch], check=True)
            subprocess.run([setfile, "-m", formatted, *batch], check=True)
        except subprocess.CalledProcessError as error:
            logging.debug("SetFile failed for %s files: %s", len(batch), error)


def update_group(entries: list[ImageEntry], dry_run: bool) -> bool:
    target = determine_group_date(entries=entries)
    if target is None:
        logging.warning("No date found for group %s", entries[0].group_key)
        return False
    for entry in entries:
        apply_timestamp(path=entry.path, target=target, dry_run=dry_run)
    apply_setfile_batch(entries=entries, target=target, dry_run=dry_run)
    action = "DRY-RUN" if dry_run else "UPDATED"
    logging.info(
        "%s %s (%s files) -> %s",
        action,
        entries[0].group_key,
        len(entries),
        target.isoformat(sep=" "),
    )
    return True


def process_directory(directory: Path, extensions: set[str], recursive: bool, dry_run: bool) -> None:
    groups: dict[str, list[ImageEntry]] = collections.defaultdict(list)
    for path in collect_image_paths(directory=directory, extensions=extensions, recursive=recursive):
        parsed = parse_group_and_date(path=path)
        if parsed is None:
            continue
        group_key, filename_date = parsed
        groups[group_key].append(ImageEntry(path=path, group_key=group_key, filename_date=filename_date))
    updated = 0
    skipped = 0
    for entries in groups.values():
        if update_group(entries=entries, dry_run=dry_run):
            updated += len(entries)
        else:
            skipped += len(entries)
    logging.info("Updated %s files; skipped %s files", updated, skipped)


def normalize_extensions(raw: str) -> set[str]:
    pieces = re.split(r"[;,]", raw)
    results: set[str] = set()
    for piece in pieces:
        trimmed = piece.strip().lower()
        if not trimmed:
            continue
        if not trimmed.startswith("."):
            trimmed = f".{trimmed}"
        results.add(trimmed)
    return results


def main() -> None:
    args = parse_args()
    configure_logging(verbose=args.verbose)
    extensions = normalize_extensions(raw=args.extensions)
    if not extensions:
        logging.error("No valid extensions provided")
        raise SystemExit(1)
    if not args.directory.exists():
        logging.error("Directory %s does not exist", args.directory)
        raise SystemExit(1)
    process_directory(
        directory=args.directory,
        extensions=extensions,
        recursive=args.recursive,
        dry_run=args.dry_run,
    )


if __name__ == "__main__":
    main()