"""Read and write data from/to the ANU CTLab netcdf data format.

This is an optional extra module, and must be explicitly installed to be used (e.g., ``pip install anu_ctlab_io[netcdf]``)."""

import importlib.util
import os
import re
from pathlib import Path
from typing import Any

import xarray as xr

from anu_ctlab_io._dataset import Dataset
from anu_ctlab_io._datatype import DataType
from anu_ctlab_io._parse_history import parse_history
from anu_ctlab_io._voxel_properties import VoxelUnit

if importlib.util.find_spec("h5netcdf") is None:
    raise ImportError(
        "h5netcdf is required. Install it with: pip install anu_ctlab_io[netcdf]"
    )

from anu_ctlab_io.netcdf._writer import dataset_to_netcdf

__all__ = [
    "dataset_from_netcdf",
    "dataset_to_netcdf",
]


def dataset_from_netcdf(
    path: Path, *, parse_history: bool = True, **kwargs: Any
) -> Dataset:
    """Loads a :any:`Dataset` from the path to a netcdf.

    This method is used by :any:`Dataset.from_path`, by preference call that constructor directly.

    :param Path: The path to the netcdf or directory of split netcdf blocks to be loaded.
    :param parse_history: Whether to parse the history of the netcdf file. Defaults to ``True``, but disableable because the parser is currently not guaranteed to succeed.
    :param kwargs: Currently this method consumes no kwargs, but will pass provided kwargs to ``Xarray.open_mfdataset``.
    :raises lark.exceptions.UnexpectedInput: Raised if ``parse_history=True`` and the parser fails to parse the specific history provided."""
    datatype = DataType.infer_from_path(path)
    dataset = _read_netcdf(path, datatype, **kwargs)
    dataset = dataset.rename(_transform_data_vars(dataset, datatype))
    dataset["data"] = dataset.data.astype(datatype.dtype)
    dataset.attrs = _update_attrs(dataset.attrs, parse_history)
    dataset_id = dataset.attrs.get("dataset_id")
    return Dataset(
        data=dataset.data.data,
        dimension_names=tuple(
            map(str, dataset.data.dims)
        ),  # Use dims from data var only
        datatype=datatype,
        voxel_unit=VoxelUnit.from_str(dataset.attrs["voxel_unit"]),
        voxel_size=dataset.attrs["voxel_size"],
        history=dataset.history,
        dataset_id=dataset_id,
    )


def _transform_data_vars(dataset: xr.Dataset, datatype: DataType) -> dict[str, str]:
    attr_transform = {f"{datatype}_{dim}dim": dim for dim in ["x", "y", "z"]}
    for k in dataset.data_vars.keys():
        match k:
            case a if isinstance(a, str) and a.find(str(datatype)) == 0:
                attr_transform[k] = "data"
    return attr_transform


def _update_attrs(attrs: dict[str, Any], parse_history_p: bool) -> dict[str, Any]:
    new_attrs: dict[str, Any] = {"history": {}}
    for k, v in attrs.items():
        match k:
            case a if a.find("history") == 0:
                if parse_history_p:
                    new_attrs["history"][k[len("history") + 1 :]] = parse_history(v)
                else:
                    new_attrs["history"][k[len("history") + 1 :]] = v
            case a if a.find("dim") != -1:
                new_attrs[re.sub("([x|y|z])dim", "\\1", k)] = v
            case a if a in ["number_of_files", "zdim_total", "total_grid_size_xyz"]:
                pass
            case a if a.find("_xyz"):
                new_attrs[re.sub("(.*)_xyz", "\\1", k)] = v
            case _:
                new_attrs[k] = attrs[k]
    return new_attrs


def _read_netcdf(path: Path | str, datatype: DataType, **kwargs: Any) -> xr.Dataset:
    path = os.path.normpath(os.path.expanduser(path))
    if os.path.isdir(path):
        possible_files = [os.path.join(path, p) for p in os.listdir(path)]
        files = sorted(list(filter(os.path.isfile, possible_files)))
        last_exc: OSError | None = None
        for engine in _read_engines():
            try:
                return xr.open_mfdataset(
                    files,
                    engine=engine,
                    combine="nested",
                    concat_dim=[f"{datatype}_zdim"],
                    combine_attrs="drop_conflicts",
                    coords="minimal",
                    compat="override",
                    mask_and_scale=False,
                    data_vars=[f"{datatype}"],
                    **kwargs,
                )
            except OSError as e:
                last_exc = e
        raise OSError(
            f"Could not read netCDF files in {path} with any available engine."
        ) from last_exc
    else:
        chunks = kwargs.pop("chunks", -1)
        last_exc = None
        for engine in _read_engines():
            try:
                return xr.open_dataset(
                    path,
                    engine=engine,
                    mask_and_scale=False,
                    chunks=chunks,
                    **kwargs,
                )
            except OSError as e:
                last_exc = e
        raise OSError(
            f"Could not read netCDF file {path} with any available engine."
        ) from last_exc


# Maybe netCDF3 support becomes optional at some point?
def _read_engines() -> list[str]:
    engines = ["h5netcdf"]
    if importlib.util.find_spec("netCDF4") is not None:
        engines.append("netCDF4")
    return engines
