# coding=utf-8
__author__ = "Dimitris Karkalousos"

# Taken and adapted from: https://github.com/wdika/NeMo/blob/main/nemo/core/utils/process_launcher/launcher.py

import copy
import hashlib
import os
import subprocess
import sys
import threading
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence

import torch
from hydra.core.config_store import ConfigStore
from hydra.core.hydra_config import HydraConfig
from hydra.core.plugins import Plugins
from hydra.core.singleton import Singleton
from hydra.core.utils import JobReturn, JobStatus, configure_log, filter_overrides, setup_globals
from hydra.plugins.launcher import Launcher
from hydra.types import HydraContext, TaskFunction
from omegaconf import DictConfig, OmegaConf, open_dict

from atommic.utils import logging


# monkey-patch hydra func
def is_in_toplevel_plugins_module(*args, **kwargs) -> bool:  # pylint: disable=unused-argument
    """Monkey-patch Hydra to allow for plugins to be loaded from any module."""
    return True


# Monkey-patch Hydra
Plugins.instance().is_in_toplevel_plugins_module = is_in_toplevel_plugins_module


@dataclass
class ProcessLauncherConfig:
    """Configuration for the ProcessLauncher plugin."""

    _target_: str = "atommic.core.utils.process_launcher.launcher.ProcessLauncher"
    num_gpus: int = -1
    jobs_per_gpu: int = 1


def execute_job(
    idx: int,
    overrides: Sequence[str],
    hydra_context: HydraContext,
    config: DictConfig,
    singleton_state: Dict[Any, Any],
    gpu_idx: int,
):
    """Creates a process that launches a "single" job that is identical in config + updated with sweep hyperparams.
    Since a different process is being used, CUDA can work in non-ddp mode without issue.
    Attempting ddp when using this script will not work as ddp cannot be used in shared contexts.

    Parameters
    ----------
        idx : int
            Global index of the job.
        overrides : list
            List of str overrides that correspond to this job.
        hydra_context : HydraContext
            Hydra Context used to load the sweep params into the global config.
        config : DictConfig
            Global config that will be updated with sweep hyper parameters.
        singleton_state : Dict[Any, Any]
            Hydra state.
        gpu_idx : int
            The GPU ID on which this process will be run.

    Returns
    -------
    tuple
        - The Process object that corresponds to this sweep.
        - The JobReturn object holding some metadata about this run.
    """
    # Required by Hydra (lookup other Hydra Launchers for details)
    setup_globals()
    Singleton.set_state(singleton_state)

    # Update base config with overrides to create sweep config
    sweep_config = hydra_context.config_loader.load_sweep_config(config, list(overrides))
    with open_dict(sweep_config):
        sweep_config.hydra.job.id = f"{sweep_config.hydra.job.name}_{idx}"
        sweep_config.hydra.job.num = idx
    HydraConfig.instance().set_config(sweep_config)

    # Set up a directory where the config will temporarily be stored.
    script_path = os.path.join(os.getcwd(), sys.argv[0])
    script_path = os.path.abspath(script_path)

    hash_salt = "|".join([script_path, str(OmegaConf.to_yaml(config))]).encode("utf-8")
    hash_val = hashlib.sha256(hash_salt).hexdigest()

    config_dir = os.path.join(os.getcwd(), "hydra_cfg", str(hash_val))
    if not os.path.exists(config_dir):
        os.makedirs(config_dir, exist_ok=True)

    task_cfg = copy.deepcopy(sweep_config)

    # Remove hydra from sweep config
    # This is done to prevent recursive call to multirun launcher in Hydra.
    with open_dict(task_cfg):
        task_cfg.pop("hydra", "")

    # Save the current jobs config to directory
    temp_config_name = f"config_{idx}.yaml"
    temp_config = os.path.join(config_dir, temp_config_name)
    OmegaConf.save(task_cfg, temp_config)

    # Compute the overides as a dict
    overrides = OmegaConf.to_container(config.hydra.overrides.task)

    # Check and replace trainer.devices in config with gpu_idx
    found_devices = False
    gpu_override = f"trainer.devices=[{gpu_idx}]"
    for oidx, val in enumerate(overrides):
        if "trainer.devices" in val:
            overrides[oidx] = gpu_override  # type: ignore
            found_devices = True

    if not found_devices:
        overrides.append(gpu_override)  # type: ignore

    # Build launch command
    # Note: We depend on PTL doing the right thing since this command has global visibility of all CUDA_VISIBLE_DEVICES
    cmd = [
        "python",
        script_path,
        "--config-path",
        config_dir,
        "--config-name",
        temp_config_name,
        *overrides,
    ]

    # Launch the subprocess; pipe the stderr
    # NOTE: If this hangs due to some reason after prolonged training, it means that the stderr pipe buffer
    # has become full at the OS level and we need to explicitly empty it (either parallel thread or manually
    # call proc.communicate(). It should not happen in general case as stderr is filled only in case retcode != 0
    # If it does happen though, implement the code here
    # https://stackoverflow.com/questions/39607172/python-subprocess-popen-poll-seems-to-hang-but-communicate-works
    proc = subprocess.Popen(cmd, stderr=subprocess.PIPE)  # pylint: disable=consider-using-with

    # Setup data thread for stderr
    std_error_buffer: List = []
    # Trivial thread just reads lines from stdout into the list
    drainerthread = threading.Thread(target=std_error_buffer.extend, args=(proc.stderr,))
    drainerthread.daemon = True
    drainerthread.start()

    # Construct JobReturn object for Hydra
    res = JobReturn()
    res.cfg = task_cfg
    res.overrides = overrides
    res.hydra_cfg = config
    res.working_dir = os.getcwd()
    res.return_value = None

    return proc, res, (std_error_buffer, drainerthread)


