#!/usr/bin/env python3
"""Create a checksummed luminance-normalized YOLO dataset without changing labels/splits."""

from __future__ import annotations

import argparse
import hashlib
import json
import shutil
from pathlib import Path

from PIL import Image


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--summary", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--force", action="store_true")
    args = parser.parse_args()
    if args.output_dir.exists():
        if not args.force:
            raise SystemExit(f"Output exists: {args.output_dir}")
        shutil.rmtree(args.output_dir)
    args.output_dir.mkdir(parents=True)

    summary = json.loads(args.summary.read_text(encoding="utf-8"))
    converted = 0
    for tile in summary["tiles"]:
        split = tile["split"]
        source_image = Path(tile["image_path"])
        source_label = Path(tile["label_path"])
        image_target = args.output_dir / "images" / split / source_image.name
        label_target = args.output_dir / "labels" / split / source_label.name
        image_target.parent.mkdir(parents=True, exist_ok=True)
        label_target.parent.mkdir(parents=True, exist_ok=True)
        with Image.open(source_image) as image:
            image.convert("L").convert("RGB").save(image_target)
        shutil.copyfile(source_label, label_target)
        tile["image_path"] = str(image_target)
        tile["label_path"] = str(label_target)
        converted += 1

    dataset_yaml = args.output_dir / "dataset.yaml"
    dataset_yaml.write_text(
        f"path: {args.output_dir}\ntrain: images/train\nval: images/val\nnames:\n  0: building\n",
        encoding="utf-8",
    )
    summary.update(
        {
            "output_dir": str(args.output_dir),
            "dataset_yaml": str(dataset_yaml),
            "preprocessing": "luminance_rgb_replicated",
            "source_summary": str(args.summary),
            "source_summary_sha256": sha256(args.summary),
        }
    )
    output_summary = args.output_dir / "yolo_tile_dataset_summary.json"
    output_summary.write_text(json.dumps(summary, indent=2), encoding="utf-8")
    evidence = {
        "schema_version": 1,
        "status": "ok",
        "preprocessing": "luminance_rgb_replicated",
        "source_summary": str(args.summary),
        "source_summary_sha256": sha256(args.summary),
        "converted_tile_count": converted,
        "output_summary": str(output_summary),
        "output_summary_sha256": sha256(output_summary),
        "dataset_yaml": str(dataset_yaml),
    }
    (args.output_dir / "grayscale-dataset-evidence.json").write_text(
        json.dumps(evidence, indent=2), encoding="utf-8"
    )
    print(json.dumps(evidence, indent=2))
    return 0


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