"""Audit forest-canopy example tiles for missing files and inconsistent shapes."""

from __future__ import annotations

import argparse
from pathlib import Path
from typing import Iterable

import rasterio
import yaml


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Audit the forest-canopy example tiles referenced by a config file."
    )
    parser.add_argument(
        "--config",
        required=True,
        type=Path,
        help="Path to a forest-canopy YAML config file.",
    )
    return parser


def _load_config(config_path: Path) -> dict:
    with config_path.open("r", encoding="utf-8") as handle:
        return yaml.safe_load(handle)


def _resolve_years(file_path: str, year1: dict, year2: dict) -> tuple[int, int]:
    path_lower = file_path.lower()
    year2_lookup = {str(key).lower(): value for key, value in year2.items()}
    for site_key, year1_value in year1.items():
        normalized_key = str(site_key).lower()
        if normalized_key in path_lower:
            if normalized_key not in year2_lookup:
                raise ValueError(f"Missing year2 entry for site '{site_key}' in config.")
            return year1_value, year2_lookup[normalized_key]

    raise ValueError(f"Could not resolve year pair for tile path: {file_path}")


def _feature_file_endings(data_generator_config: dict) -> list[tuple[str, str]]:
    endings: list[tuple[str, str]] = []
    for branch in data_generator_config.get("branch_inputs", []):
        for feature_name, feature_info in branch.get("features", {}).items():
            endings.append((feature_name, feature_info["file_ending"]))
    return endings


def _find_matching_file(tile_path: Path, file_ending: str) -> Path | None:
    matches = sorted(path for path in tile_path.iterdir() if path.name.endswith(file_ending))
    if not matches:
        return None
    if len(matches) > 1:
        raise ValueError(
            f"Multiple files in {tile_path} match required suffix '{file_ending}': "
            f"{[match.name for match in matches]}"
        )
    return matches[0]


def _shape_for(path: Path) -> tuple[int, int]:
    with rasterio.open(path) as src:
        return src.height, src.width


def _expected_truth_file(tile_path: Path, relative_tile: str, data_generator_config: dict) -> Path:
    truth = data_generator_config.get("truth", "chm")
    if truth not in {"chc", "chc_maybe"}:
        match = _find_matching_file(tile_path, "_CHM.tif")
        if match is None:
            raise FileNotFoundError(f"Missing truth CHM file in {tile_path}")
        return match

    year1, year2 = _resolve_years(
        str(tile_path),
        data_generator_config.get("year1", {}),
        data_generator_config.get("year2", {}),
    )
    year2_relative = relative_tile.replace(str(year1), str(year2), 1)
    year2_tile_path = Path(data_generator_config["data_path"]) / year2_relative
    if not year2_tile_path.is_dir():
        raise FileNotFoundError(
            f"Expected comparison tile for truth '{truth}' does not exist: {year2_tile_path}"
        )
    match = _find_matching_file(year2_tile_path, "_CHM.tif")
    if match is None:
        raise FileNotFoundError(f"Missing comparison CHM file in {year2_tile_path}")
    return match


def audit_config(config: dict) -> tuple[list[str], list[str]]:
    data_generator_config = config.get("data_generator", {})
    data_root = Path(data_generator_config["data_path"])
    relative_tiles: Iterable[str] = data_generator_config.get("file_paths", [])
    required_features = _feature_file_endings(data_generator_config)

    errors: list[str] = []
    successes: list[str] = []

    for relative_tile in relative_tiles:
        tile_path = data_root / relative_tile
        if not tile_path.is_dir():
            errors.append(f"Missing tile directory: {tile_path}")
            continue

        tile_shapes: dict[str, tuple[int, int]] = {}
        missing_features: list[str] = []

        for feature_name, file_ending in required_features:
            match = _find_matching_file(tile_path, file_ending)
            if match is None:
                missing_features.append(f"{feature_name} ({file_ending})")
                continue
            tile_shapes[match.name] = _shape_for(match)

        if missing_features:
            errors.append(f"{relative_tile}: missing {', '.join(missing_features)}")
            continue

        try:
            truth_file = _expected_truth_file(tile_path, relative_tile, data_generator_config)
            tile_shapes[f"truth::{truth_file.name}"] = _shape_for(truth_file)
        except FileNotFoundError as exc:
            errors.append(str(exc))
            continue

        unique_shapes = sorted(set(tile_shapes.values()))
        if len(unique_shapes) != 1:
            formatted_shapes = ", ".join(f"{name}={shape}" for name, shape in tile_shapes.items())
            errors.append(f"{relative_tile}: inconsistent raster shapes: {formatted_shapes}")
            continue

        successes.append(f"{relative_tile}: ok shape={unique_shapes[0]}")

    return successes, errors


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)

    config = _load_config(args.config)
    successes, errors = audit_config(config)

    for line in successes:
        print(line)

    if errors:
        print("\nAudit failed:")
        for line in errors:
            print(f"- {line}")
        return 1

    print(f"\nAudit passed for {len(successes)} tiles.")
    return 0


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