def launch(  # noqa: MC0001
    launcher, job_overrides: Sequence[Sequence[str]], initial_job_idx: int
) -> Sequence[JobReturn]:
    """Launches a list of jobs with the given config.

    Parameters
    ----------
        launcher : class
            Reference to the Launched subclass.
        job_overrides : list
            A List of List<String>, where each inner list is the arguments for one job run.
        initial_job_idx : int
            Initial job idx in batch.

    Returns
    -------
    list
        A list of JobReturn objects.
    """
    # Needed for Hydra, lookup JoblibLauncher in Hydra
    setup_globals()
    assert launcher.config is not None
    assert launcher.task_function is not None
    assert launcher.hydra_context is not None

    configure_log(launcher.config.hydra.hydra_logging, launcher.config.hydra.verbose)
    sweep_dir = Path(str(launcher.config.hydra.sweep.dir))
    sweep_dir.mkdir(parents=True, exist_ok=True)

    # Extract the runner's config (it is actually a DictConfig, but type is used for autocomplete)
    runner_cfg = launcher.runner  # type: ProcessLauncherConfig

    logging.info(
        f"ProcessLauncher({','.join([f'{k}={v}' for k, v in runner_cfg.items()])}) "  # type: ignore
        f"is launching {len(job_overrides)} jobs."
    )
    logging.info(f"Launching jobs, sweep output dir : {sweep_dir}")
    for idx, overrides in enumerate(job_overrides):
        logging.info(f"\t#{idx} : {' '.join(filter_overrides(overrides))}")

    # Needed by Hydra
    singleton_state = Singleton.get_state()

    # Process the runner's config to build up the multiplex config
    num_gpus = runner_cfg.get("num_gpus", -1)  # type: ignore
    jobs_per_gpu = runner_cfg.get("jobs_per_gpu", 1)  # type: ignore

    # Only GPUs are supported for now.
    if num_gpus <= 0:
        if torch.cuda.is_available():
            num_gpus = torch.cuda.device_count()
        else:
            raise ValueError(f"{launcher.__class__.__name__} only supports GPU operations.")

    # Setup arguments for multiplex runner
    overrides = list(job_overrides)  # type: ignore
    num_overrides = len(overrides)

    job_idx = 0
    batch_size = num_gpus * jobs_per_gpu
    gpu_idx = 0

    ret: List = []  # List of returned JobResult
    subprocess_list: List = []  # Buffer of subprocess
    results = []  # Buffer of JobResult

    # STD ERROR cache
    std_error_buffers = []  # type: List[List[str]]
    std_error_threads: threading.Thread = []  # type: ignore

    # Run over all job combinations
    while job_idx < num_overrides:
        # Fill up subprocess buffer while its size is smaller than multiplex batch size
        while len(subprocess_list) < batch_size:
            # If we run out of jobs, stop trying to submit more jobs
            if job_idx >= num_overrides:
                break

            # Submit a job as a new process
            process, res, error_tup = execute_job(
                initial_job_idx + job_idx,
                overrides[job_idx],
                launcher.hydra_context,
                launcher.config,
                singleton_state,
                gpu_idx % num_gpus,  # This will evenly distribute GPU load
            )

            # Store the subprocesses and JobResults
            subprocess_list.append(process)
            results.append(res)

            # Manage stderror thread data
            std_error_buffers.append(error_tup[0])
            std_error_threads.append(error_tup[1])  # type: ignore

            job_idx += 1
            gpu_idx += 1

        # Poll for samples in batch to finish.
        if len(subprocess_list) > 0:
            finished_processes = [0] * len(subprocess_list)

            # Check if all processes are completed or not
            # This is busy waiting, this is actually quite necessary
            # Turns out that when you do proc.communicate(), you block all other threads immediately.
            # IE they may fill up their buffers entirely, and hang while they wait for the first thread
            # who called communicate() to finish its work or crash.
            # Effectively it entirely stops multiprocessing jobs or multiplexed runs.
            # Must poll and busy wait to keep threads alive, along with drain the pipes with thread buffers.
            while sum(finished_processes) < len(subprocess_list):
                # Check all processes to make sure they have a retcode (doesn't matter yet if 0 or not)
                for proc_idx, proc in enumerate(subprocess_list):
                    # poll() is cheaper op than communicate()
                    retcode = proc.poll()

                    if retcode is not None:
                        # Log that the process with some ID has finished
                        if finished_processes[proc_idx] == 0:
                            logging.info(f"Processed job : {len(ret) + proc_idx} :: Ret code = {retcode}")

                        finished_processes[proc_idx] = 1

                        # Join this thread and merge its stderror buffer
                        proc.wait()
                        std_error_threads[proc_idx].join()  # type: ignore
                        error_data = std_error_buffers[proc_idx]
                        error_data = [
                            str(data, encoding="utf-8").encode("utf-8").decode("utf-8").encode("utf-8")  # type: ignore
                            for data in error_data
                        ]

                        std_error_buffers[proc_idx] = error_data

                time.sleep(1.0)

            # Process all the subprocess results
            for proc_idx, (proc, res) in enumerate(zip(subprocess_list, results)):
                # Wait until completion of process
                _ = proc.communicate()

                # 0 is for successful run
                if proc.returncode == 0:
                    res.status = JobStatus.COMPLETED
                else:
                    # > 0 is for error, log the error.
                    # Note: For the sake of efficiency while we log the error and raise an exception,
                    # It will only raise the 1st wrong job in all the jobs.
                    # If multiple jobs fail, it will still try to execute every job first before
                    # raising the error for the first one.
                    # This is done so that even if some jobs fail (say OOM or something),
                    # other jobs can still run.
                    err_buffer = std_error_buffers[proc_idx]
                    if isinstance(err_buffer, (list, tuple)):
                        err_string = ""
                        for err_line in err_buffer:
                            err_string = (
                                err_string
                                + f"{str(err_line, encoding='utf-8').encode('utf-8').decode('utf-8')}"  # type: ignore
                            )
                    else:
                        err_string = err_buffer

                    error_msg = (
                        f"\nHyperparameter Arguments : {proc.args}\n"
                        f"Process Return code : {proc.returncode}\n"
                        f"Error Trace :\n"
                        f"{err_string}"
                    )
                    res.return_value = Exception(error_msg)
                    res.status = JobStatus.FAILED

                logging.info(f"Finished executing job : {len(ret)}. Return Code = {proc.returncode}")
                ret.append(res)

            # Reset for next batch
            subprocess_list.clear()
            results.clear()

    return ret


