#!/usr/bin/env python3
from __future__ import annotations

import argparse
import logging
import re
import shutil
from collections import Counter
from pathlib import Path
from typing import Iterable

from PIL import Image


DEFAULT_SRC_ROOT = Path("/mnt/nfs/mono3d/ydong_data/Detection/2Ddetection_20260402")
DEFAULT_DST_IMAGES_DIR = Path(
    "/mnt/nfs/mono3d/xdzhu_data/Mono3d/Mono3d_4face_2m_g1m3/driving_png_20260202/Detection2D/D4Q_2D/images"
)
DEFAULT_DST_LABELS_DIR = Path(
    "/mnt/nfs/mono3d/xdzhu_data/Mono3d/Mono3d_4face_2m_g1m3/driving_png_20260320/Detection2D/D4Q/labels"
)
IMAGE_SUFFIXES = {".jpg", ".jpeg"}
LABEL_SUFFIXES = {".txt"}
DEFAULT_NAME_TOKENS = ("D01", "D4Q", "D4Q2", "D01P")
DEFAULT_SOURCE_SIZE = (3840, 2160)
DEFAULT_TARGET_SIZE = (1920, 1080)

try:
    RESAMPLE_LANCZOS = Image.Resampling.LANCZOS
except AttributeError:  # pragma: no cover - Pillow < 9.1 compatibility
    RESAMPLE_LANCZOS = Image.LANCZOS


def parse_size(value: str) -> tuple[int, int]:
    match = re.fullmatch(r"\s*(\d+)\s*[xX]\s*(\d+)\s*", value)
    if match is None:
        raise argparse.ArgumentTypeError(f"Invalid size '{value}'. Expected WIDTHxHEIGHT, e.g. 3840x2160.")

    width, height = (int(group) for group in match.groups())
    if width <= 0 or height <= 0:
        raise argparse.ArgumentTypeError(f"Image size must be positive, got {value}.")
    return width, height


def parse_name_tokens(raw_tokens: Iterable[str] | None) -> tuple[str, ...]:
    if raw_tokens is None:
        return DEFAULT_NAME_TOKENS

    tokens: list[str] = []
    for raw_token in raw_tokens:
        tokens.extend(piece.strip() for piece in raw_token.split(","))

    normalized = tuple(dict.fromkeys(token for token in tokens if token))
    if not normalized:
        raise ValueError("At least one non-empty --name-token value is required.")
    return normalized


def build_name_pattern(name_tokens: tuple[str, ...]) -> re.Pattern[str]:
    alternatives = "|".join(re.escape(token.lower()) for token in sorted(name_tokens, key=len, reverse=True))
    # Match dataset tokens as standalone parts of the file name so strings like "...17b78d01..."
    # are not treated as a valid D01 sample.
    return re.compile(rf"(^|[^0-9a-z])(?:{alternatives})(?=[^0-9a-z]|$)", re.IGNORECASE)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Copy Detection2D images whose names match the requested tokens, resize the image, "
            "and copy the paired labels."
        )
    )
    parser.add_argument("--src-root", type=Path, default=DEFAULT_SRC_ROOT, help="Source Detection2D dataset root.")
    parser.add_argument(
        "--dst-images-dir",
        type=Path,
        default=DEFAULT_DST_IMAGES_DIR,
        help="Destination directory for resized image files.",
    )
    parser.add_argument(
        "--dst-labels-dir",
        type=Path,
        default=DEFAULT_DST_LABELS_DIR,
        help="Destination directory for paired label files.",
    )
    parser.add_argument(
        "--name-token",
        action="append",
        dest="name_tokens",
        default=None,
        help=(
            "Case-insensitive dataset token to match in the file name. Repeat the flag or provide "
            "a comma-separated list. Defaults to D01, D4Q, D4Q2, D01P."
        ),
    )
    parser.add_argument(
        "--source-size",
        type=parse_size,
        default=DEFAULT_SOURCE_SIZE,
        help="Only process images with this source size (WIDTHxHEIGHT).",
    )
    parser.add_argument(
        "--target-size",
        type=parse_size,
        default=DEFAULT_TARGET_SIZE,
        help="Resize matching images to this size (WIDTHxHEIGHT).",
    )
    parser.add_argument(
        "--overwrite",
        action="store_true",
        help="Overwrite destination files if they already exist.",
    )
    parser.add_argument(
        "--limit",
        type=int,
        default=None,
        help="Only inspect the first N matching image names.",
    )
    parser.add_argument(
        "--log-every",
        type=int,
        default=1000,
        help="Log progress every N inspected image files.",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Print planned work without writing files.",
    )
    return parser.parse_args()


