"""
@package pmsco.schedule
Job schedule interface

This module defines common infrastructure to submit a PMSCO calculation job to a job scheduler such as Slurm.

The schedule can be defined as part of the run-file (see pmsco module).
Users may derive sub-classes in a separate module to adapt to their own computing cluster.

The basic call sequence is:
1. Create a schedule object.
2. Initialize its properties with job parameters.
3. Validate()
4. Submit()

@author Matthias Muntwiler, matthias.muntwiler@psi.ch

@copyright (c) 2015-23 by Paul Scherrer Institut @n
Licensed under the Apache License, Version 2.0 (the "License"); @n
  you may not use this file except in compliance with the License.
  You may obtain a copy of the License at
  http://www.apache.org/licenses/LICENSE-2.0
"""

import collections.abc
import datetime
import logging
import os
from pathlib import Path
import shutil
import subprocess
import sys
from typing import Any, Callable, Dict, Generator, Iterable, Iterator, List, Mapping, Optional, Sequence, Set, Tuple, Union

try:
    import commentjson as json
except ImportError:
    import json

import pmsco.config

logger = logging.getLogger(__name__)


class JobSchedule(pmsco.config.ConfigurableObject):
    """
    Base class for job schedule

    Usage:
    1. Create object, assigning a project instance.
    2. Set attributes.
    3. Assign run_dict.
    4. Call validate.
    5. Call submit.

    An example can be seen in pmsco.schedule_project.

    This class defines the abstract interface, common actions and some utilities.
    Derived classes may override any method, but should call the inherited method.
    """

    ## @var enabled (bool)
    #
    # This parameter signals whether pmsco should schedule a job or run the calculation.
    # It is not directly used by the schedule classes but by the pmsco module.
    # It must be defined in the run file and set to true to submit the job to a scheduler.
    # It is set to false in the run file copied to the job directory so that the job script starts the calculation.

    def __init__(self, project):
        super(JobSchedule, self).__init__()
        self.project = project

        # Specifies whether a job script is created (True) or not (False, default).
        self.enabled = False

        # Specifies whether the job script is submitted manually by the user (True, default)
        # or as a part of this class' execution (False).
        self.manual = True

        # Specifies whether an existing job directory is overwritten (True).
        # Otherwise, an exception is raised if it exists (False, default).
        self.overwrite_job_dir = False

        # Content of the run-file.
        # This must be set by the caller before the object is executed.
        self.run_dict = {}

        # Path to the destination run-file.
        # This value is set by self.validate.
        self.run_file = Path()

        # Job work directory.
        # This value is set by self.validate.
        self.job_dir = Path()

        # Dest path of the shell script to be submitted to the queue.
        # This value is set by self.validate.
        self.job_file = Path()

        # Directory that contains the pmsco code.
        # This value is set by self.validate.
        self.pmsco_root_dir = Path(__file__).parent.parent

        # Directory that contains the project module
        # This value is set by self.validate.
        self.project_dir = Path()

        # Project files to be copied to the job_dir.
        # Paths should be absolute or relative to project_dir.
        # The project module is appended by self.validate.
        # Other files have to be added by the caller.
        self.project_files = []

        # Name of the conda environment to activate (`source activate` command).
        # Either conda_env or virtual_env should be specified.
        # If both are empty, the environment is detected from the environment of the current process.
        self.conda_env = ""

        # Path of the virtual environment to activate (must contain the `activate` command).
        # Either conda_env or virtual_env should be specified.
        # If both are empty, the environment is detected from the environment of the current process.
        self.virtual_env = ""

    def validate(self):
        """
        validate the job parameters.

        make sure all object attributes are correct for submission.

        @return: None
        """

        assert self.run_dict, "run_dict not set"

        self.pmsco_root_dir = Path(self.project.directories['pmsco']).parent
        self.project_dir = Path(self.project.directories['project'])
        output_dir = Path(self.project.directories['output'])

        assert self.pmsco_root_dir.is_dir()
        assert (self.pmsco_root_dir / "pmsco").is_dir(), "can't find pmsco directory (source code)"
        assert self.project_dir.is_dir(), "can't find project directory"
        assert output_dir.is_dir(), "can't find output directory"
        assert output_dir.is_absolute(), "output directory must be an absolute path"
        assert self.project.job_name, "job_name is undefined"

        if output_dir.name == self.project.job_name:
            self.job_dir = output_dir
        else:
            self.job_dir = output_dir / self.project.job_name
        try:
            self.job_dir.mkdir(parents=True, exist_ok=self.overwrite_job_dir)
        except FileExistsError:
            logger.error("job directory exists - check job name or clean up manually")
            raise

        self.job_file = (self.job_dir / self.project.job_name).with_suffix(".sh")
        self.run_file = (self.job_dir / self.project.job_name).with_suffix(".json")

        self.project_files.append(sys.modules[self.project.__module__].__file__)

    def submit(self):
        """
        submit the job to the scheduler.

        as of this class, the method does to following:

        1. copy source files
        2. copy a patched version of the run file.
        3. write the job file (_write_job_file must be implemented by a derived class).

        @return: None
        """
        self._copy_source()
        self._fix_run_file()
        self._write_run_file()
        self._write_job_file()

    def _copy_source(self):
        """
        copy the source files to the job directory.

        the files to copy must be listed explicitly in the project_files attribute.
        the files are copied to the job_dir directory.
        the directory must exist.

        @return: None
        """

        files = set((self.project_dir.joinpath(pf) for pf in self.project_files))
        for f in files:
            shutil.copy2(f, self.job_dir)

    def _fix_run_file(self):
        """
        fix the run file.

        patch some entries of self.run_dict so that it can be used as run file.
        the following changes are made:
        1. set schedule.enabled to false so that the calculation is run.
        2. set the output directory to the job directory.
        3. set the log file to the job directory.

        @return: None
        """
        self.run_dict['schedule']['enabled'] = False
        self.run_dict['project']['directories']['output'] = str(self.job_dir)
        self.run_dict['project']['job_name'] = self.project.job_name
        self.run_dict['project']['log_file'] = str((self.job_dir / self.project.job_name).with_suffix(".log"))

    def _write_run_file(self):
        """
        copy the run file.

        this is a JSON dump of self.run_dict to the self.run_file file.

        @return: None
        """
        with open(self.run_file, "wt") as f:
            json.dump(self.run_dict, f, indent=2)

    def _write_job_file(self):
        """
        create the job script.

        this method must be implemented by a derived class.
        the script must be written to the self.job_file file.
        don't forget to make the file executable.

        @return: None
        """
        pass

    @staticmethod
    def detect_env() -> Dict[str, os.PathLike]:
        """
        detect the python environment

        determines the current python environment.

        examples:
        - /das/work/p17/p17274/conda/envs/pmsco310/bin
        - /home/user/envs/pmsco-uv

        @return: dictionary type -> path containing one or zero items.
            type is either 'conda', 'venv' or 'system';
            path is the bin directory containing python.
        """

        pp = Path(sys.executable).parent
        if (pp / 'activate').is_file():
            return {"venv": pp}
        for parent in pp.parents:
            if (parent / "condabin").is_dir():
                return {"conda": pp}
        else:
            return {"system": pp}


