from os import path
from pickle import dump as pickle_dump  # nosec B403
from pickle import load as pickle_load  # nosec B403

from silx.io.url import DataUrl

from .utils import get_identifier, partition_list


def get_partial_output_file(
    output_file,
    scan_file,
    partial_output_files_subfolder=None,
):
    partial_file_basename = (
        path.basename(path.splitext(output_file)[0]) + "_" + scan_file.split("_")[-1]
    )
    partial_file_dirname = path.join(
        path.dirname(output_file),
        partial_output_files_subfolder or "",
    )
    return path.join(partial_file_dirname, partial_file_basename)


class ProcessingConfig:
    """
    A class that holds information on
      - input datasets
      - final output files
      - partial output files
      - Azimuthal Integration parameters
      - Computation distribution parameters
    """

    def __init__(
        self,
        ai_config,
        computations_config,
        scans_list,
        output_files_list,
        conf_dict=None,
        partial_output_files_subfolder=None,
        workspace_directory=None,
    ):
        self.ai_config = ai_config
        self.computations_config = computations_config
        self.scans_list = scans_list
        self.output_files_list = output_files_list
        self.conf_dict = conf_dict
        self.workspace_directory = workspace_directory
        self.partial_output_files_subfolder = partial_output_files_subfolder or ""
        # Sanitize
        if self.partial_output_files_subfolder.startswith("/"):
            self.partial_subfolder = self.partial_output_files_subfolder[1:]
        self.build_tasks_distribution()

    def build_tasks_distribution(self):
        """
        Build a list of task to assign to each worker.

        There are two levels of tasks subdivision:
          - Each "multi-processes worker" (MP) handles one or several scans (thousands of images).
            This MP spawns several workers (W).
          - Each worker (W) handles one or several acquisition file (hundreds of images per file).
            It saves results to partial files (to avoid parallel write).


        Returns
        -------
        workers_tasks: list of lists of tuple
            List of n_workers items, where each item if a list of (input_url, output_fname).
            i.e workers_tasks[worker_id] is a list of tuple in the form (input_url, output_fname)

        scans_info: dict
            Dictionary holding information on all the scans.
            For each scan, scans_info[scan_url_txt] has several fields:
              - "scan_structure": the associated HDF5Dataset object
              - "partial_scan_urls": list of DataUrl: of files constituting the scan
              - "partial_output_files": list of str: partial output files
              - "output_file": str: final output file

        """
        out_subfolder = self.partial_output_files_subfolder or ""

        # Other helper structures
        # partial_scan_urls_to_main_scan = {}
        all_partial_scans_and_outputs = []

        # scan_info[scan_url] = {"partial_scan_urls": ..., "partial_output_files": ..., "output_file": ...}
        scans_info = {}
        for scan, out_file in zip(self.scans_list, self.output_files_list):
            partial_scan_urls = []
            partial_output_files = []
            for partial_scan_file, h5_path in scan.get_virtual_sources().items():
                partial_scan_url = DataUrl(
                    file_path=partial_scan_file,
                    data_path=h5_path,
                )
                partial_scan_urls.append(partial_scan_url)
                partial_output_file = get_partial_output_file(
                    out_file,
                    partial_scan_file,
                    partial_output_files_subfolder=out_subfolder,
                )
                partial_output_files.append(partial_output_file)

                # partial_scan_urls_to_main_scan[partial_scan_url.path()] = scan
                all_partial_scans_and_outputs.append(
                    (partial_scan_url, partial_output_file),
                )

            scans_info[scan.dataset_hdf5_url.path()] = {
                "scan_structure": scan,
                "partial_scan_urls": partial_scan_urls,
                "partial_output_files": partial_output_files,
                "output_file": out_file,
            }

        # Tasks distribution among workers
        # Simply divide all partial files between multi-workers (assuming homogeneous computing power)
        workers_tasks = partition_list(
            all_partial_scans_and_outputs,
            self.computations_config.n_multi_worker,
        )

        return workers_tasks, scans_info

    def dump_tasks_to_disk(self, tasks):
        workspace_dir = self.workspace_directory or path.realpath(
            ".integrator_workspace",
        )
        fname = path.join(
            workspace_dir,
            "tasks_%s.pickle" % (get_identifier(use_date=False)),
        )
        structure_to_save = {
            "ai_config": self.ai_config,
            "computations_config": self.computations_config,
            "tasks": tasks,
        }
        with open(fname, "wb") as f:
            pickle_dump(structure_to_save, f)
        return fname


def load_work_config(fname):
    # WARNING: Using pickle to deserialize data is unsafe if the source is untrusted.
    # Only load pickle files from sources that you trust.
    with open(fname, "rb") as f:
        work_config = pickle_load(f)  # nosec B301
    return work_config
