from .cli_logger import get_cli_logger
from nes.load_nes import open_netcdf
from numpy import isinf, isnan
from typing import Optional, List


logger = get_cli_logger("checker")


def run_checks(
    input_file: str,
    check_nan: Optional[bool] = True,
    check_inf: Optional[bool] = True,
    var_list: Optional[List[str]] = None,
) -> bool:
    """
    Check for NaN and/or Inf values in selected variables of a NetCDF file.

    This function uses the NES `open_netcdf()` interface to load a NetCDF dataset,
    optionally restricts the check to a subset of variables (`var_list`),
    reads all selected variables into memory, and checks for the presence of NaN or Inf values.

    Parameters
    ----------
    input_file : str
        Path to the input NetCDF file.
    check_nan : Optional[bool]
        If True, check for NaN (Not a Number) values in the dataset. Default is True.
    check_inf : Optional[bool]
        If True, check for Inf (infinite) values in the dataset. Default is True.
    var_list : Optional[List[str]]
        List of variable names to check. If None, all variables will be checked.

    Returns
    -------
    bool
        True if no NaN or Inf values are found in any selected variable.

    Raises
    ------
    ValueError
        If any selected variable contains NaN or Inf values, a ValueError is raised
        indicating the variable name.

    Notes
    -----
    This function loads selected variable data into memory. Use with caution on large datasets.
    """
    # Open the NetCDF file with NES (metadata only)
    logger.info(f"Opening NetCDF file: {input_file}")
    dataset = open_netcdf(input_file)

    if var_list is not None:
        dataset.keep_vars(var_list)

    logger.info("Loading variable data into memory...")
    dataset.load()

    # Loop over all variables in the dataset
    for var_name, var_info in dataset.variables.items():
        logger.debug(f"Checking variable: {var_name}")
        # Check for Inf values, if requested
        has_inf = isinf(var_info["data"]).any() if check_inf else False

        # Check for NaN values, if requested
        has_nan = isnan(var_info["data"]).any() if check_nan else False

        # Raise an error if problematic values are found
        if has_nan or has_inf:
            logger.warning(f"Variable '{var_name}' contains NaN or Inf values.")
            raise ValueError(f"Variable '{var_name}' contains NaN or Inf values.")

    logger.info("All variables passed NaN/Inf checks.")
    return True