class SlurmSchedule(JobSchedule):
    """
    job schedule for a slurm scheduler.

    this class implements commonly used features of the slurm scheduler.
    host-specific features and the creation of the job file should be done in a derived class.
    derived classes must, in particular, implement the _write_job_file method.
    they can override other methods, too, but should call the inherited method first.

    1. copy the source trees (pmsco and projects) to the job directory
    2. copy a patched version of the run file.
    3. call the submission command

    the public properties of this class should be assigned from the run file.
    """
    def __init__(self, project):
        super(SlurmSchedule, self).__init__(project)
        self.host = ""
        self.nodes = 1
        self.tasks = 8
        self.wall_time = datetime.timedelta(hours=1)
        self.signal_time = 600

    @staticmethod
    def parse_timedelta(td):
        """
        parse time delta input formats

        converts a string or dictionary from run-file into datetime.timedelta.

        @param td:
            str: [days-]hours[:minutes[:seconds]]
            dict: days, hours, minutes, seconds - at least one needs to be defined. values must be numeric.
            datetime.timedelta - native type
        @return: datetime.timedelta
        """
        if isinstance(td, str):
            dt = {}
            d = td.split("-")
            if len(d) > 1:
                dt['days'] = float(d.pop(0))
            t = d[0].split(":")
            try:
                dt['hours'] = float(t.pop(0))
                dt['minutes'] = float(t.pop(0))
                dt['seconds'] = float(t.pop(0))
            except (IndexError, ValueError):
                pass
            td = datetime.timedelta(**dt)
        elif isinstance(td, collections.abc.Mapping):
            td = {k: float(v) for k, v in td.items()}
            td = datetime.timedelta(**td)
        return td

    @property
    def tasks_per_node(self):
        return self.tasks // self.nodes

    @tasks_per_node.setter
    def tasks_per_node(self, value):
        self.tasks = value * self.nodes

    def validate(self):
        super(SlurmSchedule, self).validate()
        self.wall_time = self.parse_timedelta(self.wall_time)

    def submit(self):
        """
        call the sbatch command

        if manual is true, the job files are generated but the job is not submitted.

        @return: None
        """
        super(SlurmSchedule, self).submit()
        args = ['sbatch', str(self.job_file)]
        print(" ".join(args))
        if self.manual:
            print("manual run - job files created but not submitted")
        else:
            cp = subprocess.run(args)
            cp.check_returncode()