def iter_matching_images(
    src_images_dir: Path,
    suffixes: set[str],
    name_pattern: re.Pattern[str],
    limit: int | None = None,
) -> list[Path]:
    matches: list[Path] = []
    for path in src_images_dir.rglob("*"):
        if not path.is_file() or path.suffix.lower() not in suffixes:
            continue
        if not name_pattern.search(path.stem.lower()):
            continue

        matches.append(path)
        if limit is not None and len(matches) >= limit:
            break

    return matches


def read_image_size(image_path: Path) -> tuple[int, int]:
    with Image.open(image_path) as image:
        return image.size


def resize_one_image(
    src_path: Path,
    dst_dir: Path,
    target_size: tuple[int, int],
    overwrite: bool,
    dry_run: bool,
) -> tuple[str, Path, str]:
    dst_path = dst_dir / src_path.name

    if dst_path.exists() and not overwrite:
        return "skipped", dst_path, ""

    if dry_run:
        return "planned", dst_path, ""

    dst_path.parent.mkdir(parents=True, exist_ok=True)
    try:
        with Image.open(src_path) as image:
            if image.mode not in {"RGB", "L"}:
                image = image.convert("RGB")
            resized = image.resize(target_size, RESAMPLE_LANCZOS)
            save_kwargs = {"quality": 95} if dst_path.suffix.lower() in {".jpg", ".jpeg"} else {}
            resized.save(dst_path, **save_kwargs)
    except OSError as exc:
        return "failed", dst_path, str(exc)

    return "copied", dst_path, ""


def copy_one_label(src_path: Path, dst_dir: Path, overwrite: bool, dry_run: bool) -> tuple[str, Path, str]:
    dst_path = dst_dir / src_path.name

    if dst_path.exists() and not overwrite:
        return "skipped", dst_path, ""

    if dry_run:
        return "planned", dst_path, ""

    dst_path.parent.mkdir(parents=True, exist_ok=True)
    try:
        shutil.copy2(src_path, dst_path)
    except OSError as exc:
        return "failed", dst_path, str(exc)

    return "copied", dst_path, ""


def maybe_log_plan(image_path: Path, label_path: Path, dst_images_dir: Path, dst_labels_dir: Path, index: int) -> None:
    if index >= 5:
        return

    logging.info("plan image: %s -> %s", image_path, dst_images_dir / image_path.name)
    logging.info("plan label: %s -> %s", label_path, dst_labels_dir / label_path.name)


def log_progress(index: int, total: int, counts: Counter) -> None:
    logging.info(
        (
            "progress %d/%d | matched=%d wrong_size=%d missing_label=%d "
            "images(copied=%d skipped=%d planned=%d failed=%d) "
            "labels(copied=%d skipped=%d planned=%d failed=%d)"
        ),
        index,
        total,
        counts["pairs_matched"],
        counts["images_wrong_size"],
        counts["labels_missing"],
        counts["images_copied"],
        counts["images_skipped"],
        counts["images_planned"],
        counts["images_failed"],
        counts["labels_copied"],
        counts["labels_skipped"],
        counts["labels_planned"],
        counts["labels_failed"],
    )


