#!/usr/bin/env python3
"""Derive a checksum-bound, training-disabled non-overlapping YOLO val view."""

from __future__ import annotations

import argparse
import hashlib
import json
import shutil
from collections import defaultdict
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

from PIL import Image

try:
    from scripts.audit_yolo_cross_tile_repetition import tile_offsets
except ModuleNotFoundError:  # Standalone operator-tool copy beside the auditor.
    from audit_yolo_cross_tile_repetition import tile_offsets


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


def select_nonoverlap_tiles(
    tiles: list[dict[str, Any]], tile_size: int
) -> list[dict[str, Any]]:
    selected: list[dict[str, Any]] = []
    for tile in tiles:
        if not isinstance(tile, dict) or not tile.get("kept", True):
            continue
        if str(tile.get("split") or "") != "val":
            continue
        row_offset, column_offset = tile_offsets(str(tile.get("image_path") or ""))
        if row_offset % tile_size == 0 and column_offset % tile_size == 0:
            selected.append(tile)
    return sorted(
        selected,
        key=lambda tile: (
            str(tile.get("sample_slug") or ""),
            int(tile.get("tile_index") or 0),
        ),
    )


def aggregate_hash(paths: list[Path], root: Path) -> str:
    digest = hashlib.sha256()
    for path in sorted(paths):
        digest.update(path.relative_to(root).as_posix().encode("utf-8"))
        digest.update(b"\0")
        digest.update(sha256_file(path).encode("ascii"))
        digest.update(b"\n")
    return digest.hexdigest()


