import copy
import os
from glob import glob
from pathlib import Path as pathlib_Path

from nabu.pipeline.config import (
    overwrite_config,
    validate_config,
)
from nabu.pipeline.config import (
    parse_nabu_config_file as parse_config_file,
)
from nabu.resources.utils import extract_parameters
from silx.io.url import DataUrl
from tomoscan.esrf.scan.fscan import FscanDataset, list_datasets

from ..hdf5 import is_same_config
from ..integrator import AIConfiguration
from ..multiproc_integration import ComputationResources
from ..utils import file_mode_to_octal, get_additional_metadata, safe_format
from .integrator_config import integrator_config, renamed_keys

path = os.path


def get_computations_config(conf):
    """
    Converts a dictionary to a ClusterConfiguration object
    """
    if conf.get("n_workers", None) is None:
        if conf.get("computations distribution", None) is None:
            raise ValueError(
                "Passed parameter does not look like a valid configuration dictionary. Please use a dictionary built from the configuration file.",
            )
        else:
            conf = conf["computations distribution"]

    # Get python executable to use on workers
    computation_resources = ComputationResources(
        n_multi_worker=conf["n_workers"],
        n_workers_per_multiworker=conf["ai_engines_per_worker"],
        n_threads_per_worker=conf["cores_per_worker"],
        memory=conf["memory_per_worker"],
        walltime=conf["time"],
        queue=conf["partition"],
    )
    return computation_resources


def get_ai_config(conf):
    """
    Converts a dictionary to a AIConfiguration object
    """
    ai_conf_dict = conf["azimuthal integration"]
    ai_config = AIConfiguration(
        n_pts=ai_conf_dict["n_pts"],
        unit=ai_conf_dict["unit"],
        polarization_factor=ai_conf_dict["polarization_factor"],
        correct_solid_angle=bool(ai_conf_dict["correct_solid_angle"]),
        detector=ai_conf_dict["detector"],
        poni_file=ai_conf_dict["poni_file"],
        mask_file=ai_conf_dict["mask_file"],
        flatfield_file=ai_conf_dict["flatfield_file"],
        dark_file=ai_conf_dict["dark_file"],
        error_model=ai_conf_dict["error_model"],
        azimuthal_range=ai_conf_dict["azimuthal_range"],
        radial_range=ai_conf_dict["radial_range"],
        ai_method=ai_conf_dict["ai_method"],
        trim_method=ai_conf_dict["trim_method"],
        trim_n_pts=ai_conf_dict["trim_n_pts"],
        trim_bounds=ai_conf_dict["trim_bounds"],
        pixel_splitting=ai_conf_dict["pixel_splitting"],
        average_xy=ai_conf_dict["average_xy"],
        correct_detector_attenuation=ai_conf_dict["correct_detector_attenuation"],
    )
    if ai_config.azimuthal_range[-1] > 1 and ai_config.trim_method is not None:
        raise ValueError(
            "Multi-slice integration cannot be used with a trim method. Please either disable 'trim_method' or set the last number of 'azimuthal_range' to 1.",
        )
        # if ai_conf_dict["average_xy"] is not None:
        # Current limitation
        # raise ValueError("Cannot use multiple azimuthal slices with average_xy=True")
    if ai_config.trim_method == "sigma_clip" and ai_config.pixel_splitting != "no":
        print(
            "Combining sigma-clip and pixel splitting is not recommended. Disabling pixel splitting !",
        )
        ai_config.pixel_splitting = "no"

    return ai_config