def process_dataset(
    image_paths: list[Path],
    src_labels_dir: Path,
    dst_images_dir: Path,
    dst_labels_dir: Path,
    source_size: tuple[int, int],
    target_size: tuple[int, int],
    overwrite: bool,
    dry_run: bool,
    log_every: int,
) -> Counter:
    counts: Counter[str] = Counter()
    total = len(image_paths)

    for index, image_path in enumerate(image_paths, start=1):
        label_path = src_labels_dir / f"{image_path.stem}.txt"
        if label_path.suffix.lower() not in LABEL_SUFFIXES:
            counts["labels_invalid_suffix"] += 1
        elif not label_path.exists():
            counts["labels_missing"] += 1
            logging.warning("missing label for image: %s", image_path)
        else:
            try:
                image_size = read_image_size(image_path)
            except OSError as exc:
                counts["images_failed"] += 1
                logging.error("failed to read image size: %s | %s", image_path, exc)
            else:
                if image_size != source_size:
                    counts["images_wrong_size"] += 1
                    if counts["images_wrong_size"] <= 10:
                        logging.info("skip image with unexpected size %s: %s", image_size, image_path)
                else:
                    counts["pairs_matched"] += 1
                    if dry_run:
                        maybe_log_plan(image_path, label_path, dst_images_dir, dst_labels_dir, counts["pairs_matched"] - 1)

                    image_status, dst_image_path, image_message = resize_one_image(
                        src_path=image_path,
                        dst_dir=dst_images_dir,
                        target_size=target_size,
                        overwrite=overwrite,
                        dry_run=dry_run,
                    )
                    counts[f"images_{image_status}"] += 1
                    if image_status == "failed":
                        logging.error("image write failed: %s -> %s | %s", image_path, dst_image_path, image_message)
                    else:
                        label_status, dst_label_path, label_message = copy_one_label(
                            src_path=label_path,
                            dst_dir=dst_labels_dir,
                            overwrite=overwrite,
                            dry_run=dry_run,
                        )
                        counts[f"labels_{label_status}"] += 1
                        if label_status == "failed":
                            logging.error("label copy failed: %s -> %s | %s", label_path, dst_label_path, label_message)

        if index % log_every == 0 or index == total:
            log_progress(index, total, counts)

    return counts


def main() -> int:
    args = parse_args()
    logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(message)s")

    src_root = args.src_root.resolve()
    src_images_dir = src_root / "images"
    src_labels_dir = src_root / "labels"
    dst_images_dir = args.dst_images_dir.resolve()
    dst_labels_dir = args.dst_labels_dir.resolve()
    name_tokens = parse_name_tokens(args.name_tokens)
    name_pattern = build_name_pattern(name_tokens)

    if not src_root.exists():
        raise FileNotFoundError(f"Source root does not exist: {src_root}")
    if not src_images_dir.exists():
        raise FileNotFoundError(f"Source images directory does not exist: {src_images_dir}")
    if not src_labels_dir.exists():
        raise FileNotFoundError(f"Source labels directory does not exist: {src_labels_dir}")
    if args.log_every <= 0:
        raise ValueError("--log-every must be a positive integer.")
    if args.limit is not None and args.limit <= 0:
        raise ValueError("--limit must be a positive integer when provided.")

    image_paths = iter_matching_images(src_images_dir, IMAGE_SUFFIXES, name_pattern, limit=args.limit)

    logging.info("source root: %s", src_root)
    logging.info("destination images dir: %s", dst_images_dir)
    logging.info("destination labels dir: %s", dst_labels_dir)
    logging.info("file name tokens: %s", ", ".join(name_tokens))
    logging.info("required source size: %sx%s", *args.source_size)
    logging.info("target resize: %sx%s", *args.target_size)
    logging.info(
        "labels are copied without rewriting because the Detection2D labels are normalized; "
        "a full-image resize keeps normalized coordinates unchanged."
    )
    logging.info("found %d candidate image(s) by name", len(image_paths))

    counts = process_dataset(
        image_paths=image_paths,
        src_labels_dir=src_labels_dir,
        dst_images_dir=dst_images_dir,
        dst_labels_dir=dst_labels_dir,
        source_size=args.source_size,
        target_size=args.target_size,
        overwrite=args.overwrite,
        dry_run=args.dry_run,
        log_every=args.log_every,
    )

    if counts["images_failed"] or counts["labels_failed"]:
        logging.error(
            "copy finished with failures | image_failed=%d label_failed=%d",
            counts["images_failed"],
            counts["labels_failed"],
        )
        return 1

    logging.info(
        (
            "copy finished successfully | matched=%d wrong_size=%d missing_label=%d "
            "images(copied=%d skipped=%d planned=%d) labels(copied=%d skipped=%d planned=%d)"
        ),
        counts["pairs_matched"],
        counts["images_wrong_size"],
        counts["labels_missing"],
        counts["images_copied"],
        counts["images_skipped"],
        counts["images_planned"],
        counts["labels_copied"],
        counts["labels_skipped"],
        counts["labels_planned"],
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())