import os
from dataclasses import dataclass
from multiprocessing import Process
from platform import processor
from traceback import print_exception

from nabu.resources.logger import LoggerOrPrint, PrinterLogger

from .utils import (
    get_number_of_available_gpus,
    get_opencl_devices,
    get_worker_threads_affinity,
    partition_list,
)

path = os.path


@dataclass
class ComputationResources:
    n_multi_worker: int = 1
    n_workers_per_multiworker: int = 8
    n_threads_per_worker: int = 4
    memory: str = "100GB"
    walltime: str = "02:00:00"
    queue: str = "gpu"


def worker_scan(
    worker_name,
    ai_config,
    n_threads,
    scans,
    output_files,
    extra_options=None,
    n_retries=1,
):
    # Do it before importing StackIntegrator/pyFAI/h5plugin, which will mess with CPU affinity
    os.environ["OMP_NUM_THREADS"] = str(n_threads)
    #
    # Configure extra options
    extra_options = extra_options or {}
    # do_stack_mean = extra_options.get("do_stack_mean", False)
    cpu_affinity = extra_options.get("cpu_affinity", None)
    # Configure logger
    # logger = Logger(
    #     "integrator_mp_%s" % worker_name,
    #     level=extra_options.get("log_level", "info"),
    #     logfile=path.join(extra_options.get("log_dir", ""), "integrator_mp_%s.log" % worker_name)
    # )
    logger = PrinterLogger()
    # Configure CPU affinity
    if cpu_affinity is not None:
        cpu_affinity = list(map(int, cpu_affinity))  # fix unsupported numpy.int
        logger.debug("Will use affinity %s" % (str(cpu_affinity)))
        os.sched_setaffinity(0, cpu_affinity)
    # Pick opencl device
    if processor() == "ppc64le":
        excluded_platform = "NVIDIA CUDA"  # no nvidia opencl compiler on power9
    else:
        excluded_platform = (
            "Portable Computing Language"  # too much trouble for old versions
        )
    opencl_devices = get_opencl_devices(
        device_type="gpu",
        excluded_platform=excluded_platform,
    )
    dev_num = extra_options.pop("worker_device_number", 0)
    if len(opencl_devices) > 0:
        extra_options["target_device"] = opencl_devices[dev_num]
        logger.info("Using openCL device: %s" % str(extra_options["target_device"]))

    from integrator.integrator import StackIntegrator

    # TODO "existing_output"
    scan_integrator = StackIntegrator(
        ai_config,
        logger=logger,
        extra_options=extra_options,
    )

    for scan, output_file in zip(scans, output_files):
        # TODO more defensive - use n_retries
        try:
            scan_integrator.process_data(scan, output_file)
        except Exception as exc:
            print("Something went wrong while processing dataset %s" % scan)
            print_exception(exc)


class MultiprocIntegrator:
    _default_extra_options = {
        "device_type": "gpu",
        "log_dir": "logs",
    }

    def __init__(
        self,
        ai_config,
        computations_resources,
        logger=None,
        extra_options=None,
        # **integrator_kwargs,
    ):
        self.logger = LoggerOrPrint(logger)
        self.ai_config = ai_config
        self._configure_extra_options(extra_options)
        self._configure_computations_resources(computations_resources)

    def _configure_extra_options(self, user_extra_options):
        self.extra_options = self._default_extra_options.copy()
        self.extra_options.update(user_extra_options or {})

    def _configure_computations_resources(self, computations_resources):
        self.resources = computations_resources
        self.n_workers = self.resources.n_workers_per_multiworker
        self.n_threads_per_worker = self.resources.n_threads_per_worker
        self.workers_affinity = get_worker_threads_affinity(
            self.n_workers,
            self.n_threads_per_worker,
        )
        if self.extra_options["device_type"] == "gpu":
            self.n_devices = get_number_of_available_gpus(on_error="raise")

    def process_scans(self, scans, output_files):
        """
        Process one or more scans

        Parameters
        ----------
        scans: list of lists of DataUrl
            List where each item is the DataUrl of a partial scan file.
        output_files: list of str
            List where each item is the path to the partial output file.
        """
        if not isinstance(scans, list):
            scans = [scans]
        if not isinstance(output_files, list):
            output_files = [output_files]

        workers_scans = partition_list(scans, self.n_workers)
        workers_outputs = partition_list(output_files, self.n_workers)

        processes = []
        for i in range(self.n_workers):
            scan_files = workers_scans[i]
            partial_output_files = workers_outputs[i]
            if len(scan_files) == 0:
                # Nothing to do for this worker
                continue
            worker_name = "%04d" % i
            p = Process(
                target=worker_scan,
                args=(
                    worker_name,
                    self.ai_config,
                    self.resources.n_threads_per_worker,
                    scan_files,
                    partial_output_files,
                ),
                kwargs={
                    "extra_options": {
                        "cpu_affinity": self.workers_affinity[i],
                        "worker_device_type": self.extra_options["device_type"],
                        "worker_device_number": i % self.n_devices,
                        "log_level": getattr(self.logger, "level", "debug"),  # info
                        "log_dir": self.extra_options.get("log_dir", ""),
                    },
                },
            )
            p.start()
            processes.append(p)
        self.processes = processes

    def close_processes(self, timeout=None):
        for process in self.processes:
            process.join(timeout=timeout)
