"""Single and multi-threaded executors."""

import datetime
import functools
import logging
import math
import os
import sys
import threading
from abc import ABCMeta, abstractmethod
from collections.abc import Iterable, MutableSequence
from threading import Lock
from typing import Optional, cast

import psutil
from mypy_extensions import mypyc_attr
from schema_salad.exceptions import ValidationException
from schema_salad.sourceline import SourceLine

from .command_line_tool import CallbackJob, ExpressionJob
from .context import RuntimeContext, getdefault
from .cuda import cuda_version_and_device_count
from .errors import WorkflowException
from .job import JobBase
from .loghandler import _logger
from .mutation import MutationManager
from .process import Process, cleanIntermediate, relocateOutputs
from .task_queue import TaskQueue
from .update import ORIGINAL_CWLVERSION
from .utils import CWLObjectType, JobsType
from .workflow import Workflow
from .workflow_job import WorkflowJob, WorkflowJobStep

TMPDIR_LOCK = Lock()


@mypyc_attr(allow_interpreted_subclasses=True)
class JobExecutor(metaclass=ABCMeta):
    """Abstract base job executor."""

    def __init__(self) -> None:
        """Initialize."""
        self.final_output: MutableSequence[CWLObjectType | None] = []
        self.final_status: list[str] = []
        self.output_dirs: set[str] = set()

    def __call__(
        self,
        process: Process,
        job_order_object: CWLObjectType,
        runtime_context: RuntimeContext,
        logger: logging.Logger = _logger,
    ) -> tuple[CWLObjectType | None, str]:
        return self.execute(process, job_order_object, runtime_context, logger)

    def output_callback(self, out: CWLObjectType | None, process_status: str) -> None:
        """Collect the final status and outputs."""
        self.final_status.append(process_status)
        self.final_output.append(out)

    @abstractmethod
    def run_jobs(
        self,
        process: Process,
        job_order_object: CWLObjectType,
        logger: logging.Logger,
        runtime_context: RuntimeContext,
    ) -> None:
        """Execute the jobs for the given Process."""

    def execute(
        self,
        process: Process,
        job_order_object: CWLObjectType,
        runtime_context: RuntimeContext,
        logger: logging.Logger = _logger,
    ) -> tuple[CWLObjectType | None, str]:
        """Execute the process."""

        self.final_output = []
        self.final_status = []

        if not runtime_context.basedir:
            raise WorkflowException("Must provide 'basedir' in runtimeContext")

        def check_for_abstract_op(tool: CWLObjectType) -> None:
            if tool["class"] == "Operation":
                raise SourceLine(tool, "class", WorkflowException, runtime_context.debug).makeError(
                    "Workflow has unrunnable abstract Operation"
                )

        process.visit(check_for_abstract_op)

        finaloutdir = None  # Type: Optional[str]
        original_outdir = runtime_context.outdir
        if isinstance(original_outdir, str):
            finaloutdir = os.path.abspath(original_outdir)
        runtime_context = runtime_context.copy()
        outdir = runtime_context.create_outdir()
        self.output_dirs.add(outdir)
        runtime_context.outdir = outdir
        runtime_context.mutation_manager = MutationManager()
        runtime_context.toplevel = True
        runtime_context.workflow_eval_lock = threading.Condition(threading.RLock())

        job_reqs: list[CWLObjectType] | None = None
        if "https://w3id.org/cwl/cwl#requirements" in job_order_object:
            if process.metadata.get(ORIGINAL_CWLVERSION) == "v1.0":
                raise WorkflowException(
                    "`cwl:requirements` in the input object is not part of CWL "
                    "v1.0. You can adjust to use `cwltool:overrides` instead; or you "
                    "can set the cwlVersion to v1.1"
                )
            job_reqs = cast(
                list[CWLObjectType],
                job_order_object["https://w3id.org/cwl/cwl#requirements"],
            )
        elif "cwl:defaults" in process.metadata and "https://w3id.org/cwl/cwl#requirements" in cast(
            CWLObjectType, process.metadata["cwl:defaults"]
        ):
            if process.metadata.get(ORIGINAL_CWLVERSION) == "v1.0":
                raise WorkflowException(
                    "`cwl:requirements` in the input object is not part of CWL "
                    "v1.0. You can adjust to use `cwltool:overrides` instead; or you "
                    "can set the cwlVersion to v1.1"
                )
            job_reqs = cast(
                Optional[list[CWLObjectType]],
                cast(CWLObjectType, process.metadata["cwl:defaults"])[
                    "https://w3id.org/cwl/cwl#requirements"
                ],
            )
        if job_reqs is not None:
            for req in job_reqs:
                process.requirements.append(req)

        self.run_jobs(process, job_order_object, logger, runtime_context)
        if runtime_context.validate_only is True:
            return (None, "ValidationSuccess")

        if self.final_output and self.final_output[0] is not None and finaloutdir is not None:
            self.final_output[0] = relocateOutputs(
                self.final_output[0],
                finaloutdir,
                self.output_dirs,
                runtime_context.move_outputs,
                runtime_context.make_fs_access(""),
                getdefault(runtime_context.compute_checksum, True),
                path_mapper=runtime_context.path_mapper,
            )

        if runtime_context.rm_tmpdir:
            if not runtime_context.cachedir:
                output_dirs: Iterable[str] = self.output_dirs
            else:
                output_dirs = filter(
                    lambda x: not x.startswith(runtime_context.cachedir),  # type: ignore
                    self.output_dirs,
                )
            cleanIntermediate(output_dirs)

        if self.final_output and self.final_status:
            if (
                runtime_context.research_obj is not None
                and isinstance(process, (JobBase, Process, WorkflowJobStep, WorkflowJob))
                and process.parent_wf
            ):
                process_run_id: str | None = None
                name = "primary"
                process.parent_wf.generate_output_prov(self.final_output[0], process_run_id, name)
                process.parent_wf.document.wasEndedBy(
                    process.parent_wf.workflow_run_uri,
                    None,
                    process.parent_wf.engine_uuid,
                    datetime.datetime.now(),
                )
                process.parent_wf.finalize_prov_profile(name=None)
            return (self.final_output[0], self.final_status[0])
        return (None, "permanentFail")


