# Copyright NTESS. See COPYRIGHT file for details.
#
# SPDX-License-Identifier: MIT

import importlib.resources
import logging
import math
import multiprocessing
import multiprocessing.synchronize
import os
import time
from concurrent.futures import CancelledError
from datetime import timedelta
from typing import Any

import flux  # type: ignore
from flux import Flux  # type: ignore
from flux.job import FluxExecutor  # type: ignore
from flux.job import FluxExecutorFuture  # type: ignore
from flux.job import Jobspec  # type: ignore
from flux.job import JobspecV1  # type: ignore

from hpc_connect.config import Config
from hpc_connect.config import ConfigScope
from hpc_connect.submit import HPCProcess
from hpc_connect.submit import HPCSubmissionFailedError
from hpc_connect.submit import HPCSubmissionManager
from hpc_connect.util import time_in_seconds

from .discover import read_resource_info

logger = logging.getLogger("hpc_connect")


class FluxProcess(HPCProcess):
    JOB_TIMEOUT_CODE = 66

    def __init__(self, name: str, future: FluxExecutorFuture) -> None:
        self.fh = Flux()
        self.name = name
        self.fut: FluxExecutorFuture = future
        self._jobid: int | None = None
        self._rc: int | None = None

        def set_returncode(fut: FluxExecutorFuture):
            try:
                info = flux.job.result(self.fh, fut.jobid())
                self.returncode = info.returncode
            except (CancelledError, Exception):
                self.returncode = 1

        def set_jobid(fut: FluxExecutorFuture):
            try:
                self.jobid = fut.jobid()
                logger.debug(f"submitted job {self.jobid} for {self.name}")
            except CancelledError:
                self.returncode = 1
            except Exception as e:
                raise HPCSubmissionFailedError(e)

        self.fut.add_jobid_callback(set_jobid)
        self.fut.add_done_callback(set_returncode)

    @property
    def jobid(self) -> int | None:
        return self._jobid

    @jobid.setter
    def jobid(self, arg: int) -> None:
        self._jobid = arg

    @property
    def returncode(self) -> int | None:
        return self._rc

    @returncode.setter
    def returncode(self, arg: int) -> None:
        self._rc = arg

    def poll(self) -> int | None:
        return self.returncode

    def cancel(self) -> None:
        logger.warning(f"Canceling flux job {self.jobid}")
        try:
            flux.job.cancel(self.fh, self.jobid)
        except OSError:
            logger.debug(f"Job {self.jobid} is inactive, cannot cancel")
        except Exception:
            logger.error(f"Failed to cancel job {self.jobid}")
        self.returncode = 1


class FluxMultiProcess(HPCProcess):
    def __init__(
        self,
        lock: multiprocessing.synchronize.RLock,
        procs: list[FluxProcess] | None = None,
    ) -> None:
        self.lock = lock
        self.procs = procs or []

    @property
    def returncode(self) -> int | None:
        rcs = [p.returncode for p in self.procs if p is not None]
        if not rcs:
            return None
        return max(rcs)  # type: ignore

    @returncode.setter
    def returncode(self, arg: int) -> None:
        raise NotImplementedError

    def append(self, proc: FluxProcess) -> None:
        self.procs.append(proc)

    def pop(self, /, i: int = -1) -> FluxProcess:
        return self.procs.pop(i)

    def cancel(self) -> None:
        with self.lock:
            for proc in self.procs:
                proc.cancel()

    def poll(self) -> int | None:
        stat: list[int | None] = []
        for proc in self.procs:
            stat.append(proc.poll())
        if any([_ is None for _ in stat]):
            return None
        return max(stat)  # type: ignore


