import sys
from multiprocessing import Process
from os import chmod, getcwd, mkdir, path
from pprint import pprint

from nabu.app.utils import parse_params_values
from nabu.pipeline.config_validators import convert_to_bool
from nabu.resources.logger import Logger
from tomoscan.io import HDF5File

from ..monitoring import monitor_scans
from ..processconfig import ProcessingConfig
from ..writer import watch_scans_and_create_output_files
from .utils import get_distributed_integrator_params

CheckDatasetsConfig = {
    "config_file": {
        "help": "Path to the configuration file.",
        "default": "",
        "mandatory": True,
    },
    "check_output": {
        "help": "Whether to check output file (shape, links, etc). Default is True.",
        "default": "1",
    },
    "list_missing_datasets": {
        "help": "Whether to display datasets that are not integrated yet. Default is True.",
        "default": "1",
    },
    "list_integrated_datasets": {
        "help": "Whether to print integrated datasets. Default is True.",
        "default": "1",
    },
}


def frame_print(msg, width=80, frame_chr="-"):
    print(frame_chr * width)
    print(msg)
    print(frame_chr * width)


def check_datasets_cli():

    args = parse_params_values(
        CheckDatasetsConfig,
        parser_description="Check the progression of datasets integration",
    )
    conf_file = args["config_file"]

    # Check boolean arguments
    bool_args = {}
    for arg_name in [
        "check_output",
        "list_missing_datasets",
        "list_integrated_datasets",
    ]:
        bool_args[arg_name], err = convert_to_bool(args["check_output"])
        if err is not None:
            print("Error: invalid argument for parameter '%s': %s" % (arg_name, err))

    # Folder for log files
    logdir = path.join(path.relpath(getcwd()), "logs")
    if not path.isdir(logdir):
        mkdir(logdir)
        chmod(logdir, 0o755)  # nosec B103

    # List all datasets
    conf, ai_config, comp_resources, datasets, output_files = (
        get_distributed_integrator_params(
            conf_file,
            overwritten_params={"output": {"existing_output": "pass"}},
            parsing_error_handling="pass",
            ignore_already_processed=False,
        )
    )
    if datasets[0] is None or output_files[0] is None:
        print("Error: dataset location and output location have to be provided")
        sys.exit(0)

    # List remaining datasets
    _, _, _, datasets_remaining, output_files_remaining = (
        get_distributed_integrator_params(
            conf_file,
            print_already_processed=False,
            parsing_error_handling="pass",
            # overwritten_params={"output": {"existing_output": "skip"}},
        )
    )

    # Print difference
    n_frames_tot = sum([d.data_shape[0] for d in datasets]) // 1000
    n_frames_rem = sum([d.data_shape[0] for d in datasets_remaining]) // 1000
    frame_print(
        "Integrated datasets: %d/%d (%dk/%dk frames)"
        % (
            len(datasets) - len(datasets_remaining),
            len(datasets),
            n_frames_tot - n_frames_rem,
            n_frames_tot,
        ),
    )

    if bool_args["list_integrated_datasets"]:
        pprint(
            [
                path.basename(out_file)
                for out_file in output_files
                if out_file not in output_files_remaining
            ],
        )

    if bool_args["check_output"]:
        datasets_integrated = [d for d in datasets if d not in datasets_remaining]
        output_integrated = [
            outp for outp in output_files if outp not in output_files_remaining
        ]
        check_integrated_datasets(datasets_integrated, output_integrated, ai_config)

    if bool_args["list_missing_datasets"]:
        print("*" * 80)
        print(
            "Remaining datasets: %d (%dk frames)"
            % (len(datasets_remaining), n_frames_rem),
        )
        pprint([path.basename(d.fname) + ":" + d.entry for d in datasets_remaining])

    if len(datasets_remaining) == 0:
        return

    # Monitor (progressbar for on-going integration)
    process_config = ProcessingConfig(ai_config, comp_resources, datasets, output_files)
    _, scans_info = process_config.build_tasks_distribution()

    # Create final output files if needed
    def do_in_process(func, *args, **kwargs):
        p = Process(target=func, args=args, kwargs=kwargs)
        p.start()
        return p
        # func(*args, **kwargs)

    logger_for_out_files_creation = Logger(
        "out_files_creation",
        level=conf["pipeline"]["verbosity"],
        logfile=path.join(logdir, "out_files_creation.log"),
        console=False,
    )

    remaining_scans_info = {
        d.dataset_hdf5_url.path(): scans_info[d.dataset_hdf5_url.path()]
        for d in datasets_remaining
    }
    do_in_process(
        watch_scans_and_create_output_files,
        remaining_scans_info,
        ai_config,
        conf["dataset"]["layout"],
        period=1,
        logger=logger_for_out_files_creation,
    )
    monitor_scans(remaining_scans_info, period=1, n_threads=4)


def check_integrated_datasets(datasets, output_files, ai_config):

    error_model = ai_config.error_model or None

    def _check_h5_fields(fdesc, expected_shape):
        err = None
        shape_intensities, shape_error = None, None
        try:
            shape_intensities = fdesc["intensities"].shape
            if error_model is not None:
                shape_error = fdesc["error"].shape
        except Exception as exc:
            err = "Cannot access intensities/error: %s" % (str(exc))
        if err is not None:
            return err
        if shape_intensities != expected_shape:
            err = "Incorrect output shape for 'intensities' (expected %s, got %s)" % (
                expected_shape,
                shape_intensities,
            )
        if error_model is not None and shape_error != expected_shape:
            err = "Incorrect output shape for 'error' (expected %s, got %s)" % (
                expected_shape,
                shape_error,
            )
        return err

    for dataset, output_file in zip(datasets, output_files):
        expected_shape = (dataset.data_shape[0], ai_config.n_pts)
        err = None
        with HDF5File(output_file, "r") as f:
            for azim_slice in f:
                if azim_slice.startswith("azim_"):
                    err = err or _check_h5_fields(f[azim_slice], expected_shape)
        if err is not None:
            print(
                "Invalid integration %s -> %s: %s" % (dataset.fname, output_file, err),
            )