class PsiRaSchedule(SlurmSchedule):
    """
    job shedule for the Ra cluster at PSI.

    this class selects specific features of the Ra cluster,
    such as the partition and node type.
    it also implements the _write_job_file method.

    for information about the Ra cluster, see
    https://www.psi.ch/en/photon-science-data-services/offline-computing-facility-for-sls-and-swissfel-data-analysis

    COMPUTE NODES

    | NodeName | Weight | CPUs | RealMemory | MemSpecLimit | Sockets | CoresPerSocket | ThreadsPerCore |
    | --- | :---: | :---: | :---: | :---: | :---: | :---: | :---: |
    | ra-c-[033-048] | 100 | 36 | 256000 | 21502 | 2 | 18 | 1 |
    | ra-c-[049-072] | 100 | 40 | 386000 | 21502 | 2 | 20 | 1 |
    | ra-c-[073-084] |  | 52 |  |  | 2 | 26 | 1 |
    | ra-c-[085-096] |  | 56 |  |  | 2 | 28 | 1 |

    PARTITIONS

    | PartitionName | Nodes | MaxTime | DefaultTime | Priority |
    | --- | :---: | :---: | :---: | :---: | :---: |
    | hour    |  | 0-01:00:00 | 0-01:00:00 | 4 |
    | day     |  | 1-00:00:00 | 0-08:00:00 | 2 |
    | week    |  | 8-00:00:00 | 2-00:00:00 | 1 |

    The option --ntasks-per-node is meant to be used with the --nodes option.
    (For the --ntasks option, the default is one task per node, use the --cpus-per-task option to change this default.)
    """

    ## @var partition (str)
    #
    # the partition is selected based on wall time and number of tasks by the validate() method.
    # it should not be listed in the run file.

    ## @var modules (list of str)
    #
    # names of the software modules to load (`module load` command)

    def __init__(self, project):
        super(PsiRaSchedule, self).__init__(project)
        self.partition = ""
        self.modules = []
        self.default_env = {}

    def validate(self):
        """
        check validity of parameters and detect the environment

        the validity is not checked in detail - just intercept some common or severe mistakes.

        detect whether we are in a conda or virtual environment
        so that we can add the necessary activation to the job script.

        detect which modules are active. depending on the python environment, LOADEDMODULES is:
            - venv: openssl/3.4.1:TclTk/8.6.16:xz/5.8.0:Python/3.12.9:gcc/9.3.0:libfabric/1.18.0:cuda/11.1.0:openmpi/4.0.5_slurm
            - conda: miniforge/2025-03-25:gcc/9.3.0:libfabric/1.18.0:cuda/11.1.0:openmpi/4.0.5_slurm

        other potentially useful but currently unused environment variables are:
            - PMODULES_LOADED_COMPILER: openmpi/4.0.5_slurm
            - PMODULES_LOADED_TOOLS: openssl/3.4.1:xz/5.8.0
            - PMODULES_LOADED_LIBRARIES: libfabric/1.18.0
            - PMODULES_LOADED_PROGRAMMING: TclTk/8.6.16:Python/3.12.9:gcc/9.3.0:cuda/11.1.0
            - PMODULES_LOADED_TOMCAT: miniforge/2025-03-25

        note: openssl, TclTk, xz and cuda were not requested explicitly.
        """

        super(PsiRaSchedule, self).validate()

        # check that the submission is sane - the values are not firm
        assert self.nodes <= 2
        assert self.tasks <= 64
        assert 30 * 60 <= self.wall_time.total_seconds() <= 8 * 24 * 60 * 60

        if self.wall_time.total_seconds() > 24 * 60 * 60:
            self.partition = "week"
        else:
            self.partition = "day"

        assert self.partition in ["day", "week"]

        self.default_env = self.detect_env()

        if len(self.modules) == 0:
            self.modules = os.environ["LOADEDMODULES"].split(":")

    def _write_job_file(self):
        """
        write a job file for the ra cluster

        @return: None
        """
        lines = []

        lines.append('#!/bin/bash')
        lines.append('#SBATCH --export=NONE')
        lines.append(f'#SBATCH --job-name="{self.project.job_name}"')
        lines.append(f'#SBATCH --partition={self.partition}')
        lines.append(f'#SBATCH --time={int(self.wall_time.total_seconds() / 60)}')
        lines.append(f'#SBATCH --nodes={self.nodes}')
        lines.append(f'#SBATCH --ntasks-per-node={self.tasks_per_node}')
        # 0 - 65535 seconds
        # currently, PMSCO does not react to signals properly
        # lines.append(f'#SBATCH --signal=TERM@{self.signal_time}')
        lines.append(f'#SBATCH --output="{self.project.job_name}.o.%j"')
        lines.append(f'#SBATCH --error="{self.project.job_name}.e.%j"')

        # environment
        if self.modules:
            lines.append('module use unstable || true')
            lines.append('module use Libraries || true')
        for module in self.modules:
            lines.append(f'module load {module}')

        conda_env = self.conda_env or self.default_env.get("conda")
        virtual_env = self.virtual_env or self.default_env.get("venv")
        if conda_env:
            lines.append(f'source activate {conda_env}')
        elif virtual_env:
            activate_script = Path(virtual_env) / "activate"
            if activate_script.is_file():
                lines.append(f'source {activate_script}')

        lines.append('env')
        lines.append('')

        # run
        lines.append(f'cd "{self.job_dir}"')
        lines.append(f'mpirun python -m pmsco -r {self.run_file}')
        lines.append('')

        # clean up
        lines.append(f'cd "{self.job_dir}"')
        lines.append('exit 0')

        self.job_file.write_text("\n".join(lines))
        self.job_file.chmod(0o755)
