#!/usr/bin/env python3
"""Read-only checks for a flat PNG/JPEG/WebP + UTF-8 .txt training directory.

Requires Pillow (included by the training environment). Writes no files, starts
no training, and uploads nothing. Exit 1 means structural errors were found.
Warnings do not fail the check. This cannot judge caption accuracy or quality.
"""

import argparse
import hashlib
import json
from pathlib import Path
import sys
import warnings

from PIL import Image, UnidentifiedImageError


IMAGE_SUFFIXES = {".png", ".jpg", ".jpeg", ".webp"}


def inspect_dataset(directory):
    directory = Path(directory)
    report = {"images": 0, "errors": [], "warnings": []}
    if not directory.is_dir():
        report["errors"].append("Dataset directory does not exist.")
        return report

    entries = sorted(directory.iterdir())
    images = [p for p in entries if p.is_file() and p.suffix.lower() in IMAGE_SUFFIXES]
    report["images"] = len(images)
    if not images:
        report["errors"].append("No PNG, JPEG or WebP images found at the top level.")
    if any(p.is_dir() for p in entries):
        report["errors"].append("Use a flat image directory; nested folders are not checked.")

    stems = {}
    hashes = {}
    for path in images:
        key = path.stem.casefold()
        if key in stems:
            report["errors"].append(f"{path.name}: caption name collides with {stems[key]}.")
        stems[key] = path.name
        try:
            with path.open("rb") as source:
                digest = hashlib.file_digest(source, "sha256").hexdigest()
            if digest in hashes:
                report["warnings"].append(f"{path.name}: byte-identical to {hashes[digest]}.")
            hashes[digest] = path.name
            with warnings.catch_warnings():
                warnings.simplefilter("error", Image.DecompressionBombWarning)
                with Image.open(path) as im:
                    im.verify()
                with Image.open(path) as im:
                    im.load()
                    if min(im.size) < 512:
                        report["warnings"].append(f"{path.name}: short side below 512 px ({im.width}x{im.height}).")
                    if im.mode != "RGB":
                        report["warnings"].append(f"{path.name}: mode {im.mode}; review transparency/color conversion.")
                    if getattr(im, "n_frames", 1) != 1:
                        report["errors"].append(f"{path.name}: animated images are outside this example.")
                    if im.getexif().get(274, 1) != 1:
                        report["warnings"].append(f"{path.name}: apply EXIF orientation before training.")
        except (OSError, ValueError, UnidentifiedImageError, Image.DecompressionBombError, Image.DecompressionBombWarning) as exc:
            report["errors"].append(f"{path.name}: cannot decode safely ({type(exc).__name__}).")

        caption = path.with_suffix(".txt")
        try:
            text = caption.read_text(encoding="utf-8")
            if not text.strip():
                report["errors"].append(f"{caption.name}: caption is empty.")
            if "\x00" in text or text.startswith("\ufeff"):
                report["errors"].append(f"{caption.name}: remove NUL/BOM characters.")
        except (OSError, UnicodeError):
            report["errors"].append(f"{caption.name}: missing, unreadable or not UTF-8.")

    for path in entries:
        if path.is_file() and path.suffix.lower() == ".txt" and path.stem.casefold() not in stems:
            report["warnings"].append(f"{path.name}: caption has no matching supported image.")
    return report


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("directory", type=Path)
    args = parser.parse_args()
    try:
        result = inspect_dataset(args.directory)
    except OSError as exc:
        print(f"Cannot read dataset: {exc}", file=sys.stderr)
        return 1
    print(json.dumps(result, ensure_ascii=False, indent=2))
    return 1 if result["errors"] else 0


if __name__ == "__main__":
    sys.exit(main())