class FluxSubmissionManager(HPCSubmissionManager):
    """Setup and submit jobs to the Flux scheduler"""

    name = "flux"
    lock: multiprocessing.synchronize.RLock = multiprocessing.RLock()

    def __init__(self, config: Config | None = None) -> None:
        super().__init__(config=config)
        self.flux: FluxExecutor | None = FluxExecutor()
        self.fh = Flux()
        if self.config.get("machine:resources") is None:
            if info := read_resource_info():
                scope = ConfigScope("flux", None, {"machine": {"resources": [info]}})
                self.config.push_scope(scope)
            else:
                logger.warning("Unable to determine system configuration from flux, using default")

    @property
    def supports_subscheduling(self) -> bool:
        return True

    @staticmethod
    def matches(name: str | None) -> bool:
        return name is not None and name.lower() == "flux"

    @property
    def submission_template(self) -> str:
        if "HPCC_FLUX_SUBMIT_TEMPLATE" in os.environ:
            return os.environ["HPCC_FLUX_SUBMIT_TEMPLATE"]
        return str(importlib.resources.files("hpcc_flux").joinpath("templates/submit.sh.in"))

    def create_jobspec(
        self,
        name: str,
        script: str,
        duration: float | timedelta | None = None,
        variables: dict[str, str | None] | None = None,
        output: str | None = None,
        error: str | None = None,
        #
        cpus: int | None = None,
        gpus: int | None = None,
        nodes: int | None = None,
        exclusive: bool = True,
    ) -> Jobspec:
        """
        Create the flux jobspec for the executable script ``script``.  The number of tasks
        ``tasks`` is the number of MPI tasks that ``script`` launches, not the number of
        copies of ``script`` that flux should launch.

        """
        kwds: dict[str, Any] = {"command": [script], "exclusive": exclusive}
        alloc = self.get_alloc_settings(cpus, gpus, nodes)
        kwds.update(alloc)
        jobspec = JobspecV1.from_nest_command(**kwds)
        jobspec.setattr("system.job.name", name)
        jobspec.stdout = output or "job-ouput.txt"
        jobspec.stderr = error or output or "job-error.txt"
        jobspec.duration = duration or 60.0  # duration is in seconds
        env = os.environ.copy()
        if variables:
            for key, val in variables.items():
                if val is None:
                    env.pop(key, None)
                else:
                    env[key] = val
        jobspec.environment = env
        return jobspec

    def get_alloc_settings(
        self,
        cpus: int | None = None,
        gpus: int | None = None,
        nodes: int | None = None,
    ) -> dict[str, Any]:
        alloc: dict[str, Any] = {}
        if nodes is not None:
            if cpus is None:
                cpus = nodes * self.config.count_per_node("cpu")
            if gpus is None:
                gpus = nodes * self.config.count_per_node("gpu", default=0)
        else:
            cpus = cpus or 1
            gpus = gpus or 0
            nodes = self.config.nodes_required(max_cpus=cpus, max_gpus=gpus)

        alloc["num_nodes"] = nodes
        alloc["num_slots"] = nodes
        if nodes > 1:
            cpus = max(1, math.ceil(cpus / nodes))
            gpus = max(0, math.ceil(gpus / nodes))

        alloc["cores_per_slot"] = cpus
        alloc["gpus_per_slot"] = gpus
        return alloc

    def shutdown(self):
        with self.lock:
            if self.flux:
                time.sleep(0.25)
                if flux.job:
                    jobs = flux.job.JobList(self.fh, filters=["running", "pending"]).jobs()
                    if jobs:
                        logger.warning(f"{len(jobs)} active jobs remain after shutdown")
                self.flux.shutdown(wait=False, cancel_futures=True)
                self.flux = None

    def submit(
        self,
        name: str,
        args: list[str],
        scriptname: str | None = None,
        qtime: float | None = None,
        submit_flags: list[str] | None = None,
        variables: dict[str, str | None] | None = None,
        output: str | None = None,
        error: str | None = None,
        nodes: int | None = None,
        cpus: int | None = None,
        gpus: int | None = None,
        exclusive: bool = True,
        **kwargs: Any,
    ) -> FluxProcess:
        cpus = cpus or kwargs.get("tasks")
        duration = timedelta(seconds=time_limit_in_seconds(qtime, pad=60))
        script = self.write_submission_script(
            name,
            args,
            scriptname,
            qtime=max(1, duration.total_seconds() / 60.0),  # time limit in minutes in the script
            submit_flags=submit_flags,
            variables=variables,
            output=output,
            error=error,
            nodes=nodes,
            cpus=cpus,
            gpus=gpus,
        )
        assert script is not None
        jobspec = self.create_jobspec(
            name,
            script,
            duration=duration,
            variables=variables,
            output=output,
            error=error,
            nodes=nodes,
            cpus=cpus,
            gpus=gpus,
            exclusive=exclusive,
        )
        fut = self.flux.submit(jobspec)  # type: ignore
        return FluxProcess(name, future=fut)

    def submitn(
        self,
        name: list[str],
        args: list[list[str]],
        scriptname: list[str] | None = None,
        qtime: list[float] | None = None,
        submit_flags: list[list[str]] | None = None,
        variables: list[dict[str, str | None]] | None = None,
        output: list[str] | None = None,
        error: list[str] | None = None,
        nodes: list[int] | None = None,
        cpus: list[int] | None = None,
        gpus: list[int] | None = None,
        **kwargs: Any,
    ) -> FluxMultiProcess:
        cpus = cpus or kwargs.get("tasks")  # backward compatible
        assert len(name) == len(args)
        procs = FluxMultiProcess(self.lock)
        submission_delay: float = 0.0
        if t := os.getenv("HPC_CONNECT_FLUX_SUBMITN_DELAY"):
            submission_delay = time_in_seconds(t)
        with self.lock:
            for i in range(len(name)):
                proc = self.submit(
                    name[i],
                    args[i],
                    select(scriptname, i, f"flux-submit-{i}"),
                    qtime=select(qtime, i),
                    submit_flags=select(submit_flags, i),
                    variables=select(variables, i),
                    output=select(output, i, f"flux-out-{i}"),
                    error=select(error, i, f"flux-err-{i}"),
                    nodes=select(nodes, i),
                    cpus=select(cpus, i),
                    gpus=select(gpus, i),
                    exclusive=False,
                )
                if submission_delay:
                    time.sleep(submission_delay)
                procs.append(proc)
        return procs

    def format_submission_data(
        self,
        name: str,
        args: list[str],
        qtime: float | None = None,
        submit_flags: list[str] | None = None,
        variables: dict[str, str | None] | None = None,
        output: str | None = None,
        error: str | None = None,
        #
        nodes: int | None = None,
        cpus: int | None = None,
        gpus: int | None = None,
    ) -> dict[str, Any]:
        data = super().format_submission_data(
            name,
            args,
            qtime=qtime,
            submit_flags=submit_flags,
            variables=variables,
            output=output,
            error=error,
            nodes=nodes,
            cpus=cpus,
            gpus=gpus,
        )
        alloc = self.get_alloc_settings(cpus=cpus, gpus=gpus, nodes=nodes)
        data.update(alloc)
        return data


def select(arg: Any, i: int, default: Any = None) -> Any:
    if isinstance(arg, list):
        return arg[i]
    return arg or default


def time_limit_in_seconds(qtime: float | None, pad: int = 0) -> int:
    """Return the time limit in seconds. Guarenteed return value >= 1."""
    limit = math.ceil(qtime or 1)
    if pad > 0:
        limit += pad
    return limit
