"""Validate NetCDF variables, dimensions, and statistics for FV3GFS inputs."""

import argparse
from pathlib import Path

import xarray as xr


REQUIRED_VARIABLES = {
    "PRESsfc",
    "surface_temperature",
    "DSWRFtoa",
    "HGTsfc",
    "ocean_fraction",
    *{f"air_temperature_{i}" for i in range(8)},
    *{f"specific_total_water_{i}" for i in range(8)},
    *{f"eastward_wind_{i}" for i in range(8)},
    *{f"northward_wind_{i}" for i in range(8)},
}


def validate_data(data_path: str, latitude: int, longitude: int) -> int:
    files = sorted(Path(data_path).rglob("*.nc"))
    if not files:
        print(f"ERROR: no NetCDF files found under {data_path}")
        return 1
    failed = False
    for path in files:
        with xr.open_dataset(path) as dataset:
            variables = set(dataset.data_vars)
            missing = sorted(REQUIRED_VARIABLES - variables)
            dimensions = {name: int(size) for name, size in dataset.sizes.items()}
            print(f"{path}: dimensions={dimensions}, variables={len(variables)}")
            if missing:
                print(f"ERROR: missing required variables: {', '.join(missing)}")
                failed = True
            if latitude not in dimensions.values() or longitude not in dimensions.values():
                print(f"ERROR: expected spatial dimensions containing {latitude} and {longitude}")
                failed = True
            time_sizes = [dimensions[name] for name in dimensions if "time" in name.lower()]
            if not time_sizes or max(time_sizes) < 7:
                print("ERROR: expected a time dimension with at least 7 frames")
                failed = True
    return int(failed)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--data-dir", required=True)
    parser.add_argument("--latitude", type=int, default=180)
    parser.add_argument("--longitude", type=int, default=360)
    args = parser.parse_args()
    raise SystemExit(validate_data(args.data_dir, args.latitude, args.longitude))


if __name__ == "__main__":
    main()