@mypyc_attr(allow_interpreted_subclasses=True)
class SingleJobExecutor(JobExecutor):
    """Default single-threaded CWL reference executor."""

    def run_jobs(
        self,
        process: Process,
        job_order_object: CWLObjectType,
        logger: logging.Logger,
        runtime_context: RuntimeContext,
    ) -> None:
        process_run_id: str | None = None

        # define provenance profile for single commandline tool
        if not isinstance(process, Workflow) and runtime_context.research_obj is not None:
            process.provenance_object = runtime_context.research_obj.initialize_provenance(
                full_name=runtime_context.cwl_full_name,
                # following are only set from main when directly command line tool
                # when nested in a workflow, they should be disabled since they would
                # already have been provided/initialized by the parent workflow prov-obj
                host_provenance=runtime_context.prov_host,
                user_provenance=runtime_context.prov_user,
                orcid=runtime_context.orcid,
                # single tool execution, so RO UUID = wf UUID = tool UUID
                run_uuid=runtime_context.research_obj.ro_uuid,
                fsaccess=runtime_context.make_fs_access(""),
            )
            process.parent_wf = process.provenance_object

        jobiter = process.job(job_order_object, self.output_callback, runtime_context)

        try:
            for job in jobiter:
                if job is not None:
                    if runtime_context.builder is not None and hasattr(job, "builder"):
                        job.builder = runtime_context.builder
                    if job.outdir is not None:
                        self.output_dirs.add(job.outdir)
                    if runtime_context.research_obj is not None:
                        if not isinstance(process, Workflow):
                            prov_obj = process.provenance_object
                        else:
                            prov_obj = job.prov_obj
                        if prov_obj:
                            runtime_context.prov_obj = prov_obj
                            prov_obj.fsaccess = runtime_context.make_fs_access("")
                            prov_obj.evaluate(
                                process,
                                job,
                                job_order_object,
                                runtime_context.research_obj,
                            )
                            process_run_id = prov_obj.record_process_start(process, job)
                            runtime_context = runtime_context.copy()
                        runtime_context.process_run_id = process_run_id
                    if runtime_context.validate_only is True:
                        if isinstance(job, WorkflowJob):
                            name = job.tool.lc.filename
                        else:
                            name = getattr(job, "name", str(job))
                        print(
                            f"{name} is valid CWL. No errors detected in the inputs.",
                            file=runtime_context.validate_stdout,
                        )
                        return
                    job.run(runtime_context)
                else:
                    logger.error("Workflow cannot make any more progress.")
                    break
        except (
            ValidationException,
            WorkflowException,
        ):  # pylint: disable=try-except-raise
            raise
        except Exception as err:
            logger.exception("Got workflow error")
            raise WorkflowException(str(err)) from err


