import logging
import os
import posixpath
import shlex
from ast import literal_eval
from contextlib import ContextDecorator
from datetime import datetime
from pathlib import Path
from subprocess import CalledProcessError, check_output  # nosec B404
from threading import Timer
from time import sleep

import numpy as np
from h5py import File as H5pyFile
from nabu.io.writer import export_dict_to_h5, merge_hdf5_files
from nabu.resources.logger import LoggerOrPrint
from psutil import cpu_count
from scipy import constants
from tomoscan.utils.hdf5 import get_data_sources

from .resources.utils import get_resource_file

path = os.path

ai_units = {
    "q_a^-1": "q_A^-1",
    "q_A^-1": "q_A^-1",
    "2th_deg": "2th_deg",
    "r_mm": "r_mm",
}

ALLOWED_COMMANDS = {
    "nvidia-smi",
    "ls",
    "echo",
    "integrate-mp",
    "integrate-mp-worker",
    "integrate-slurm",
    "sbatch",
    "squeue",
    "scontrol",
    "scancel",
    "sinfo",
    "sacct",
}


class PeriodicAction:
    def __init__(self, interval, function, args=None, kwargs=None):
        self.function = function
        self.interval = interval
        self._args = args
        self._kwargs = kwargs

    def _function_wrapper(self, *args, **kwargs):
        self.function(*args, **kwargs)
        self.start()

    def start(self):
        self.timer = Timer(
            self.interval,
            self._function_wrapper,
            args=self._args,
            kwargs=self._kwargs,
        )
        self.timer.start()

    def stop(self):
        self.timer.cancel()

    cancel = stop


class PatchLogLevel(ContextDecorator):
    def __init__(self, prefix, loglevel):
        self.prefix = prefix
        self.new_loglevel = loglevel
        self.base_loglevels = {
            k: v.getEffectiveLevel()
            for k, v in logging.root.manager.loggerDict.items()
            if k.startswith(prefix) and isinstance(v, logging.Logger)
        }

    def __enter__(self):
        for logger_name in self.base_loglevels:
            logging.getLogger(logger_name).setLevel(self.new_loglevel)

    def __exit__(self, exc_type, exc, exc_tb):
        for logger_name, old_log_level in self.base_loglevels.items():
            logging.getLogger(logger_name).setLevel(old_log_level)


"""
    with PatchLogLevel("nxtomomill", logging.CRITICAL):
        pass
"""


def make_azims(azim_min, azim_max, nb_slices):
    cake = (float(azim_min), float(azim_max))
    npieces = int(nb_slices)
    # -- calculate azim slices
    if npieces == 1:
        azims = [(cake[0], cake[1])]
    # -- calculate azim slices if n.pieces > 1
    elif npieces > 1:
        vals = np.linspace(cake[0], cake[1], npieces + 1)
        azims = [(vals[i], vals[i + 1]) for i in range(npieces)]
    return np.array(azims)


def parse_poni_file(poni_file):
    conf = {}
    with open(poni_file) as opened_file:
        for line in opened_file:
            line = line.strip()
            if line.startswith("#") or (":" not in line):
                continue
            words = line.split(":", 1)
            key = words[0].strip().lower()
            try:
                value = words[1].strip()
            except Exception as error:
                print("Error %s with line: %s", error, line)
            conf[key] = value
    return conf


def check_poni_file_detector(poni_file, detector, force_replace=False):
    """
    Check a poni file so that the detector is replaced with another (user-provided) detector.
    This is useful when the detector file name within the poni file does not exist, or cannot be
    accessed.

    Parameters
    ----------
    poni_file: str
        Path to the poni file
    detector: str
        Name or path of the detector used to replace the one in the poni file
    force_replace: bool, optional
        Whether to use the provided detector even if the one in the poni file can be accessed
    """
    conf = parse_poni_file(poni_file)
    if "detector_config" not in conf:
        return conf
    det_conf = literal_eval(conf["detector_config"]) or {}
    det_fname = det_conf.get("filename", None)
    if force_replace or (det_fname is not None and not path.isfile(det_fname)):
        conf["detector_config"] = conf["detector_config"].replace(
            det_conf["filename"],
            detector,
        )
    return conf