def browse_datasets_wildcard(
    wildcards,
    entry=None,
    detector_name=None,
    detector_suffix=None,
    errors="print",
    extra_metadata_fields=None,
    logger=None,
):
    """
    Browse folders and look for datasets.
    This assumes that the datasets are stored in this layout:
    /path/to/samplename_num/samplename_num.h5

    Parameters
    ----------
    wildcards: str or list of str
        Wildcard or list of wildcards,
    entry: str, optional
        Specific HDF5 entry to open (e.g. "3.1"); if None, scans all X.1 entries
    detector_name: str
        Detector to use
    detector_suffix: str, optional
        Suffix to use after detector name that can be found in the measurements section of the HDF5 file (Ex: raw_frame)
    errors: {"print","raise","ignore"}
        How to handle errors
    extra_metadata_fields: list[str], optional
        Which extra metadata axes to pull via `get_additional_metadata`
    logger : Logger or None
        Logger instance or None; will default to printing.
    """

    def handle_error(msg):
        if errors == "print":
            print(msg)
        elif errors == "raise":
            raise ValueError(msg)
        else:
            pass

    if isinstance(wildcards, str):
        wildcards = [wildcards]
    datasets = []
    for wc in wildcards:
        for fname in sorted(glob(wc.strip())):
            if not path.isfile(fname):
                handle_error(f"No such file: {fname}")
                continue

            try:
                file_entries = list_datasets(fname)
            except RuntimeError as exc:
                handle_error(f"Error with file {fname}: {exc} - ignoring")
                continue

            entries = [entry] if entry is not None else file_entries

            for d_entry in entries:
                try:
                    # 1) build the “master” FscanDataset
                    master = FscanDataset(
                        fname,
                        entry=d_entry,
                        detector_name=detector_name,
                        detector_suffix=detector_suffix,
                        on_missing_metadata="ignore",
                    )

                    # 2) pull in any extra metadata
                    scan_num = master.entry.split(".")[0]
                    extra = (
                        get_additional_metadata(
                            scan_num,
                            master.fname,
                            extra_metadata_fields,
                            logger=logger,
                        )
                        or {}
                    )
                    master.metadata.update(extra)

                    # 3) split into one sub_dataset per VDS file
                    vs_map = master.get_virtual_sources(remove_nonexisting=True)
                    for sub_fname, sub_path in vs_map.items():
                        sub = copy.copy(master)
                        sub.fname = sub_fname
                        sub.data_path = sub_path
                        sub.dataset_hdf5_url = DataUrl(
                            file_path=sub_fname,
                            data_path=sub_path,
                            scheme="silx",
                        )
                        # clear cache so downstream calls to get_virtual_sources() will recompute if needed
                        sub._virtual_sources = None
                        datasets.append(sub)

                except (ValueError, NameError, KeyError, IndexError, OSError) as exc:
                    handle_error(f"Error with file {fname}.\n{exc}\nIgnoring file.")

    return datasets


def get_output_folders(
    datasets,
    output_prefix,
    name_template="azint_{dataset}_{scan_num}_{detector}.h5",
    create_dirs=False,
    mode=None,
):
    """
    From a list of datasets, returns a list of output paths corresponding to each file name,
    mirroring the directory structure under RAW_DATA into output_prefix.

    Parameters
    ----------
    datasets: list of HDF5Dataset
        List HDF5Dataset objects. The file path is expected to be in the ID15A layout "proposal/RAW_DATA/sample/dataset/scan(s)"
        example: /data/id15/inhouse6/ihxxx/id15a/RAW_DATA/sampletest/sampletest_0008/sampletest_0008.h5 (possibly several scans therein)
    output_prefix: str
        Common path where output files will be saved
    name_template: str
        Template name. For example azint_{dataset}_{scan}_{detector}.h5 for azint_ACL9011001b_0009_scan0001_pilatus.h5
    create_dirs: bool, optional
        Whether to create output directories. Default is False.

    Returns
    -------
    output_files: list of str
        List of output files corresponding to each dataset
    datasets_info: dict
        Nested dictionary with the information on samples/datasets/scans
    """
    output_files = []
    datasets_info = {}
    filemode = file_mode_to_octal(mode)

    def _ensure_can_write(directory, filemode):
        try:
            os.chmod(directory, filemode)
        except PermissionError:
            print(f"Warning: cannot chmod {directory} to {oct(filemode)}")

    for dataset in datasets:
        # --- 1) render the filename as before
        file_name = safe_format(
            name_template,
            dataset=path.splitext(path.basename(dataset.fname))[0],
            scan_num=f"scan{int(dataset.entry.split('.')[0]):04d}",
            detector=dataset.detector_name,
        )

        # --- 2) figure out the subdirs under RAW_DATA
        parts = dataset.fname.split(os.sep)
        if "RAW_DATA" in parts:
            i = parts.index("RAW_DATA")
            subdirs = parts[
                i + 1 : -1
            ]  # everything between RAW_DATA/…/ and the file itself
        else:
            subdirs = []

        # --- 3) assemble the final output directory
        out_dir = path.join(output_prefix, *subdirs)

        # --- 4) full outfile path
        out_file = path.join(out_dir, file_name)
        output_files.append(out_file)

        # --- 5) create & chmod if requested
        if create_dirs:
            pathlib_Path(out_dir).mkdir(parents=True, exist_ok=True)
        _ensure_can_write(out_dir, filemode)

        # --- (optional) still fill your datasets_info if you need it
        # here we key by the datasets own HDF5 url path
        datasets_info[dataset.dataset_hdf5_url.path()] = {
            "dataset": dataset,
            "output_file": out_file,
        }

    # finally also ensure the top level prefix is writable
    _ensure_can_write(output_prefix, filemode)

    return output_files, datasets_info


