# Copyright lowRISC contributors (OpenTitan project).
# Licensed under the Apache License, Version 2.0, see LICENSE for details.
# SPDX-License-Identifier: Apache-2.0

# ----------------------------------
# SGE.py
#       _JobData Class
# ----------------------------------
import os
import pwd
import re
import subprocess
import time
from collections.abc import Generator, Sequence
from pathlib import Path

from dvsim.logging import log


class _JobData:
    """Internal helper class to manage job data from qstat."""

    def __init__(self, qstat_job_line: str) -> None:
        """Initialise a JobData object for qstat job status parsing.

        Args:
            qstat_job_line: string from qstat call

        """
        # The format of the line goes like
        # job-ID prior name  user  state submit/start at queue slots ja-task-ID

        tokens = qstat_job_line.split()

        assert len(tokens) >= 8, "Not a valid qstat line: " + qstat_job_line
        # Job line must have at least 8 tokens
        try:
            self.id = int(tokens[0])

        except ValueError:
            error_msg = "Could not convert data to an integer."
            raise ValueError(error_msg)

        self.priority = float(tokens[1])
        self.name = tokens[2]
        self.user = tokens[3]
        self.state = tokens[4]
        self.time = " ".join(tokens[5:7])

        if "@" in qstat_job_line:
            # Has queue assigned, e.g. core2.q@q2.caspian.mit.edu
            self.queue = tokens[7]
            self.slots = tokens[8]
            if len(tokens) == 9:
                self.ja_task_id = None
            elif len(tokens) == 10:
                self.ja_task_id = tokens[-1]
            else:
                msg = f"Could not parse line: {qstat_job_line}"
                raise ValueError(msg)
        else:
            # No queue assigned
            self.slots = tokens[7]
            if len(tokens) == 8:
                self.ja_task_id = None
            elif len(tokens) == 9:
                self.ja_task_id = tokens[-1]
            else:
                msg = f"Could not parse line: {qstat_job_line}"
                raise ValueError(msg)

        # Convert array indices ja_task_id into python list format
        ja_task_id = []
        if self.ja_task_id is not None:
            for blob in self.ja_task_id.split(","):
                # Parse data of form '193-349:1' or '1934'
                x = blob.split(":")
                if len(x) == 1:
                    ja_task_id += x
                else:
                    subjobs, step = x
                    begin, end = subjobs.split("-")
                    ja_task_id += range(int(begin), int(end) + 1, int(step))

        self._ja_tasklist = ja_task_id

    def __repr__(self) -> str:
        repr_str = ["{"]
        for key, value in self.__dict__.items():
            if key[0] != "_":
                repr_str.append(key + "=" + str(value))
        repr_str.append("}")
        return "\n".join(repr_str)


class JobList:
    """Internal helper class to manage job lists."""

    def __init__(self, qstat_output: str) -> None:
        """Initialise JobList."""
        self._joblist = []
        for line in qstat_output.split("\n")[2:-1]:
            self._joblist.append(_JobData(line))

    def __iter__(self) -> Generator[_JobData]:
        """Get an Iterator over the job list."""
        yield from self._joblist

    def __repr__(self) -> str:
        """Get string representation of the job list."""
        return "\n".join([str(job) for job in self._joblist])