def get_h5_config_from_pyfai(pyfai_config, ai_config, dataset):
    h5_config = {
        "ai_config": str(ai_config),
        # CHECKME not sure about the order. But for now both are the same (from intrange code)
        "x_pixel_size": pyfai_config["pixel_size"][0],
        "y_pixel_size": pyfai_config["pixel_size"][1],
        ("x_pixel_size", "unit"): "mm",
        ("y_pixel_size", "unit"): "mm",
        "S-D distance": pyfai_config["S-D distance"],
        ("S-D distance", "unit"): "mm",
        "wavelength": pyfai_config["wavelength"],
        ("wavelength", "unit"): "angstrom",
        "pyfai_method": pyfai_config["pyfai_method"],
        "exposure time": dataset.exposure_time,
        ("exposure time", "unit"): "s",
    }
    return h5_config


def create_output_masterfile(
    files,
    output_file,
    ai_config,
    do_stack_mean,
    h5_config,
    entry="/",
    process_name="azimuthal integration",
):
    local_files = files_relative_paths(files, output_file)
    azimuthal_ranges = make_azims(*ai_config.azimuthal_range)
    for azim_idx in range(len(azimuthal_ranges)):
        base_path = posixpath.join(
            entry,
            process_name,
            "results",
            str("azim_%02d" % (azim_idx + 1)),
        )
        h5_paths = [
            posixpath.join(base_path, what)
            for what in ["intensities", "radial_positions"]
        ]
        if (
            ai_config.error_model is not None and ai_config.trim_method is None
        ):  # no error propagation when doing median/sigma-clip (?)
            h5_paths.append(posixpath.join(base_path, "error"))
        if do_stack_mean:
            h5_paths.append(posixpath.join(base_path, "averaged stack"))
        for h5_path in h5_paths:
            merge_hdf5_files(
                local_files,
                h5_path,
                output_file,
                process_name,
                output_entry=entry,
                config=None,
                base_dir=path.dirname(output_file),
                overwrite=True,
                # Not very elegant. 'data_name' expects something like "intensities",
                # but we pass something like 'azim_01/intensities'.
                # This works (h5py seems to create groups recursively) but might be brittle.
                data_name=posixpath.join(
                    str("azim_%02d" % (azim_idx + 1)),
                    path.basename(h5_path),
                ),
            )
    # Save configuration
    h5_config["azimuthal_ranges"] = azimuthal_ranges
    h5_config["radial_range"] = ai_config.radial_range
    h5_config["result_files"] = local_files
    # h5_config["result_stacks"] = {}
    # for task_desc in self._tasks.values():
    # fname = path.basename(task_desc["result"])
    # h5_config["result_stacks"][fname] = task_desc["range"]
    export_dict_to_h5(
        h5_config,
        output_file,
        posixpath.join(entry, process_name, "configuration"),
    )


def browse_virtual_sources(dataset, logger):
    fname = dataset.dataset_hdf5_url.file_path()
    data_path = dataset.dataset_hdf5_url.data_path()
    virtual_sources = get_data_sources(fname, data_path)
    first_source = next(iter(virtual_sources.keys()))
    if first_source == fname:
        # not a virtual dataset
        logger.error("Not a virtual dataset: %s" % fname)
        return None
    n_tries = 0
    vs = virtual_sources[first_source]
    dataset_virtual_sources = list(vs.keys())
    # Get the shape of first non-virtual dataset, instead of browsing everything
    while not isinstance(vs, str) and n_tries < 10:
        vs_fname = next(iter(vs.keys()))
        vs = vs[vs_fname]
        n_tries += 1
    if n_tries >= 10:
        logger.error("Too many indirections in %s, stopping" % fname)
        vs_fname, vs = None, None
    dataset_first_vsource_fname = vs_fname
    dataset_first_vsource_path = vs
    return (
        dataset_virtual_sources,
        dataset_first_vsource_fname,
        dataset_first_vsource_path,
    )


