import logging
import sys
from multiprocessing import Process
from os import chmod, getcwd, mkdir, path
from pathlib import Path
from pprint import pformat

from nabu.app.utils import parse_params_values
from nabu.resources.logger import Logger
from nabu.resources.utils import extract_parameters

from .. import version as integrator_version
from ..monitoring import monitor_scans
from ..processconfig import ProcessingConfig
from ..slurm import SlurmResourcesConfiguration, submit_bash_script
from ..utils import get_identifier
from ..writer import watch_scans_and_create_output_files
from .utils import get_distributed_integrator_params

IntegrateMPConfig = {
    "config_file": {
        "help": "Path to the configuration file.",
        "default": "",
        "mandatory": True,
    },
}

task_scheduler_submission_resources = {
    "esrf_p9gpu_half": "--partition p9gpu --gres=gpu:1 -c 64 --cores-per-socket 16 --time=01:00:00 --mem=200G",
    "esrf_p9gpu_full": "--partition p9gpu --gres=gpu:2 -c 128 --cores-per-socket 16 --time=01:00:00 --mem=400G",
}


def integrate_mp_cli(run_with_arg=True, config_path=None):
    if run_with_arg and config_path is None:
        args = parse_params_values(
            IntegrateMPConfig,
            parser_description="Distribute azimuthal integration on one powerful machine",
        )
        conf_file = args["config_file"]
    else:
        conf_file = config_path

    conf, ai_config, comp_resources, datasets, output_files = (
        get_distributed_integrator_params(
            conf_file,
        )
    )
    if len(datasets) == 0:
        print("No dataset to integrate, exiting now.")
        sys.exit(1)
    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)

    # Folders for log files and workspace
    logdir = path.join(path.relpath(getcwd()), "logs")
    workspace_dir = conf["computations distribution"]["workspace_path"] or path.join(
        path.relpath(getcwd()),
        ".integrator_workspace",
    )
    for dir_ in [logdir, workspace_dir]:
        if not path.isdir(dir_):
            mkdir(dir_)
            chmod(dir_, 0o755)  # nosec B103

    logger = Logger(
        "multiproc_integration",
        level=conf["pipeline"]["verbosity"],
        logfile=path.join(logdir, "multiproc_integration.log"),
    )

    logger.info("Will process the following datasets:")
    logger.info(pformat(datasets))

    comp_cfg = conf["computations distribution"]
    partition_name = comp_cfg["partition"]

    # (TODO) Multi-process AI on the local machine
    if partition_name == "local":
        logger.info("Using the local machine to distribute computations")
        raise NotImplementedError

    # ---

    # Build the tasks distribution
    process_config = ProcessingConfig(
        ai_config,
        comp_resources,
        datasets,
        output_files,
        workspace_directory=workspace_dir,
    )
    workers_tasks, scans_info = process_config.build_tasks_distribution()
    workspace_file = process_config.dump_tasks_to_disk(workers_tasks)

    # Get the python environment to use on workers.
    python_executable = str(comp_cfg["python_executable"]).strip()
    if python_executable == "":
        raise ValueError("No python executable configured for worker jobs.")

    _, activate_env_text = determine_environment_parameters(
        python_executable,
        logger=logger,
    )

    # Distribute the work on SLURM
    _, conffile_name = path.split(path.abspath(conf_file))
    script_fname = path.join(
        workspace_dir,
        "%s_%s_integrate_mp.sh"
        % (path.splitext(conffile_name)[0], get_identifier(use_date=False)),
    )
    generate_bash_script(activate_env_text, workspace_file, script_fname)
    # TODO more flexible:
    # "cores-per-socket >= 16" should prevent from getting old nodes
    cores = 32 if partition_name != "p9gpu" else 64
    slurm_resources = SlurmResourcesConfiguration(
        partition=partition_name,
        time=comp_resources.walltime,
        memory=comp_resources.memory,
        cpus_per_task=cores,
        gres="gpu:1",
        cores_per_socket=16,
    )
    ret = submit_bash_script(
        slurm_resources,
        script_fname,
        logdir,
        workspace_dir,
        slurm_command_args=[str(v) for v in range(comp_resources.n_multi_worker)],
    )
    for r in ret:
        if isinstance(r, bytes):
            print(r.decode())

    #
    # Monitor integration progress
    #

    def do_in_process(func, *args, **kwargs):
        p = Process(target=func, args=args, kwargs=kwargs)
        p.start()
        return p

    logger_for_out_files_creation = Logger(
        "out_files_creation",
        level=conf["pipeline"]["verbosity"],
        logfile=path.join(logdir, "out_files_creation.log"),
        console=False,
    )
    do_in_process(
        watch_scans_and_create_output_files,
        scans_info,
        ai_config,
        conf["dataset"]["layout"],
        period=1,
        logger=logger_for_out_files_creation,
        extra_metadata={"config_file": conf_file, "configuration": conf},
        repack_output=conf["output"]["repack_output"],
        repack_output_options={
            "python_env": activate_env_text,
            "slurm_resources": slurm_resources,
            "log_dir": logdir,
            "workspace_dir": workspace_dir,
        },
    )

    monitor_scans(scans_info, period=1, n_threads=4)


