from __future__ import annotations

import argparse
from datetime import UTC, datetime
from pathlib import Path
from typing import Iterable

import xarray as xr
from xarray.coders import CFDatetimeCoder


EXCLUDE_DIRS = {"_figures", "_notebooks", "compressed"}
START_DATE = "2000-01-01"
DEFAULT_COMPLEVEL = 1
NETCDF_ENGINE = "netcdf4"
TIME_CODER = CFDatetimeCoder(use_cftime=True)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Rewrite CMIP6 monthly time coordinates to month starts, keep only "
            "data from 2000 onward, and write compressed copies under ./compressed."
        )
    )
    parser.add_argument(
        "--root",
        type=Path,
        default=Path(__file__).resolve().parent,
        help="Workspace root (where compressed/ will be created).",
    )
    parser.add_argument(
        "--input-root",
        type=Path,
        default=None,
        help="Directory containing the model folders. Defaults to <root>/raw.",
    )
    parser.add_argument(
        "--output-root",
        type=Path,
        default=None,
        help="Destination root. Defaults to <root>/compressed.",
    )
    parser.add_argument(
        "--complevel",
        type=int,
        default=DEFAULT_COMPLEVEL,
        choices=range(0, 10),
        metavar="0-9",
        help="zlib compression level. 1 is lossless and usually the best speed/size tradeoff here.",
    )
    parser.add_argument(
        "--overwrite",
        action="store_true",
        help="Overwrite outputs that already exist.",
    )
    parser.add_argument(
        "--limit",
        type=int,
        default=None,
        help="Process at most this many files. Useful for validation runs.",
    )
    return parser.parse_args()


def iter_input_files(input_root: Path, output_root: Path) -> Iterable[Path]:
    for model_dir in sorted(input_root.iterdir()):
        if not model_dir.is_dir():
            continue
        if model_dir.name in EXCLUDE_DIRS:
            continue
        if model_dir.resolve() == output_root.resolve():
            continue
        yield from sorted(model_dir.glob("*_era5corrected.nc"))


def to_month_start(value):
    if hasattr(value, "calendar"):
        return type(value)(value.year, value.month, 1)
    return datetime(value.year, value.month, 1)


def rewrite_time(ds: xr.Dataset) -> xr.Dataset:
    new_time = [to_month_start(value) for value in ds.time.values]
    return ds.assign_coords(time=new_time)


def build_encoding(ds: xr.Dataset, complevel: int, time_encoding: dict) -> dict:
    encoding: dict[str, dict] = {
        "time": {
            "units": time_encoding.get("units", "days since 1850-01-01"),
            "calendar": time_encoding.get("calendar", ds.time.attrs.get("calendar", "standard")),
            "dtype": time_encoding.get("dtype", "float64"),
        }
    }

    for name, variable in ds.data_vars.items():
        var_encoding = {
            "zlib": True,
            "complevel": complevel,
            "shuffle": True,
        }
        if variable.dtype.kind != "f":
            var_encoding["dtype"] = variable.dtype

        source_chunks = variable.encoding.get("chunksizes")
        if source_chunks:
            var_encoding["chunksizes"] = source_chunks

        fill_value = variable.encoding.get("_FillValue")
        if fill_value is not None:
            var_encoding["_FillValue"] = fill_value

        encoding[name] = var_encoding

    return encoding


def process_file(src_path: Path, input_root: Path, output_root: Path, complevel: int, overwrite: bool) -> None:
    rel_path = src_path.relative_to(input_root)
    dst_path = output_root / rel_path
    dst_path.parent.mkdir(parents=True, exist_ok=True)

    if dst_path.exists() and not overwrite:
        print(f"[skip] {dst_path} already exists")
        return
    if dst_path.exists() and overwrite:
        print(f"[overwrite] removing existing {dst_path}")
        dst_path.unlink()

    print(f"[open] {src_path}")
    with xr.open_dataset(src_path, decode_times=False, engine=NETCDF_ENGINE) as ds_raw:
        ds = xr.decode_cf(ds_raw, decode_times=TIME_CODER)
        if "time" not in ds.coords:
            raise ValueError(f"No time coordinate found in {src_path}")

        original_count = ds.sizes.get("time", 0)
        time_encoding = dict(ds.time.encoding)

        ds_out = ds.sel(time=slice(START_DATE, None))
        kept_count = ds_out.sizes.get("time", 0)
        if kept_count == 0:
            print(f"[skip] {src_path} has no data on or after {START_DATE}")
            return

        ds_out = rewrite_time(ds_out)
        ds_out.attrs = dict(ds.attrs)
        history_line = (
            f"{datetime.now(UTC).strftime('%Y-%m-%d %H:%M:%S')} UTC: "
            "rewrote monthly timestamps to first-of-month, filtered to >= 2000-01-01, "
            "and saved with lossless zlib compression"
        )
        previous_history = ds_out.attrs.get("history")
        ds_out.attrs["history"] = (
            f"{previous_history}\n{history_line}" if previous_history else history_line
        )

        encoding = build_encoding(ds_out, complevel=complevel, time_encoding=time_encoding)

        first_time = ds_out.time.values[0]
        last_time = ds_out.time.values[-1]
        print(
            f"[write] {dst_path} | time: {original_count} -> {kept_count} | "
            f"range: {first_time} .. {last_time}"
        )
        ds_out.to_netcdf(dst_path, engine=NETCDF_ENGINE, encoding=encoding)


def main() -> None:
    args = parse_args()
    root = args.root.resolve()
    input_root = (args.input_root or (root / "raw")).resolve()
    output_root = (args.output_root or (root / "compressed")).resolve()
    output_root.mkdir(parents=True, exist_ok=True)

    files = list(iter_input_files(input_root, output_root))
    if args.limit is not None:
        files = files[: args.limit]

    print(f"Found {len(files)} input files under {input_root}")
    print(f"Writing outputs under {output_root}")

    for index, src_path in enumerate(files, start=1):
        print(f"[{index}/{len(files)}] processing {src_path.relative_to(input_root)}")
        process_file(
            src_path=src_path,
            input_root=input_root,
            output_root=output_root,
            complevel=args.complevel,
            overwrite=args.overwrite,
        )

    print("Finished processing all files")


if __name__ == "__main__":
    main()