class SGE:
    """External system call handler for Sun Grid Engine environment."""

    def __init__(self, q=None, path: str | Path = "") -> None:
        """Initialise Sun Grid Engine proxy."""
        path = Path(path)
        if q is None:
            # No queue specified. By default, submit to all available queues.
            self.cmd_qconf = path / "qconf"

            try:
                qliststr = _exec(f"{self.cmd_qconf} -sql")
            except OSError:
                error_msg = "Error querying queue configuration"
                log.exception(error_msg)
                raise OSError(error_msg)

            self.q = qliststr.replace("\n", ",")[:-1]

            log.info(
                "Sun Grid Engine handler initialized\nQueues detected: %s",
                self.q,
            )

        else:
            self.q = q

        self.cmd_qsub = path / "qsub"
        self.cmd_qstat = path / "qstat"

    def wait(
        self,
        jobid: int,
        interval: int = 10,
        name: str | None = None,
        pbar: None = None,
    ) -> None:
        """Wait for job running on SGE Grid Engine environment to finish.

        If you are just waiting for one job, this becomes a dumb substitute for
        the -sync y option which can be specified to qsub.

        Args:
            jobid: job identifier
            interval: polling interval of SGE queue, in seconds. (Default: 10)
            name: name of the job
            pbar: progress bar (not implemented?)

        """
        dowait = True
        while dowait:
            p = subprocess.Popen(  # noqa: S602
                self.cmd_qstat,
                shell=True,
                stdout=subprocess.PIPE,
            )
            pout, _ = p.communicate()

            if pbar is not None:
                log.error("Progress bar handling not implemented")

            dowait = False
            for line in pout.split("\n"):
                t = line.split()
                if len(t) >= 5 and t[0] == str(jobid):
                    # Find a line with useful info
                    if re.search(t[4], "qwrt"):
                        # Job must be queued, running or being transferred
                        dowait = True
                        break
                    if re.search(t[4], "acuE"):
                        # Job or host in error state
                        log.warning("Job %d in error state", str(jobid))

            if dowait:
                time.sleep(interval)
                if name is None:
                    log.info("Time %s: waiting for jobid %s to finish", time.ctime(), str(jobid))
                else:
                    log.info(
                        "Time %s: waiting for job '%s' (jobid %s) to \
finish",
                        time.ctime(),
                        name,
                        str(jobid),
                    )

    def submit(
        self,
        job,
        array=False,
        useenvironment=True,
        usecwd=True,
        name=None,
        stdin=None,
        stdout=None,
        stderr=None,
        joinstdouterr=True,
        nproc=1,
        wait=True,
        lammpi=True,
    ) -> int:
        """Submit a job to SGE.

        Returns:
            jobid as a number.

        """
        log.info(
            "Submitting job:   "
            + str(job)
            + " stdout: %s \
Stderr: %s",
            stdout,
            stderr,
        )

        # Parameters to qsub specified as the header of the job specified on
        # STDIN
        lamstring = (lammpi and f" -pe lammpi {nproc}") or ""
        qsuboptslist = ["-cwd -V ", lamstring]

        if name is not None:
            qsuboptslist.append("-N " + name)

        if stdin is not None:
            qsuboptslist.append("-i " + stdin)

        if stdout is not None:
            qsuboptslist.append("-o " + stdout)

        if stderr is not None:
            qsuboptslist.append("-e " + stderr)

        if joinstdouterr:
            qsuboptslist.append("-j")

        if wait:
            qsuboptslist.append("-sync y")

        if usecwd:
            qsuboptslist.append("-cwd")

        if useenvironment:
            qsuboptslist.append("-V")

        if array is not False:
            try:
                n = int(array[0])
            except IndexError:
                n = int(array)
                msg = "List is empty!"
                raise IndexError(msg)
            except ValueError:
                error_msg = "array[0] being an out of bounds access."
                log.exception(error_msg)
                raise ValueError(error_msg)

            try:
                m = int(array[1])
            except ValueError:
                m = None
                msg = "Could not convert data to an integer."
                raise ValueError(msg)
            except IndexError:
                m = None
                msg = "array[1] being an out of bounds access."
                raise IndexError(msg)

            try:
                s = int(array[2])
            except IndexError:
                s = None
                msg = "array[2] being an out of bounds access."
                raise IndexError(msg)
            except ValueError:
                s = None
                msg = "Could not convert data to an integer."
                raise ValueError(msg)

            if m == s is None:
                qsuboptslist.append(f"-t {n}")
            elif s is None:
                qsuboptslist.append(f"-t {n}-{m}")
            else:
                qsuboptslist.append(f"-t {n}-{m}:{s}")

        qsubopts = " ".join(qsuboptslist)

        pout = _exec(
            command=str(self.cmd_qsub),
            stdin=qsubopts + "\n" + job,
            print_command=False,
        )

        try:
            # Next to last line should be
            # "Your job 1389 (name) has been submitted"
            # parse for job id
            return int(pout.split("\n")[-2].split()[2])

        # except (ValueErrorValueError, IndexError, AttributeError) (e):
        except (ValueError, IndexError, AttributeError):
            error_msg = f"""Error submitting SGE job:
{qsubopts}
{job}

Output was:
{pout}"""
            log.exception(error_msg)
            raise OSError(error_msg)

    def getuserjobs(self, user=pwd.getpwuid(os.getuid())[0]) -> Sequence[_JobData]:
        """Return a list of SGE jobids run by a specific user.

        Args:
            user: SGE user to poll (Default = '', i.e. current user)
            qstat: path to qstat binary (Default = 'qstat')

        """
        p = subprocess.Popen(  # noqa: S602
            f"{self.cmd_qstat} -u {user}",
            shell=True,
            stdout=subprocess.PIPE,
        )
        qstat_output, _ = p.communicate()
        joblist = JobList(qstat_output)
        return [job for job in joblist if job.user == user]

    def run_job(
        self,
        command: str,
        name: str = "default",
        log_filename: str = "default.log",
        *,
        wait: bool = True,
    ) -> int:
        """Run job on SGE with piped logging."""
        return self.submit(
            command,
            name=name,
            stdout=log_filename,
            stderr=log_filename,
            wait=wait,
        )

    def get_queue_instance_status(self):
        """Get loads for each queue instance."""
        output = _exec(f"{self.cmd_qstat} -f")

        data = []
        for line in output.split("\n")[1:]:
            t = line.split()
            if len(t) != 5:
                continue

            nodename = t[0].split("@")[1].split(".")[0]
            maxslots = int(t[2].split("/")[2])
            load = float(t[3])

            data.append({"name": nodename, "maxslots": maxslots, "load": load})

        return data


def _exec(
    command: str,
    log_path: str | None = None,
    stdin: str = "",
    *,
    print_command: bool = False,
) -> bytes:
    """Run command line using subprocess, optionally returning stdout."""

    def _call_cmd(command: str, stdin: str = "") -> bytes:
        p = subprocess.Popen(  # noqa: S602
            command,
            shell=True,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
        )
        output, _ = p.communicate(stdin)
        return output

    if print_command:
        log.info("Executing process: \x1b[1;92m%-50s\x1b[0m Logfile: %s", command, log_path)

    output = ""
    if log_path is not None:
        try:
            with Path(log_path).open("a") as f:
                if print_command:
                    pass

                output = _call_cmd(command, stdin)
                f.write(output)

        except OSError:
            error_msg = "Error: File: " + str(log_path) + " does not appear to exist."
            log.exception(error_msg)
            raise OSError(error_msg)

    else:
        output = _call_cmd(command, stdin)

    log.info("Output of command is:\n%s", output)

    return output