class ProcessLauncher(Launcher):
    """Process Launcher

    Based on the JoblibLauncher, but uses processes to scatter jobs in a multiplexed manner across some number of GPUs
    on a single machine.
    """

    def __init__(self, **kwargs: Any) -> None:
        """Inits :class:`ProcessLauncher`."""
        self.config: Optional[DictConfig] = None
        self.task_function: Optional[TaskFunction] = None
        self.hydra_context: Optional[HydraContext] = None
        self.runner: ProcessLauncherConfig = kwargs  # type: ignore

    def setup(
        self,
        *,
        hydra_context: HydraContext,
        task_function: TaskFunction,
        config: DictConfig,
    ) -> None:
        """Setups the launcher.

        Parameters
        ----------
        hydra_context : HydraContext
            Hydra context object.
        task_function : TaskFunction
            Task function.
        config : DictConfig
            Launcher config.
        """
        self.config = config
        self.task_function = task_function
        self.hydra_context = hydra_context

    def launch(self, job_overrides: Sequence[Sequence[str]], initial_job_idx: int) -> Sequence[JobReturn]:
        """Launch jobs.

        Parameters
        ----------
        job_overrides : Sequence[Sequence[str]]
            List of List<String>, where each inner list is the arguments for one job run.
        initial_job_idx : int
            Initial job idx.

        Returns
        -------
        Sequence[JobReturn]
            Sequence of JobReturn objects, where each object represents the return value of one job.
        """
        return launch(launcher=self, job_overrides=job_overrides, initial_job_idx=initial_job_idx)


ConfigStore.instance().store(
    group="hydra/launcher",
    name="atommic_launcher",
    node=ProcessLauncherConfig,
    provider="atommic_process_launcher",
)

Plugins.instance(ProcessLauncher)