def image_has_low_visual_variance(image_path: Path, blank_range_threshold: int) -> bool:
    image = Image.open(image_path).convert("L")
    minimum, maximum = image.getextrema()
    return maximum - minimum <= blank_range_threshold


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--source-dir", required=True, type=Path)
    parser.add_argument("--output-dir", required=True, type=Path)
    parser.add_argument(
        "--exclude-low-variance",
        action=argparse.BooleanOptionalAction,
        default=True,
    )
    parser.add_argument("--blank-range-threshold", type=int, default=3)
    args = parser.parse_args()

    source_dir = args.source_dir.expanduser().resolve(strict=True)
    output_dir = args.output_dir.expanduser().resolve(strict=False)
    if output_dir.exists():
        parser.error(f"output directory already exists: {output_dir}")
    source_summary_path = source_dir / "yolo_tile_dataset_summary.json"
    source_summary = json.loads(source_summary_path.read_text(encoding="utf-8"))
    tiles = source_summary.get("tiles")
    tile_size = source_summary.get("tile_size")
    if not isinstance(tiles, list):
        raise ValueError("source summary must contain a tiles list")
    if not isinstance(tile_size, int) or tile_size <= 0:
        raise ValueError("source summary must contain a positive integer tile_size")
    selected_candidates = select_nonoverlap_tiles(tiles, tile_size)
    if not selected_candidates:
        raise ValueError("no non-overlapping validation tiles were selected")

    selected: list[dict[str, Any]] = []
    excluded_tiles: list[dict[str, Any]] = []
    for tile in selected_candidates:
        source_image = Path(str(tile.get("image_path") or "")).resolve(strict=True)
        if args.exclude_low_variance and image_has_low_visual_variance(
            source_image, args.blank_range_threshold
        ):
            excluded_tiles.append(
                {
                    "sample_slug": str(tile.get("sample_slug") or "unknown"),
                    "tile_index": int(tile.get("tile_index") or 0),
                    "source_image_path": str(source_image),
                    "reason": "low_visual_variance_no_data",
                }
            )
            continue
        selected.append(tile)
    if not selected:
        raise ValueError("every non-overlapping validation tile was excluded")

    offsets_by_sample: dict[str, set[tuple[int, int]]] = defaultdict(set)
    source_val_counts: dict[str, int] = defaultdict(int)
    for tile in tiles:
        if (
            isinstance(tile, dict)
            and tile.get("kept", True)
            and tile.get("split") == "val"
        ):
            source_val_counts[str(tile.get("sample_slug") or "unknown")] += 1
    for tile in selected:
        offsets_by_sample[str(tile.get("sample_slug") or "unknown")].add(
            tile_offsets(str(tile.get("image_path") or ""))
        )
    accounted_samples = set(offsets_by_sample) | {
        tile["sample_slug"] for tile in excluded_tiles
    }
    if accounted_samples != set(source_val_counts):
        raise ValueError("non-overlap selection has an unaccounted validation sample")

    output_dir.mkdir(parents=True)
    (output_dir / "images" / "train").mkdir(parents=True)
    (output_dir / "labels" / "train").mkdir(parents=True)
    output_tiles: list[dict[str, Any]] = []
    output_files: list[Path] = []
    for tile in selected:
        source_image = Path(str(tile.get("image_path") or "")).resolve(strict=True)
        source_label = Path(str(tile.get("label_path") or "")).resolve(strict=True)
        target_image = output_dir / "images" / "val" / source_image.name
        target_label = output_dir / "labels" / "val" / source_label.name
        target_image.parent.mkdir(parents=True, exist_ok=True)
        target_label.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(source_image, target_image)
        shutil.copy2(source_label, target_label)
        output_files.extend((target_image, target_label))
        derived_tile = dict(tile)
        derived_tile.update(
            {
                "image_path": str(target_image),
                "label_path": str(target_label),
                "derived_from_image_path": str(source_image),
                "derived_from_label_path": str(source_label),
            }
        )
        output_tiles.append(derived_tile)

    dataset_yaml = output_dir / "dataset.yaml"
    dataset_yaml.write_text(
        "\n".join(
            (
                f"path: {output_dir}",
                "train: images/train",
                "val: images/val",
                "names:",
                "  0: building",
                "",
            )
        ),
        encoding="utf-8",
    )
    no_training = {
        "schema_version": 1,
        "training_prohibited": True,
        "reason": "Immutable non-overlapping evaluation view; train directories are intentionally empty.",
        "source_summary_sha256": sha256_file(source_summary_path),
    }
    no_training_path = output_dir / "NO_TRAINING.json"
    no_training_path.write_text(
        json.dumps(no_training, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    summary = {
        "schema_version": 1,
        "status": "ready_evaluation_only",
        "claim_boundary": (
            "Non-protected, non-overlapping tile-level validation view; no "
            "training, protected-test, object-level independence or promotion claim."
        ),
        "dataset_yaml": str(dataset_yaml),
        "output_dir": str(output_dir),
        "source_dataset_dir": str(source_dir),
        "source_summary_path": str(source_summary_path),
        "source_summary_sha256": sha256_file(source_summary_path),
        "class_names": source_summary.get("class_names", ["building"]),
        "tile_size": tile_size,
        "stride": tile_size,
        "source_stride": source_summary.get("stride"),
        "training_prohibited": True,
        "training_marker": str(no_training_path),
        "tile_count": len(output_tiles),
        "val_tile_count": len(output_tiles),
        "train_tile_count": 0,
        "positive_tile_count": sum(
            int(tile.get("label_count") or 0) > 0 for tile in output_tiles
        ),
        "negative_tile_count": sum(
            int(tile.get("label_count") or 0) == 0 for tile in output_tiles
        ),
        "label_count": sum(int(tile.get("label_count") or 0) for tile in output_tiles),
        "selected_offsets_by_sample": {
            sample: [list(offset) for offset in sorted(offsets)]
            for sample, offsets in sorted(offsets_by_sample.items())
        },
        "source_val_tile_counts_by_sample": dict(sorted(source_val_counts.items())),
        "excluded_tile_count": len(excluded_tiles),
        "excluded_tiles": excluded_tiles,
        "excluded_sample_slugs": sorted(
            {tile["sample_slug"] for tile in excluded_tiles}
        ),
        "exclude_low_variance": args.exclude_low_variance,
        "blank_range_threshold": args.blank_range_threshold,
        "tiles": output_tiles,
    }
    summary_path = output_dir / "yolo_tile_dataset_summary.json"
    summary_path.write_text(
        json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    manifest = {
        "schema_version": 1,
        "generated_at": datetime.now(UTC).isoformat(),
        "status": "complete_evaluation_only",
        "source_summary_path": str(source_summary_path),
        "source_summary_sha256": sha256_file(source_summary_path),
        "output_summary_path": str(summary_path),
        "output_summary_sha256": sha256_file(summary_path),
        "dataset_yaml_sha256": sha256_file(dataset_yaml),
        "evaluation_file_set_aggregate_sha256": aggregate_hash(
            output_files, output_dir
        ),
        "training_marker_sha256": sha256_file(no_training_path),
        "training_prohibited": True,
        "selected_tile_count": len(output_tiles),
        "selected_sample_count": len(offsets_by_sample),
        "excluded_tile_count": len(excluded_tiles),
        "excluded_tiles": excluded_tiles,
    }
    manifest_path = output_dir / "evaluation_view_manifest.json"
    manifest_path.write_text(
        json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    print(json.dumps(manifest, indent=2))
    return 0


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