import re
import warnings

import numpy as np
import xarray as xr

from access_moppy.base import CMORiser
from access_moppy.derivations import custom_functions, evaluate_expression
from access_moppy.utilities import (
    calculate_latitude_bounds,
    calculate_longitude_bounds,
    calculate_time_bounds,
)


class Atmosphere_CMORiser(CMORiser):
    """
    Handles CMORisation of NetCDF datasets for Atmosphere/Land variables across CMIP versions.
    """

    def calculate_missing_bounds_variables(self, bnds_required):
        """Calculate missing bounds variables for coordinates."""
        for bnds_var in bnds_required:
            # Extract coordinate name by removing "_bnds" suffix
            coord_name = bnds_var.replace("_bnds", "")
            if bnds_var not in self.ds.data_vars and bnds_var not in self.ds.coords:
                if coord_name not in self.ds.coords:
                    raise ValueError(
                        f"Cannot calculate {bnds_var}: coordinate '{coord_name}' not found in dataset"
                    )

                # Warn user that bounds are missing and will be calculated automatically
                warnings.warn(
                    f"'{bnds_var}' not found in raw data. Automatically calculating bounds for '{coord_name}' coordinate.",
                    UserWarning,
                    stacklevel=3,
                )

                # Determine which calculation function to use based on coordinate name
                if coord_name in ["time", "t"]:
                    # Calculate time bounds - atmosphere uses "bnds"
                    self.ds[bnds_var] = calculate_time_bounds(
                        self.ds,
                        time_coord=coord_name,
                        bnds_name="bnds",  # Atmosphere uses "bnds"
                    )

                elif coord_name in ["lat", "latitude", "y"]:
                    # Calculate latitude bounds - use "bnds" for atmosphere data
                    self.ds[bnds_var] = calculate_latitude_bounds(
                        self.ds, coord_name, bnds_name="bnds"
                    )

                elif coord_name in ["lon", "longitude", "x"]:
                    # Calculate longitude bounds - use "bnds" for atmosphere data
                    self.ds[bnds_var] = calculate_longitude_bounds(
                        self.ds, coord_name, bnds_name="bnds"
                    )

                else:
                    # For other coordinates, we could add more handlers or skip
                    warnings.warn(
                        f"No automatic calculation available for '{bnds_var}'. This may cause CMIP compliance issues.",
                        UserWarning,
                        stacklevel=3,
                    )
                    continue
            # Ensure the coordinate's bounds attribute always points to the bounds variable,
            # regardless of whether it was just calculated or already existed in the input data.
            if coord_name in self.ds.coords or coord_name in self.ds.data_vars:
                self.ds[coord_name].attrs["bounds"] = bnds_var

    def remove_spurious_time_dimensions(self, required_vars):
        """
        Remove spurious time dimensions from coordinate and auxiliary variables.

        This method addresses a common issue in xarray when combining datasets:
        spatial bounds (lat_bnds, lon_bnds) and other coordinate variables can incorrectly
        gain time dimensions during multi-file dataset operations, even though they are time-invariant.

        Why this is necessary:
        - When using xr.open_mfdataset() with combine_coords="time", xarray
          conservatively assumes all coordinate-linked variables might vary with time
        - This causes spatial bounds and coordinates to be broadcasted along the time dimension
        - Results in redundant data storage and non-CF-compliant files

        Why this is reasonable for ACCESS Models:
        - ACCESS Models use static grids throughout model runs
        - Latitude, longitude coordinates (and their bounds) are time-invariant
        - The grid definition remains constant across all timesteps
        - Only time_bnds and data variables should legitimately have a time dimension
        - This optimization is safe and improves storage efficiency

        Args:
            required_vars (list): Variables that should keep their time dimension
        """
        # Identify all variables that have gained spurious time dimensions
        # Include bounds variables and any other coordinate variables
        problematic_vars = [
            name
            for name in self.ds.variables
            if "time" not in name  # Don't touch time_bnds or time coordinate
            and name not in required_vars  # Don't touch required data variables
            and name in self.ds
            and "time" in self.ds[name].coords
            and self.ds[name].dims != ("time",)  # Skip pure time variables
        ]

        if problematic_vars:
            # Process all problematic variables efficiently in a single operation
            corrections = {
                name: self.ds[name].isel(time=0).drop_vars("time")
                for name in problematic_vars
            }
            self.ds = self.ds.assign(corrections)

    def select_and_process_variables(self):
        # Check if this is an internal calculation that doesn't need input variables
        calc = self.mapping[self.cmor_name]["calculation"]

        if calc["type"] == "internal":
            # For internal calculations, we don't need to load any input data
            # Call the internal calculation function directly
            func_name = calc["function"]
            if func_name not in custom_functions:
                raise ValueError(
                    f"Internal calculation function '{func_name}' not found in custom_functions"
                )

            # Execute the internal function to generate the variable data
            self.ds = custom_functions[func_name](**calc.get("kwargs", {}))

            self.vocab._get_axes(
                self.mapping
            )  # Ensure axes are loaded for renaming later

            # Ensure the CMOR variable exists
            if self.cmor_name not in self.ds:
                raise ValueError(
                    f"Internal calculation function '{func_name}' did not generate variable '{self.cmor_name}'"
                )

            return

        # Original logic for other calculation types
        # Select input variables required for the CMOR variable
        required_vars = self.mapping[self.cmor_name]["model_variables"]

        required_axes, axes_rename_map = self.vocab._get_axes(self.mapping)
        required_bounds, bounds_rename_map = self.vocab._get_required_bounds_variables(
            self.mapping
        )

        required = set(
            list(required_vars)
            + list(axes_rename_map.keys())
            + list(bounds_rename_map.keys())
        )
        self.load_dataset(required_vars=required)

        # Validate that all required model variables were actually loaded.
        # Without this, a missing variable is only caught at the rename/formula
        # step, producing a cryptic ValueError. Raise early with actionable context.
        missing_model_vars = [v for v in required_vars if v not in self.ds]
        if missing_model_vars:
            available = sorted(self.ds.data_vars)
            raise KeyError(
                f"Required model variable(s) {missing_model_vars} not found in the "
                f"input files for '{self.cmor_name}'. "
                f"Available data variables: {available}. "
                f"Check the 'model_variables' entry in the mapping."
            )

        # Remove spurious time dimensions from spatial bounds and coordinates
        self.remove_spurious_time_dimensions(required_vars)

        # Ensure time dimension is sorted
        self.sort_time_dimension()

        # Handle the calculation type
        if calc["type"] == "direct":
            # If the calculation is direct, just rename the variable
            self.ds = self.ds.rename({required_vars[0]: self.cmor_name})
        elif calc["type"] == "formula":
            context = {var: self.ds[var] for var in required_vars}
            context.update(custom_functions)

            # Save original time attrs before formula (decode_cf moves them to encoding)
            orig_time_attrs = self.ds["time"].attrs.copy() if "time" in self.ds else {}
            result = evaluate_expression(calc, context)

            # Check whether the time interval/frequency has changed (e.g. daily → monthly)
            if "time" in result.dims and result.sizes["time"] != self.ds.sizes.get(
                "time", result.sizes["time"]
            ):
                # If the temporal resolution changes, rebuild self.ds while preserving variables that are not time-dependent
                time_indep = {
                    v: self.ds[v]
                    for v in self.ds.data_vars
                    if "time" not in self.ds[v].dims and v not in required_vars
                }
                time_indep_coords = {
                    c: self.ds[c]
                    for c in self.ds.coords
                    if "time" not in self.ds[c].dims and c != "time"
                }
                self.ds = result.to_dataset(name=self.cmor_name)
                for v, da in time_indep.items():
                    self.ds[v] = da
                self.ds = self.ds.assign_coords(
                    {
                        c: v
                        for c, v in time_indep_coords.items()
                        if c not in self.ds.coords
                    }
                )
                # Restore original time attrs so generate_filename can read units/calendar
                if orig_time_attrs and "time" in self.ds:
                    self.ds["time"].attrs.update(orig_time_attrs)
                self.calculate_missing_bounds_variables(required_bounds)
            else:
                # If the temporal resolution remains unchanged, assign directly
                self.ds[self.cmor_name] = result

            # Drop unit after calculation. update_attributes() will add the right units later on.
            self.ds[self.cmor_name].attrs.pop("units", None)
            # Drop the original input variables, except the CMOR variable and keep bounds
            self.ds = self.ds.drop_vars(
                [
                    var
                    for var in required_vars
                    if var != self.cmor_name and var not in required_bounds.keys()
                ],
                errors="ignore",
            )

        elif calc["type"] == "dataset_function":
            # Function that operates on the full dataset
            func_name = calc["function"]
            self.ds = self.ds.rename({required_vars[0]: self.cmor_name})
            self.ds = custom_functions[func_name](self.ds, **calc.get("kwargs", {}))
        else:
            raise ValueError(f"Unsupported calculation type: {calc['type']}")

        # Rename axes and bounds variables
        rename_map = {
            k: v
            for k, v in {**bounds_rename_map, **axes_rename_map}.items()
            if k in self.ds
        }

        # Drop any existing variables that have the same names as our target names
        conflicting_vars = [
            v
            for v in rename_map.values()
            if v in self.ds and v not in rename_map.keys()
        ]
        if conflicting_vars:
            self.ds = self.ds.drop_vars(conflicting_vars, errors="ignore")

        self.ds = self.ds.rename(rename_map)
        # Drop stale units from renamed coordinates; update_attributes will
        # assign the correct CMIP units from the vocabulary.
        for old_name, new_name in rename_map.items():
            if old_name != new_name and new_name in self.ds.coords:
                self.ds[new_name].attrs.pop("units", None)

        # Calculate missing bounds variables after renaming so that
        # coordinate names in self.ds match the output names in required_bounds
        self.calculate_missing_bounds_variables(required_bounds)

        # Transpose the data variable according to the CMOR dimensions
        # Handle both string and list dimension formats
        dimensions = self.vocab.variable["dimensions"]
        try:
            # Try treating as string (space-separated)
            cmor_dims = re.sub(r"\w*level", "lev", dimensions).split()
        except TypeError:
            # If re.sub() fails (TypeError for list input), it's already a list
            cmor_dims = [re.sub(r"\w*level", "lev", dim) for dim in dimensions]

        transpose_order = [
            self.vocab.axes[dim]["out_name"]
            for dim in cmor_dims
            if "value" not in self.vocab.axes[dim]
        ]

        # Squeeze singleton time dimensions not needed in output
        for dim in ("time_0", "time_1"):
            if dim in self.ds[self.cmor_name].dims and dim not in transpose_order:
                self.ds = self.ds.isel({dim: 0}, drop=True)
        # Squeeze singleton dimensions if they are not in the transpose order
        for dim in self.ds[self.cmor_name].dims:
            if dim not in transpose_order and self.ds[self.cmor_name][dim].size == 1:
                self.ds[self.cmor_name] = self.ds[self.cmor_name].squeeze(dim)

        # Enforce dimension order: time first, lat/lon last (lat before lon),
        # with any remaining dimensions (e.g. lev) in between.
        time_dims = [dim for dim in transpose_order if dim == "time"]
        middle_dims = [
            dim for dim in transpose_order if dim not in ("time", "lat", "lon")
        ]
        lat_lon_dims = [dim for dim in ("lat", "lon") if dim in transpose_order]
        transpose_order = time_dims + middle_dims + lat_lon_dims

        self.ds[self.cmor_name] = self.ds[self.cmor_name].transpose(*transpose_order)

    def update_attributes(self):
        self.ds.attrs = {
            k: v
            for k, v in self.vocab.get_required_global_attributes().items()
            if v not in (None, "")
        }

        required_coords = {
            v["out_name"] for v in self.vocab.axes.values() if "value" in v
        }.union({v["out_name"] for v in self.vocab.axes.values()})
        self.ds = self.ds.drop_vars(
            [c for c in self.ds.coords if c not in required_coords], errors="ignore"
        )

        cmor_attrs = self.vocab.variable
        self._check_units(self.cmor_name, cmor_attrs.get("units"))

        self.ds[self.cmor_name].attrs.update(
            {k: v for k, v in cmor_attrs.items() if v not in (None, "")}
        )
        var_type = cmor_attrs.get("type", "double")
        self.ds[self.cmor_name] = self.ds[self.cmor_name].astype(
            self.type_mapping.get(var_type, np.float64)
        )

        try:
            if cmor_attrs.get("valid_min") not in (None, "") and cmor_attrs.get(
                "valid_max"
            ) not in (None, ""):
                vmin = self.type_mapping.get(var_type, np.float64)(
                    cmor_attrs["valid_min"]
                )
                vmax = self.type_mapping.get(var_type, np.float64)(
                    cmor_attrs["valid_max"]
                )
                self._check_range(self.cmor_name, vmin, vmax)
        except ValueError as e:
            raise ValueError(
                f"Failed to validate value range for {self.cmor_name}: {e}"
            )

        for dim, meta in self.vocab.axes.items():
            name = meta["out_name"]
            dtype = self.type_mapping.get(meta.get("type", "double"), np.float64)
            if name in self.ds:
                if meta.get("standard_name") == "time":
                    self._check_calendar(name)
                original_units = self.ds[name].attrs.get("units", "")
                coord_attrs = {
                    k: v
                    for k, v in {
                        "standard_name": meta.get("standard_name"),
                        "long_name": meta.get("long_name"),
                        "units": meta.get("units"),
                        "axis": meta.get("axis"),
                        "positive": meta.get("positive"),
                        "valid_min": dtype(meta["valid_min"])
                        if "valid_min" in meta
                        else None,
                        "valid_max": dtype(meta["valid_max"])
                        if "valid_max" in meta
                        else None,
                    }.items()
                    if v is not None
                }
                if coord_attrs.get(
                    "units"
                ) == "days since ?" and original_units.lower().startswith("days since"):
                    coord_attrs["units"] = original_units
                # Skip astype for time coordinates containing datetime/cftime objects,
                # and for character-type coordinates (string arrays like vegtype)
                if (
                    meta.get("standard_name") == "time"
                    and (
                        self.ds[name].dtype == object
                        or np.issubdtype(self.ds[name].dtype, np.datetime64)
                    )
                ) or meta.get("type") == "character":
                    updated = self.ds[name].copy()
                else:
                    updated = self.ds[name].astype(dtype)
                updated.attrs.update(coord_attrs)
                updated.attrs.pop("_FillValue", None)
                self.ds[name] = updated
            elif "value" in meta:
                val = meta["value"]
                # Handle character type (e.g., string coordinate)
                if meta["type"] == "character":
                    arr = xr.DataArray(
                        np.array(
                            val, dtype="S"
                        ),  # ensure type is character (byte string)
                        dims=(),
                        attrs={
                            k: v
                            for k, v in {
                                "standard_name": meta.get("standard_name"),
                                "long_name": meta.get("long_name"),
                                "units": meta.get("units"),
                                "axis": meta.get("axis"),
                                "positive": meta.get("positive"),
                                "valid_min": meta.get("valid_min"),
                                "valid_max": meta.get("valid_max"),
                            }.items()
                            if v is not None
                        },
                    )
                else:
                    arr = xr.DataArray(
                        dtype(val),
                        dims=(),
                        attrs={
                            k: v
                            for k, v in {
                                "standard_name": meta.get("standard_name"),
                                "long_name": meta.get("long_name"),
                                "units": meta.get("units"),
                                "axis": meta.get("axis"),
                                "positive": meta.get("positive"),
                                "valid_min": dtype(meta["valid_min"])
                                if "valid_min" in meta
                                else None,
                                "valid_max": dtype(meta["valid_max"])
                                if "valid_max" in meta
                                else None,
                            }.items()
                            if v is not None
                        },
                    )
                self.ds = self.ds.assign_coords({name: arr})

        # CF §7.1 — bounds variables must share attrs with their parent coordinate.
        # They should NOT carry _FillValue, coordinates, or a bounds pointer.
        # After renaming (e.g. sigma_theta_bnds → b_bnds) the inherited source attrs
        # are stale; replace them with attrs derived from the (now-updated) parent.
        INHERIT_KEYS = (
            "standard_name",
            "long_name",
            "units",
            "axis",
            "positive",
            "calendar",
        )
        all_ds_vars = list(self.ds.data_vars) + list(self.ds.coords)
        for var in all_ds_vars:
            if not var.endswith("_bnds"):
                continue
            parent = var[: -len("_bnds")]
            parent_attrs = {}
            if parent in self.ds:
                parent_attrs = {
                    k: v for k, v in self.ds[parent].attrs.items() if k in INHERIT_KEYS
                }
            self.ds[var].attrs = parent_attrs