def _is_cvmfs_package(python_executable):
    """
    Checks if the python executable is contained in a cvmfs package

    Parameters
    ----------
        path: Path to the python executable
    """
    return python_executable.startswith("/cvmfs/") and "packages" in python_executable


def determine_environment_parameters(python_executable, logger=None):
    """
    Determines how to activate the environment based on the python executable path.

    Parameters
    ----------
      python_executable: Full path to the python executable
      logger: Logger to record information about python environment
    """
    if logger is None:
        logger = logging.getLogger("null")
        logger.addHandler(logging.NullHandler())

    # Derive the environment/module directory (two levels up from the python binary)
    virtual_env = Path(python_executable).parent.parent
    logger.debug("Workers will be using environment/module: %s" % virtual_env)

    if _is_cvmfs_package(python_executable):
        # load slurm module
        activate_env_path = Path(str(virtual_env).replace("packages", "modules"))
        logger.debug("Detected slurm module.")
        logger.debug(f"module load {activate_env_path}")
        activate_env_text = f"""# Load CVMFS module\nmodule load {activate_env_path}"""
    elif (virtual_env / "etc" / "profile.d" / "conda.sh").is_file():
        # activate base conda environment
        logger.debug("Detected conda environment.")
        activate_env_path = virtual_env / "etc" / "profile.d" / "conda.sh"
        activate_env_text = (
            f"""# Initialize conda\nsource {activate_env_path}\nconda activate base"""
        )
        logger.debug(f"source {activate_env_path}\nconda activate base")
    elif (virtual_env.parent.parent / "etc" / "profile.d" / "conda.sh").is_file():
        # activate a named conda environment
        logger.debug("Detected conda environment.")
        activate_env_path = virtual_env.parent.parent / "etc" / "profile.d" / "conda.sh"
        env_name = virtual_env.name
        activate_env_text = f"""# Initialize conda\nsource {activate_env_path}\nconda activate {env_name}"""
        logger.debug(f"source {activate_env_path}\nconda activate {env_name}")
    elif (virtual_env / "bin" / "activate").is_file():
        # source virtual environment
        logger.debug("Detected Python virtual environment.")
        activate_env_path = virtual_env / "bin" / "activate"
        activate_env_text = (
            f"""# Initialize virtual environment\nsource {activate_env_path}"""
        )
        logger.debug(f"source {activate_env_path}")
    else:
        raise ValueError(
            f"Unknown environment type for python executable: {python_executable}. Please enter either a conda or a python virtual environment (.venv).",
        )

    return activate_env_path, activate_env_text


def generate_bash_script(activate_env_text, workspace_file, script_path):
    """
    Generates a bash script that activates the provided conda environment
    and then calls the worker command.

    Parameters
    ----------
      activate_env_text: String required to activate virtual environment
      workspace_file: The workspace file needed by the worker
      script_path: Where to write the bash script
    """
    # The following script sources the virtual environment or conda environment, and then calls the integrate-mp-worker.
    content = f"""\
#!/bin/bash -l

{activate_env_text}

integrate-mp-worker {workspace_file} $1
"""
    with open(script_path, "w") as fdesc:
        fdesc.write(content)
