# This file is part of the Open Data Cube, see https://opendatacube.org for more information
#
# Copyright (c) 2015-2025 ODC Contributors
# SPDX-License-Identifier: Apache-2.0
import contextlib
import logging
from collections.abc import Mapping
from pathlib import Path
from typing import Any

import netCDF4
import xarray
from odc.geo import CRS
from xarray.core.coordinates import DatasetCoordinates

from datacube.storage._hdf5 import HDF5_LOCK
from datacube.utils import DatacubeException

from . import writer as netcdf_writer

_LOG: logging.Logger = logging.getLogger(__name__)


def _get_units(coord) -> str:
    """
    Guess units from coordinate
    1. Value of .units if set
    2. 'seconds since 1970-01-01 00:00:00' if coord dtype is datetime
    3. '1' otherwise
    """
    units = getattr(coord, "units", None)
    if units is not None:
        return units
    dtype = getattr(coord.values, "dtype", None)
    if dtype is None:
        return "1"

    if dtype.kind == "M":
        return "seconds since 1970-01-01 00:00:00"

    return "1"


def create_netcdf_storage_unit(
    filename: Path,
    crs: CRS,
    coordinates: DatasetCoordinates,
    variables: Mapping[str, Any],
    variable_params: Mapping[str, Mapping[str, Any]],
    global_attributes: dict | None = None,
    netcdfparams: dict | None = None,
) -> netCDF4.Dataset:
    """
    Create a NetCDF file on disk.

    :param filename: filename to write to
    :param crs: Datacube CRS object defining the spatial projection
    :param coordinates: Dict of named `datacube.model.Coordinate`s to create
    :param variables: Dict of named `datacube.model.Variable`s to create
    :param variable_params:
        Dict of dicts, with keys matching variable names, of extra parameters for variables
    :param global_attributes: named global attributes to add to output file
    :param netcdfparams: Extra parameters to use when creating netcdf file
    :return: open netCDF4.Dataset object, ready for writing to
    """
    filename = Path(filename)
    if filename.exists():
        raise RuntimeError(f"Storage Unit already exists: {filename}")

    with contextlib.suppress(OSError):
        filename.parent.mkdir(parents=True)

    _LOG.info("Creating storage unit: %s", filename)

    nco = netcdf_writer.create_netcdf(str(filename), **(netcdfparams or {}))

    for name, coord in coordinates.items():
        if coord.values.ndim > 0:  # skip CRS coordinate
            netcdf_writer.create_coordinate(
                nco, str(name), coord.values, _get_units(coord)
            )

    grid_mapping = netcdf_writer.DEFAULT_GRID_MAPPING
    netcdf_writer.create_grid_mapping_variable(nco, crs, name=grid_mapping)

    for name, variable in variables.items():
        has_crs = all(dim in variable.dims for dim in crs.dimensions)
        var_params = variable_params.get(name, {})
        data_var = netcdf_writer.create_variable(
            nco,
            name,
            variable,
            grid_mapping=grid_mapping if has_crs else None,
            **var_params,
        )

        for key, value in var_params.get("attrs", {}).items():
            setattr(data_var, key, value)

    for key, value in (global_attributes or {}).items():
        setattr(nco, key, value)

    return nco


def write_dataset_to_netcdf(
    dataset: xarray.Dataset,
    filename: Path,
    global_attributes: dict | None = None,
    variable_params: dict | None = None,
    netcdfparams=None,
) -> None:
    """
    Write a Data Cube style xarray Dataset to a NetCDF file

    Requires a spatial Dataset, with attached coordinates and global crs attribute.

    :param dataset:
    :param filename: Output filename
    :param global_attributes: Global file attributes. dict of attr_name: attr_value
    :param variable_params: dict of variable_name: {param_name: param_value, [...]}
                            Allows setting storage and compression options per variable.
                            See the `netCDF4.Dataset.createVariable` for available
                            parameters.
    :param netcdfparams: Optional params affecting netCDF file creation
    """
    global_attributes = global_attributes or {}
    variable_params = variable_params or {}
    filename = Path(filename)

    if not dataset.data_vars.keys():
        raise DatacubeException("Cannot save empty dataset to disk.")

    if dataset.odc.geobox is None:
        raise DatacubeException(
            "Dataset geobox property is None, cannot write to NetCDF file."
        )

    try:
        HDF5_LOCK.acquire(blocking=True)
        nco = create_netcdf_storage_unit(
            filename,
            dataset.odc.geobox.crs,
            dataset.coords,
            dataset.data_vars,
            variable_params,
            global_attributes,
            netcdfparams,
        )

        for name, variable in dataset.data_vars.items():
            nco[name][:] = netcdf_writer.netcdfy_data(variable.values)

        nco.close()
    finally:
        HDF5_LOCK.release()