# This function might be moved elsewhere in the future
def get_distributed_integrator_params(
    conf_file,
    overwritten_params=None,
    conf_dict=None,
    print_already_processed=True,
    ignore_already_processed=True,
    parsing_error_handling="print",
    extra_metadata_fields=None,
    logger=None,
):
    """
    Parse and validate a configuration file, return several data structures detailed below.

    Parameters
    ----------
    conf_file: str
        Path to the configuration file
    overwritten_params: dict, optional
        Dictionary of parameters overwriting some options of the configuration file
    conf_dict: dict, optional
        Configuration dict to use instead of first parameter 'conf_file'.
    print_already_processed: bool, optional
        Print datasets already processed (if existing_output = skip)
    ignore_already_processed: bool, optional
        Remove already-processed datasets from the results list. Default is True
    parsing_error_handling: str, optional
        How to handle datasets parsing errors. Can be "print", "raise" or "pass".
        Default is "print".
    extra_metadata_fields: List[str], optional
        Extra fields to grab from HDF5 file. Currently, only supports attributes
        found in "measurement", "instrument/positioners", or "instrument/fscan_parameters".
    logger : Logger or None
        Logger instance or None; will default to printing.

    Returns
    -------
    conf: dict
        Nested dictionary extracted from the configuration file
    ai_config: AIConfiguration
        Data structure with the azimuthal integration configuration
    comp_resources: ComputationResources
        Data structure with the computation resources
    datasets: list of Dataset objects
        List where each item is a data structure with information on the dataset to process
    output_files: list of str
        List where each item is the path to the output file, either a .nx file or a directory
    """
    # Parse/validate configuration file
    if conf_dict is None:
        conf_dict = parse_config_file(conf_file)
        conffile_mtime = os.path.getmtime(conf_file)
    conf = validate_config(
        conf_dict,
        integrator_config,
        renamed_keys,
    )
    overwritten_params = overwritten_params or {}
    conf = overwrite_config(conf, overwritten_params)

    # Get AI configuration
    ai_config = get_ai_config(conf)

    # Get computation resources configuration
    comp_resources = get_computations_config(conf)

    os.environ["OMP_NUM_THREADS"] = str(comp_resources.n_threads_per_worker)

    # Get detector name.
    detector_name = conf["azimuthal integration"]["detector_name"].strip()
    if conf["azimuthal integration"]["detector_suffix"] is not None:
        detector_suffix = conf["azimuthal integration"]["detector_suffix"].strip()
    else:
        detector_suffix = conf["azimuthal integration"]["detector_suffix"]

    # Get beamline-specific dataset layout
    # dataset_layout = conf["dataset"]["layout"]
    detector_name = detector_name or FscanDataset._default_detector_name

    # Browse dataset
    dataset_fname = conf["dataset"]["location"]
    if dataset_fname is None:
        # No dataset provided
        dataset_list = [None]
    else:
        datasets = dataset_fname.split(",")
        dataset_list = browse_datasets_wildcard(
            datasets,
            entry=conf["dataset"]["hdf5_entry"],
            detector_name=detector_name,
            detector_suffix=detector_suffix,
            errors=parsing_error_handling,
            extra_metadata_fields=extra_metadata_fields,
            logger=logger,
        )

    # Output file
    output_file = conf["output"]["location"]
    if output_file is None:
        if dataset_list[0] is not None:
            raise ValueError(
                "Output file/directory must be provided when dataset location is provided",
            )
        output_file_list = [None]
    output_file_list, output_folders_layout = get_output_folders(  # noqa: RUF059
        dataset_list,
        output_file,
        create_dirs=True,
        mode=conf["output"]["file_mode"],
    )

    # Skip datasets that are already processed
    what_to_do_if_exists = conf["output"]["existing_output"]
    if (
        what_to_do_if_exists in ["skip", "reprocess_if_conffile_more_recent"]
        and ignore_already_processed
    ):

        def print_if_necessary(msg):
            if print_already_processed:
                print(msg)

        new_dataset_list = []
        new_output_file_list = []
        for dataset, output_file in zip(dataset_list, output_file_list):
            skip = False
            if path.isfile(output_file) and (
                (
                    what_to_do_if_exists == "skip"
                    and is_same_config(output_file, ai_config)
                )
                or (
                    what_to_do_if_exists == "reprocess_if_conffile_more_recent"
                    and os.path.getmtime(output_file) > conffile_mtime
                )
            ):
                skip = True
                # # DEBUG
                # elif what_to_do_if_exists == "reprocess_if_conffile_more_recent" and os.path.getmtime(output_file) <= conffile_mtime:
                #     skip = False
                #     print("Conf file more recent than %s - will have to reprocess!" % output_file)
                #     from datetime import datetime
                #     print(datetime.fromtimestamp(conffile_mtime).isoformat())
                #     print(datetime.fromtimestamp(os.path.getmtime(output_file)).isoformat())
                # # ----
            if skip:
                print_if_necessary(
                    str("Discarding dataset %s: already processed" % (str(dataset))),
                )
            else:
                new_dataset_list.append(dataset)
                new_output_file_list.append(output_file)
        dataset_list = new_dataset_list
        output_file_list = new_output_file_list

    return conf, ai_config, comp_resources, dataset_list, output_file_list