def retry_create_masterfile(n_retries, logger):
    def decorator(func):
        def wrapper(*args, **kwargs):
            files_list = args[0]
            success = False
            exception_to_raise = None
            for i in range(n_retries):
                if i > 0:
                    logger.info(
                        "Retrying to create master file (attempt %d/%d)"
                        % (i, n_retries),
                    )
                try:
                    func(*args, **kwargs)
                    success = True
                except KeyError as exc:
                    success = False
                    exception_to_raise = exc
                    logger.exception("Error when creating master file !")
                    sleep(2)
                    for fname in files_list:
                        Path(fname).touch()
                    sleep(2)
                if success:
                    break
            if not (success):
                raise exception_to_raise

        return wrapper

    return decorator


def retry_n_times(n_times, logger, func, *args, **kwargs):

    def log_message(msg):
        if logger is not None:
            logger.info(msg)
        else:
            print(msg)

    ret = "n/a"
    for i in range(n_times):
        try:
            ret = func(*args, **kwargs)
        except Exception as exc:
            log_message("==== Something went wrong, re-trying %d/%d" % (i + 1, n_times))
            ret = "error"
            ret_exc = exc
        if ret not in {"n/a", "error"}:
            break
        sleep(5)
    if ret == "error":
        raise ret_exc
    return ret


def walltime_to_seconds(walltime):
    """
    Converts a "walltime" formated string to a number of seconds.

    Parameters
    ----------
    walltime: str
        Walltime, in the form "Hours:Minutes:Seconds"


    Returns
    -------
    secs: int
        Number of seconds corresponding to walltime

    Examples
    --------
    walltime_to_seconds("02:50:27") gives 10227

    """
    components = list(map(int, walltime.split(":")))
    if len(components) != 3:
        raise ValueError("Expected format H:M:S")
    return (np.array(components) * np.array([3600, 60, 1])).sum()


def whoami():
    ret = None
    try:
        ret = os.getlogin()
    except OSError:  # can happen eg. in a SLURM reservation
        ret = os.environ.get("LOGNAME", None)
    return ret


def get_identifier(use_date=True):
    """
    Return a simple identifier in the form {whoami}_{date}
    """
    if use_date:
        res = "%s_%s" % (
            whoami(),
            datetime.now().astimezone().strftime("%Y_%m_%d_%H_%M_%S"),
        )
    else:
        res = "%s" % (whoami())
    return res


# TODO move to nabu
def partition_list(list_, n_sets):
    indices = np.arange(len(list_))
    sets_indices = np.array_split(indices, n_sets)
    results = []
    for set_indices in sets_indices:
        r = [list_[idx] for idx in set_indices]
        results.append(r)
    return results


# TODO move to nabu
def files_relative_paths(files, reference_file):
    """
    From a list of (absolute) files paths, return a list of relative paths.
    Each path is relative to 'reference_file'.
    """
    reference_dir = path.dirname(reference_file)
    local_files = []
    for file in files:
        relpath = path.relpath(file, start=reference_dir)
        # Sanitize paths like ".." or "../stuff"
        relpath = relpath.removeprefix("..")
        relpath = relpath.removeprefix("../")
        local_files.append(relpath)
    return local_files


def get_worker_threads_affinity(n_workers, n_threads_per_worker):
    """
    Given a number of workers (processes) and number of threads per worker,
    returns the optimal affinity list for each worker.

    Note that it likely won't work properly if the current host is used through
    a task scheduler(eg. SLURM). In this case, one must ensure that at least
    one full socket is allocated.
    """
    threads_per_core = cpu_count(logical=True) // cpu_count(logical=False)
    # n_available_threads = len(os.sched_getaffinity(0))
    physical_cores = sorted(os.sched_getaffinity(0))[::threads_per_core]

    workers_phys_cores = partition_list(physical_cores, n_workers)
    if len(workers_phys_cores[0]) >= n_threads_per_worker:
        return [wpc[:n_threads_per_worker] for wpc in workers_phys_cores]

    wpc = np.array(workers_phys_cores)
    wt = wpc
    while wt.size < n_workers * n_threads_per_worker:
        wt = np.hstack([wt, wpc + 1])
        wpc += 1
    return list(map(sorted, wt.tolist()))


