################################################################################
# Copyright (c), The AiiDA team. All rights reserved.                          #
# This file is part of the AiiDA code.                                         #
#                                                                              #
# The code is hosted on GitHub at https://github.com/aiidateam/aiida-firecrest #
# For further information on the license, see the LICENSE.txt file             #
# For further information please visit http://www.aiida.net                    #
################################################################################
"""Scheduler interface."""

from __future__ import annotations

import datetime
import re
import string
from typing import TYPE_CHECKING, Any, ClassVar

from aiida.engine.processes.exit_code import ExitCode
from aiida.schedulers import Scheduler, SchedulerError
from aiida.schedulers.datastructures import JobInfo, JobState, JobTemplate
from aiida.schedulers.plugins.slurm import _TIME_REGEXP, SlurmJobResource
from firecrest.FirecrestException import FirecrestException

from aiida_firecrest.utils import FcPath, JobNotFoundError, convert_header_exceptions

if TYPE_CHECKING:
    from aiida_firecrest.transport import FirecrestTransport


class FirecrestScheduler(Scheduler):  # type: ignore[misc]
    """Scheduler interface for FirecREST.
    Must be used together with the 'firecrest' transport plugin."""

    transport: FirecrestTransport
    _job_resource_class = SlurmJobResource
    _features: ClassVar[dict[str, Any]] = {
        "can_query_by_user": False,
    }
    _logger = Scheduler._logger.getChild("firecrest")

    def _get_submit_script_header(self, job_tmpl: JobTemplate) -> str:
        """
        Return the submit script header, using the parameters from the
        job_tmpl.

        :params job_tmpl: an JobTemplate instance with relevant parameters set.

        TODO: truncate the title if too long
        """
        # adapted from Slurm plugin
        lines = []
        if job_tmpl.submit_as_hold:
            lines.append("#SBATCH -H")

        if job_tmpl.rerunnable:
            lines.append("#SBATCH --requeue")
        else:
            lines.append("#SBATCH --no-requeue")

        if job_tmpl.email:
            # If not specified, but email events are set, SLURM
            # sends the mail to the job owner by default
            lines.append(f"#SBATCH --mail-user={job_tmpl.email}")

        if job_tmpl.email_on_started:
            lines.append("#SBATCH --mail-type=BEGIN")
        if job_tmpl.email_on_terminated:
            lines.append("#SBATCH --mail-type=FAIL")
            lines.append("#SBATCH --mail-type=END")

        if job_tmpl.job_name:
            # The man page does not specify any specific limitation
            # on the job name.
            # Just to be sure, I remove unwanted characters, and I
            # trim it to length 128

            # I leave only letters, numbers, dots, dashes and underscores
            # Note: I don't compile the regexp, I am going to use it only once
            job_title = re.sub(r"[^a-zA-Z0-9_.-]+", "", job_tmpl.job_name)

            # prepend a 'j' (for 'job') before the string if the string
            # is now empty or does not start with a valid charachter
            if not job_title or (
                job_title[0] not in string.ascii_letters + string.digits
            ):
                job_title = f"j{job_title}"

            # Truncate to the first 128 characters
            # Nothing is done if the string is shorter.
            job_title = job_title[:128]

            lines.append(f'#SBATCH --job-name="{job_title}"')

        if job_tmpl.import_sys_environment:
            lines.append("#SBATCH --get-user-env")

        if job_tmpl.sched_output_path:
            lines.append(f"#SBATCH --output={job_tmpl.sched_output_path}")

        if job_tmpl.sched_join_files:
            # TODO: manual says:  # pylint: disable=fixme
            # By  default both standard output and standard error are directed
            # to a file of the name "slurm-%j.out", where the "%j" is replaced
            # with  the  job  allocation  number.
            # See that this automatic redirection works also if
            # I specify a different --output file
            if job_tmpl.sched_error_path:
                self.logger.info(
                    "sched_join_files is True, but sched_error_path is set in "
                    "SLURM script; ignoring sched_error_path"
                )
        else:
            if job_tmpl.sched_error_path:
                lines.append(f"#SBATCH --error={job_tmpl.sched_error_path}")
            else:
                # To avoid automatic join of files
                lines.append("#SBATCH --error=slurm-%j.err")

        if job_tmpl.queue_name:
            lines.append(f"#SBATCH --partition={job_tmpl.queue_name}")

        if job_tmpl.account:
            lines.append(f"#SBATCH --account={job_tmpl.account}")

        if job_tmpl.qos:
            lines.append(f"#SBATCH --qos={job_tmpl.qos}")

        if job_tmpl.priority:
            #  Run the job with an adjusted scheduling priority  within  SLURM.
            #  With no adjustment value the scheduling priority is decreased by
            #  100. The adjustment range is from -10000 (highest  priority)  to
            #  10000  (lowest  priority).
            lines.append(f"#SBATCH --nice={job_tmpl.priority}")

        if not job_tmpl.job_resource:
            raise ValueError(
                "Job resources (as the num_machines) are required for the SLURM scheduler plugin"
            )

        lines.append(f"#SBATCH --nodes={job_tmpl.job_resource.num_machines}")
        if job_tmpl.job_resource.num_mpiprocs_per_machine:
            lines.append(
                f"#SBATCH --ntasks-per-node={job_tmpl.job_resource.num_mpiprocs_per_machine}"
            )

        if job_tmpl.job_resource.num_cores_per_mpiproc:
            lines.append(
                f"#SBATCH --cpus-per-task={job_tmpl.job_resource.num_cores_per_mpiproc}"
            )

        if job_tmpl.max_wallclock_seconds is not None:
            try:
                tot_secs = int(job_tmpl.max_wallclock_seconds)
                if tot_secs <= 0:
                    raise ValueError
            except ValueError:
                raise ValueError(
                    "max_wallclock_seconds must be "
                    "a positive integer (in seconds)! It is instead '{}'"
                    "".format(job_tmpl.max_wallclock_seconds)
                ) from None
            days = tot_secs // 86400
            tot_hours = tot_secs % 86400
            hours = tot_hours // 3600
            tot_minutes = tot_hours % 3600
            minutes = tot_minutes // 60
            seconds = tot_minutes % 60
            if days == 0:
                lines.append(f"#SBATCH --time={hours:02d}:{minutes:02d}:{seconds:02d}")
            else:
                lines.append(
                    f"#SBATCH --time={days:d}-{hours:02d}:{minutes:02d}:{seconds:02d}"
                )

        # It is the memory per node, not per cpu!
        if job_tmpl.max_memory_kb:
            try:
                physical_memory_kb = int(job_tmpl.max_memory_kb)
                if physical_memory_kb <= 0:
                    raise ValueError
            except ValueError:
                raise ValueError(
                    "max_memory_kb must be a positive integer (in kB)! "
                    f"It is instead `{job_tmpl.max_memory_kb}`"
                ) from None
            # --mem: Specify the real memory required per node in MegaBytes.
            # --mem and  --mem-per-cpu  are  mutually exclusive.
            lines.append(f"#SBATCH --mem={physical_memory_kb // 1024}")

        if job_tmpl.custom_scheduler_commands:
            lines.append(job_tmpl.custom_scheduler_commands)

        return "\n".join(lines)

    def submit_job(self, working_directory: str, filename: str) -> str | ExitCode:
        transport = self.transport
        with convert_header_exceptions():
            try:
                result = transport.blocking_client.submit(
                    system_name=transport._machine,
                    working_dir=working_directory,
                    script_remote_path=str(
                        FcPath(working_directory).joinpath(filename)
                    ),
                )
            except FirecrestException as exc:
                raise SchedulerError(str(exc)) from exc
        return str(result["jobId"])

    def get_jobs(
        self,
        jobs: list[str] | None = None,
        user: str | None = None,
        as_dict: bool = False,
    ) -> list[JobInfo] | dict[str, JobInfo]:
        """
        Retrieve the information of the jobs from the scheduler.
        :param jobs: a list of job ids to retrieve. If None, all jobs of the user are retrieved.
        :param user: the user to retrieve the jobs for. Note: this parameter is not used in Firecrest,
            as it does not support querying by user. We just ignore it.
            It's part of the Scheduler inerface, so we cannot remove it.
        :param as_dict: if True, return a dictionary with job ids as keys and JobInfo objects as values.

        :return: a list of JobInfo objects, or a dictionary with job ids as keys and JobInfo objects as values.
        :raises SchedulerError: if there is an error retrieving the jobs from the scheduler.
        """
        results = []
        transport = self.transport

        def _send_request_and_handle_errors(
            job_id: str | None = None,
        ) -> list[Any]:
            """Send a request to the Firecrest server and handle errors."""
            try:
                with convert_header_exceptions():
                    return transport.blocking_client.job_info(  # type: ignore[no-any-return]
                        transport._machine, job_id
                    )
            except FirecrestException as exc:
                raise SchedulerError(
                    f"Error retrieving job: {job_id} from the Firecrest server: {exc}"
                ) from exc
            except JobNotFoundError:
                # If the job is not found, we just skip it
                # this is to be consistent with aiida-firecrest-v1
                return []

        # If the job is completed, while aiida expect a silent return
        if jobs:
            for job in jobs:
                # This loop is a way to battle a delay in firecrest
                # when reading the recently submitted jobs from slurm.
                # See some discussions in https://github.com/aiidateam/aiida-firecrest/issues/103
                # This is the recommended way by CSCS to ensure the factuality of a response.
                for _ in range(5):
                    response = _send_request_and_handle_errors(str(job))
                    if response:
                        results += response
                        break
        else:
            for _ in range(5):
                results = _send_request_and_handle_errors()
                if results:
                    break
            results = _send_request_and_handle_errors()
            if not results:
                return {} if as_dict else []

        job_list = []
        for raw_result in results:
            job_state_raw = raw_result["status"]["state"]
            # sometimes FirecREST returns extra comments after the state, e.g., 'CANCELLED by 23570'
            job_state_cooked = job_state_raw.split()[0]
            job_owner = raw_result["user"]

            # if the job is completed or cancelled, we skip it
            if job_state_cooked in ["COMPLETED", "CANCELLED", "FAILED", "TIMEOUT"]:
                self.logger.debug(
                    f"Skipping from `func::get_job()`! job {raw_result['jobId']} with state {job_state_cooked} "
                )
                continue

            this_job = JobInfo()

            this_job.job_owner = job_owner
            this_job.job_id = str(raw_result["jobId"])
            # TODO: firecrest does not return the annotation, so set to an empty string.
            # To be investigated how important that is.
            this_job.annotation = ""

            try:
                job_state_string = _MAP_STATUS_SLURM[job_state_cooked]
            except KeyError:
                self.logger.warning(
                    f"Unrecognized job_state '{job_state_cooked}' for job id {this_job.job_id}"
                )
                job_state_string = JobState.UNDETERMINED
            # QUEUED_HELD states are not specific states in SLURM;
            # they are instead set with state QUEUED, and then the
            # annotation tells if the job is held.
            # I check for 'Dependency', 'JobHeldUser',
            # 'JobHeldAdmin', 'BeginTime'.
            # Other states should not bring the job in QUEUED_HELD, I believe
            # (the man page of slurm seems to be incomplete, for instance
            # JobHeld* are not reported there; I also checked at the source code
            # of slurm 2.6 on github (https://github.com/SchedMD/slurm),
            # file slurm/src/common/slurm_protocol_defs.c,
            # and these seem all the states to be taken into account for the
            # QUEUED_HELD status).
            # There are actually a few others, like possible
            # failures, or partition-related reasons, but for the moment I
            # leave them in the QUEUED state.
            if job_state_string == JobState.QUEUED and this_job.annotation in [
                "Dependency",
                "JobHeldUser",
                "JobHeldAdmin",
                "BeginTime",
            ]:
                job_state_string = JobState.QUEUED_HELD

            this_job.job_state = job_state_string

            # The rest is optional
            try:
                this_job.num_machines = int(raw_result["allocationNodes"])
            except ValueError:
                self.logger.warning(
                    "The number of allocated nodes is not "
                    "an integer ({}) for job id {}!".format(
                        raw_result["allocationNodes"], this_job.job_id
                    )
                )

            # See issue https://github.com/aiidateam/aiida-firecrest/issues/39
            # try:
            #     this_job.num_mpiprocs = int(thisjob_dict['number_cpus'])
            # except ValueError:
            #     self.logger.warning(
            #         'The number of allocated cores is not '
            #         'an integer ({}) for job id {}!'.format(
            # thisjob_dict['number_cpus'], this_job.job_id)
            #     )

            # ALLOCATED NODES HERE
            # string may be in the format
            # nid00[684-685,722-723,748-749,958-959]
            # therefore it requires some parsing, that is unnecessary now.
            # I just store is as a raw string for the moment, and I leave
            # this_job.allocated_machines undefined
            if this_job.job_state == JobState.RUNNING:
                this_job.allocated_machines_raw = raw_result["allocationNodes"]

            this_job.queue_name = raw_result["partition"]

            try:
                # time_limit = self._convert_time(raw_result["time"]["limit"])
                # note: v2 already returns in integer format, no need to convert
                time_limit = raw_result["time"]["limit"]

                if time_limit is None:
                    this_job.requested_wallclock_time_seconds = 0
                else:
                    this_job.requested_wallclock_time_seconds = time_limit
            except ValueError:
                self.logger.warning(
                    f"Couldn't parse the time limit for job id {this_job.job_id}"
                )

            # Only if it is RUNNING; otherwise it is not meaningful,
            # and may be not set (in my test, it is set to zero)
            if this_job.job_state == JobState.RUNNING:
                try:
                    wallclock_time_seconds = raw_result["time"]["elapsed"]
                    if wallclock_time_seconds is not None:
                        this_job.wallclock_time_seconds = wallclock_time_seconds
                    else:
                        this_job.wallclock_time_seconds = 0
                except ValueError:
                    self.logger.warning(
                        f"Couldn't parse time_used for job id {this_job.job_id}"
                    )

            # submission_time is not returned explicitly by the FirecREST server
            # see: https://github.com/aiidateam/aiida-firecrest/issues/40

            if raw_result["time"]["start"] != 0:
                try:
                    this_job.dispatch_time = self._parse_time_string(
                        raw_result["time"]["start"]
                    )
                except ValueError:
                    self.logger.warning(
                        f"Couldn't parse dispatch_time for job id {this_job.job_id}"
                    )

            if raw_result["time"]["end"] != 0:
                try:
                    this_job.finish_time = self._parse_time_string(
                        raw_result["time"]["end"]
                    )
                except ValueError:
                    self.logger.warning(
                        f"Couldn't parse finish_time for job id {this_job.job_id}"
                    )

            this_job.title = raw_result["name"]

            # Everything goes here anyway for debugging purposes
            this_job.raw_data = raw_result

            # Double check of redundant info
            # Not really useful now, allocated_machines in this
            # version of the plugin is never set
            if (
                this_job.allocated_machines is not None
                and this_job.num_machines is not None
                and len(this_job.allocated_machines) != this_job.num_machines
            ):
                self.logger.error(
                    "The length of the list of allocated "
                    "nodes ({}) is different from the "
                    "expected number of nodes ({})!".format(
                        len(this_job.allocated_machines), this_job.num_machines
                    )
                )

            # I append to the list of jobs to return
            job_list.append(this_job)

        if as_dict:
            return {job.job_id: job for job in job_list}

        return job_list

    def kill_job(self, jobid: str) -> bool:
        transport = self.transport
        try:
            with convert_header_exceptions():
                transport.blocking_client.cancel_job(transport._machine, jobid)
        except FirecrestException as exc:
            self.logger.debug(f"Unable to kill job {jobid}, the error was {exc}")
            return False
        return True

    def _convert_time(self, string: str) -> int | None:
        """
        Note: this function was copied from the Slurm scheduler plugin
        Convert a string in the format DD-HH:MM:SS to a number of seconds.
        """
        if string == "UNLIMITED":
            return 2147483647  # == 2**31 - 1, largest 32-bit signed integer (68 years)

        if string == "NOT_SET" or string == "N/A":
            return None

        groups = _TIME_REGEXP.match(string)
        if groups is None:
            raise ValueError("Unrecognized format for time string.")

        groupdict = groups.groupdict()
        # should not raise a ValueError, they all match digits only
        days = int(groupdict["days"] if groupdict["days"] is not None else 0)
        hours = int(groupdict["hours"] if groupdict["hours"] is not None else 0)
        mins = int(groupdict["minutes"] if groupdict["minutes"] is not None else 0)
        secs = int(groupdict["seconds"] if groupdict["seconds"] is not None else 0)

        return days * 86400 + hours * 3600 + mins * 60 + secs

    def _parse_time_string(
        self, epoch: int, fmt: str = "Unix time"
    ) -> datetime.datetime:
        """
        Parse an integer time in the epoch format (seconds since 1970-01-01 00:00:00 UTC)
        returns a datetime object.
        """

        try:
            return datetime.datetime.fromtimestamp(epoch)
        except Exception as exc:
            self.logger.debug(
                f"Unable to parse epoch time {epoch}, the message was {exc}"
            )
            raise ValueError("Problem parsing the epoch time.") from exc


# see https://slurm.schedmd.com/squeue.html#lbAG
# note firecrest returns full names, not abbreviations
_MAP_STATUS_SLURM = {
    "CA": JobState.DONE,
    "CANCELLED": JobState.DONE,
    "CD": JobState.DONE,
    "COMPLETED": JobState.DONE,
    "CF": JobState.QUEUED,
    "CONFIGURING": JobState.QUEUED,
    "CG": JobState.RUNNING,
    "COMPLETING": JobState.RUNNING,
    "F": JobState.DONE,
    "FAILED": JobState.DONE,
    "NF": JobState.DONE,
    "NODE_FAIL": JobState.DONE,
    "PD": JobState.QUEUED,
    "PENDING": JobState.QUEUED,
    "PR": JobState.DONE,
    "PREEMPTED": JobState.DONE,
    "R": JobState.RUNNING,
    "RUNNING": JobState.RUNNING,
    "S": JobState.SUSPENDED,
    "SUSPENDED": JobState.SUSPENDED,
    "TO": JobState.DONE,
    "TIMEOUT": JobState.DONE,
}