if sys.platform != "darwin":

    def _max_cores() -> int:
        return len(psutil.Process().cpu_affinity())

else:

    def _max_cores() -> int:
        return psutil.cpu_count() or 1


class MultithreadedJobExecutor(JobExecutor):
    """
    Experimental multi-threaded CWL executor.

    Does simple resource accounting, will not start a job unless it
    has cores / ram available, but does not make any attempt to
    optimize usage.
    """

    def __init__(self, max_parallel: int = 0) -> None:
        """Initialize."""
        super().__init__()
        self.exceptions: list[WorkflowException] = []
        self.pending_jobs: list[JobsType] = []
        self.pending_jobs_lock = threading.Lock()

        self.max_ram = int(psutil.virtual_memory().available / 2**20)
        self.max_cores = float(max_parallel if max_parallel >= 1 else _max_cores())
        self.max_cuda = cuda_version_and_device_count()[1]
        self.allocated_ram = float(0)
        self.allocated_cores = float(0)
        self.allocated_cuda: int = 0

    def select_resources(
        self, request: dict[str, int | float], runtime_context: RuntimeContext
    ) -> dict[str, int | float]:  # pylint: disable=unused-argument
        """Naïve check for available cpu cores and memory."""
        result: dict[str, int | float] = {}
        maxrsc = {"cores": self.max_cores, "ram": self.max_ram}
        resources_types = {"cores", "ram"}
        if "cudaDeviceCountMin" in request or "cudaDeviceCountMax" in request:
            maxrsc["cudaDeviceCount"] = self.max_cuda
            resources_types.add("cudaDeviceCount")
        for rsc in resources_types:
            rsc_min = request[rsc + "Min"]
            if rsc_min > maxrsc[rsc]:
                raise WorkflowException(
                    f"Requested at least {rsc_min} {rsc} but only " f"{maxrsc[rsc]} available"
                )
            rsc_max = request[rsc + "Max"]
            if rsc_max < maxrsc[rsc]:
                result[rsc] = math.ceil(rsc_max)
            else:
                result[rsc] = maxrsc[rsc]

        result["tmpdirSize"] = math.ceil(request["tmpdirMin"])
        result["outdirSize"] = math.ceil(request["outdirMin"])

        return result

    def _runner(
        self,
        job: JobBase | WorkflowJob | CallbackJob | ExpressionJob,
        runtime_context: RuntimeContext,
        TMPDIR_LOCK: threading.Lock,
    ) -> None:
        """Job running thread."""
        try:
            _logger.debug(
                "job: {}, runtime_context: {}, TMPDIR_LOCK: {}".format(
                    job, runtime_context, TMPDIR_LOCK
                )
            )
            job.run(runtime_context, TMPDIR_LOCK)
        except WorkflowException as err:
            _logger.exception(f"Got workflow error: {err}")
            self.exceptions.append(err)
        except Exception as err:  # pylint: disable=broad-except
            _logger.exception(f"Got workflow error: {err}")
            wf_exc = WorkflowException(str(err))
            wf_exc.__cause__ = err
            wf_exc.__suppress_context__ = True
            self.exceptions.append(wf_exc)
        finally:
            if runtime_context.workflow_eval_lock:
                with runtime_context.workflow_eval_lock:
                    if isinstance(job, JobBase):
                        ram = job.builder.resources["ram"]
                        self.allocated_ram -= ram
                        cores = job.builder.resources["cores"]
                        self.allocated_cores -= cores
                        cudaDevices: int = cast(
                            int, job.builder.resources.get("cudaDeviceCount", 0)
                        )
                        self.allocated_cuda -= cudaDevices
                    runtime_context.workflow_eval_lock.notify_all()

    def run_job(
        self,
        job: JobsType | None,
        runtime_context: RuntimeContext,
    ) -> None:
        """Execute a single Job in a separate thread."""
        if job is not None:
            with self.pending_jobs_lock:
                self.pending_jobs.append(job)

        with self.pending_jobs_lock:
            n = 0
            while (n + 1) <= len(self.pending_jobs):
                # Simple greedy resource allocation strategy.  Go
                # through pending jobs in the order they were
                # generated and add them to the queue only if there
                # are resources available.
                job = self.pending_jobs[n]
                if isinstance(job, JobBase):
                    ram = job.builder.resources["ram"]
                    cores = job.builder.resources["cores"]
                    cudaDevices = cast(int, job.builder.resources.get("cudaDeviceCount", 0))
                    if ram > self.max_ram or cores > self.max_cores or cudaDevices > self.max_cuda:
                        _logger.error(
                            'Job "%s" cannot be run, requests more resources (%s) '
                            "than available on this host (already allocated ram is %d, "
                            "allocated cores is %d, allocated CUDA is %d, "
                            "max ram %d, max cores %d, max CUDA %d).",
                            job.name,
                            job.builder.resources,
                            self.allocated_ram,
                            self.allocated_cores,
                            self.allocated_cuda,
                            self.max_ram,
                            self.max_cores,
                            self.max_cuda,
                        )
                        self.pending_jobs.remove(job)
                        return

                    if (
                        self.allocated_ram + ram > self.max_ram
                        or self.allocated_cores + cores > self.max_cores
                        or self.allocated_cuda + cudaDevices > self.max_cuda
                    ):
                        _logger.debug(
                            'Job "%s" cannot run yet, resources (%s) are not '
                            "available (already allocated ram is %d, allocated cores is %d, "
                            "allocated CUDA devices is %d, "
                            "max ram %d, max cores %d, max CUDA %d).",
                            job.name,
                            job.builder.resources,
                            self.allocated_ram,
                            self.allocated_cores,
                            self.allocated_cuda,
                            self.max_ram,
                            self.max_cores,
                            self.max_cuda,
                        )
                        n += 1
                        continue

                if isinstance(job, JobBase):
                    ram = job.builder.resources["ram"]
                    self.allocated_ram += ram
                    cores = job.builder.resources["cores"]
                    self.allocated_cores += cores
                    cuda = cast(int, job.builder.resources.get("cudaDevices", 0))
                    self.allocated_cuda += cuda
                self.taskqueue.add(
                    functools.partial(self._runner, job, runtime_context, TMPDIR_LOCK),
                    runtime_context.workflow_eval_lock,
                )
                self.pending_jobs.remove(job)

    def wait_for_next_completion(self, runtime_context: RuntimeContext) -> None:
        """Wait for jobs to finish."""
        if runtime_context.workflow_eval_lock is not None:
            runtime_context.workflow_eval_lock.wait(timeout=3)
        if self.exceptions:
            raise self.exceptions[0]

    def run_jobs(
        self,
        process: Process,
        job_order_object: CWLObjectType,
        logger: logging.Logger,
        runtime_context: RuntimeContext,
    ) -> None:
        self.taskqueue: TaskQueue = TaskQueue(threading.Lock(), int(math.ceil(self.max_cores)))
        try:
            jobiter = process.job(job_order_object, self.output_callback, runtime_context)

            if runtime_context.workflow_eval_lock is None:
                raise WorkflowException("runtimeContext.workflow_eval_lock must not be None")

            runtime_context.workflow_eval_lock.acquire()
            for job in jobiter:
                if job is not None:
                    if isinstance(job, JobBase):
                        job.builder = runtime_context.builder or job.builder
                        if job.outdir is not None:
                            self.output_dirs.add(job.outdir)

                self.run_job(job, runtime_context)

                if job is None:
                    if self.taskqueue.in_flight > 0:
                        self.wait_for_next_completion(runtime_context)
                    else:
                        logger.error("Workflow cannot make any more progress.")
                        break

            self.run_job(None, runtime_context)
            while self.taskqueue.in_flight > 0:
                self.wait_for_next_completion(runtime_context)
                self.run_job(None, runtime_context)

            runtime_context.workflow_eval_lock.release()
        finally:
            self.taskqueue.drain()
            self.taskqueue.join()


class NoopJobExecutor(JobExecutor):
    """Do nothing executor, for testing purposes only."""

    def run_jobs(
        self,
        process: Process,
        job_order_object: CWLObjectType,
        logger: logging.Logger,
        runtime_context: RuntimeContext,
    ) -> None:
        pass

    def execute(
        self,
        process: Process,
        job_order_object: CWLObjectType,
        runtime_context: RuntimeContext,
        logger: logging.Logger | None = None,
    ) -> tuple[CWLObjectType | None, str]:
        return {}, "success"