# TODO move to nabu
def get_opencl_devices(device_type="gpu", excluded_platform=None):
    # Do it here, otherwise contexts are sometimes created on all platforms,
    # which will make further forks()/spawn() crash
    from pyopencl import (
        device_type as cl_device_type,
    )
    from pyopencl import (
        get_platforms as cl_get_platforms,
    )

    dev_type = {"cpu": cl_device_type.CPU, "gpu": cl_device_type.GPU}[
        device_type.lower()
    ]
    excluded_platform = excluded_platform or []
    tuples = []
    for i, platform in enumerate(cl_get_platforms()):
        if platform.name in excluded_platform:
            continue
        for j, device in enumerate(platform.get_devices()):
            if device.type == dev_type:
                tuples.append((i, j))
    return tuples


# TODO move to nabu (?)
class _DefaultFormat(dict):
    """
    https://docs.python.org/3/library/stdtypes.html
    """

    def __missing__(self, key):
        return key


def safe_format(str_, **kwargs):
    """
    Alternative to str.format(), but does not throw a KeyError when fields are missing.
    """
    return str_.format_map(_DefaultFormat(**kwargs))


def exec_shell_command(cmd):
    parts = shlex.split(cmd)
    if parts[0] not in ALLOWED_COMMANDS:
        raise ValueError(f"Untrusted command: {parts[0]}")
    try:
        res = check_output(parts)  # nosec B603
    except RuntimeError as rerr:
        print("Error: %s" % rerr)
        return None
    except OSError as oerr:
        cmd0 = parts[0]
        print(
            "Something went wrong when executing %s. Please verify that the command %s is available"
            % (cmd, cmd0),
        )
        print("Error: %s" % oerr)
        return None
    except CalledProcessError as err:
        print("Error: %s " % err)
        return None
    else:
        return res


def get_number_of_available_gpus(on_error="ignore"):
    """
    Get the number of available GPUs.

    Parameters
    ----------
    on_error: str, optional
        What to do on error. Can be "ignore", "print" or "raise".
    """
    # The current implementation is just a parsing of the "nvidia-smi" command.
    # Yes, this is terrible, and pyopencl should be used instead.
    # But in the current state of things:
    #   - pyopencl.get_platforms() sometimes (depending on the machine/setup) creates contexts
    #     all over the place, which will mess up further attemps to pick a device or simply fork/spawn.
    #   - We don't use GPUs other than Nvidia ones
    #
    # An equivalently-ugly but slightly better alternative would be to parse the output of
    # "clinfo", but (1) this command is not always available, (2) much more tedious to parse

    def handle_error(msg):
        if on_error == "raise":
            raise ValueError(msg)
        elif on_error == "print":
            print(msg)

    r = exec_shell_command("nvidia-smi -L")
    if r is None:
        handle_error("Could not detect any GPU")
        return 0
    devs = [x for x in r.decode().split(os.linesep) if len(x) > 0]
    if len(devs) == 0 or devs[0] == "No devices found.":
        handle_error("Could not detect any GPU")
        return 0
    return len(devs)


def file_mode_to_octal(file_mode, default=0o755):
    """
    Convert a file mode provided as a "regular int" (base 10) to an octal number of os.chmod()
    """
    if file_mode is None:
        return default
    try:
        res = int(str(file_mode), base=8)
    except ValueError:  # already in octal ?
        res = file_mode
    return res


#
# Get additional metadata.
# Code from Stefano Checchia in hdintrange 0.12


def get_additional_metadata(
    scan_number,
    h5abspath,
    axes=None,
    detector_name=None,
    logger=None,
):
    """
    Extract specified metadata arrays from an HDF5 dataset.

    Parameters
    ----------
    scan_number : str
        Prefix of the scan entries to search (e.g. '6').
    h5abspath : str
        Absolute path to the HDF5 file.
    axes : list of str, optional
        List of metadata field names to extract (e.g. ['hry', 'hrz']).
    detector_name : str, optional
        Name of the detector group under instrument.
    logger : Logger or None
        Logger instance or None; will default to printing.

    Returns
    -------
    out : dict
        Mapping of metadata names to NumPy arrays or scalars.
    """
    out = {}
    axes = axes or []
    logger = LoggerOrPrint(logger)
    logger.debug(f"Searching {h5abspath} for metadata fields: {axes}")

    # Determine scan entries (e.g. ['6.1', '6.2', ...])
    with H5pyFile(h5abspath, "r") as hf:
        scan_entries = [k for k in hf if k.split(".")[0] == scan_number]
        if not scan_entries:
            logger.warning(f"No entries found for scan {scan_number}")
            return out

        # Categories to search, in order
        categories = [
            "measurement",
            "instrument/positioners",
            "instrument/fscan_parameters",
        ]

        # Extract each requested axis
        for tgt in axes:
            found = False
            for cat in categories:
                for sn in scan_entries:
                    logger.debug(f"Looking for {tgt} in {sn}/{cat}")
                    path = f"{sn}/{cat}/{tgt}"
                    if path in hf:
                        data = hf[path][()]
                        # Interpolate if this is a .2 entry and not a time field
                        if sn.endswith(".2") and tgt not in (
                            "epoch",
                            "fpico2",
                            "fpico3",
                        ):
                            t2 = hf[f"{sn}/measurement/epoch"][()]
                            t1 = hf[f"{sn.replace('.2', '.1')}/measurement/epoch_trig"][
                                ()
                            ]
                            data = np.interp(t1, xp=t2, fp=data)
                        # Round and cast
                        out[tgt] = np.round(data, 6).astype(np.float64)
                        logger.debug(f"Found {tgt} at {path}")
                        found = True
                        break
                if found:
                    break
            if not found:
                logger.warning(f"Metadata '{tgt}' not found in file {h5abspath}")

    return out


#
# Detector attenuation caused by non-zero thickness
# Code from Stefano Checchia in hdintrange 0.13


def make_detector_atten(wavelength_m, tth_rad, d_cm=0.1, logger=None):
    """
    Econst (keV/A)
    X = 2theta_rad, 1D or 2D
    mu = linear attenuation coef (1/cm)
    d = sensor thickness (cm)
    K = exp((-mu*d)/cos(X))
    intens_final = intens_ini / K
    """
    ene_ev = (constants.h * constants.c / wavelength_m) / constants.e
    cdtefile = get_resource_file("cdte_mu_cm.dat")
    logger = LoggerOrPrint(logger)

    try:
        spl_ene, spl_att = np.loadtxt(
            cdtefile,
            unpack=True,
            usecols=(0, 1),
            comments="#",
        )
        mu = spl_att[np.argmin(abs(spl_ene - ene_ev))]
        logger.debug(f"Found {cdtefile}")

    except OSError as err:
        xcom_ene = 1e3 * np.array(
            [
                5,
                6,
                8,
                10,
                15,
                20,
                26.71,
                26.71,
                30,
                31.81,
                31.81,
                40,
                50,
                60,
                70,
                80,
                90,
                100,
                110,
                120,
                130,
                140,
            ],
        )
        # mass attenuation coefficients (cm2/g) for CdTe from XCOM database
        xcom_att = np.array(
            [
                8.392e02,
                5.286e02,
                2.492e02,
                1.381e02,
                4.656e01,
                2.144e01,
                9.832e00,
                2.943e01,
                2.182e01,
                1.874e01,
                3.494e01,
                1.930e01,
                1.067e01,
                6.542e00,
                4.321e00,
                3.019e00,
                2.206e00,
                1.671e00,
                1.305e00,
                1.046e00,
                8.566e-01,
                7.152e-01,
            ],
        )
        sensor_density = 5.85
        spl = np.interp(ene_ev, xcom_ene, np.log10(xcom_att))
        mu = 10**spl * sensor_density
        logger.warning(f"{err} Recalculating value for {np.round(ene_ev):.0f} eV")

    logger.debug(f"CdTe LAC at {np.round(ene_ev):.0f} eV = {mu:.1f} / cm")
    abs_corr_fact_k = (1 - np.exp(-mu * d_cm / np.cos(tth_rad))) / (
        1 - np.exp(-mu * d_cm)
    )

    return abs_corr_fact_k
