"""Galaxy job runners to use Amazon AWS native compute resources, such as AWS Batch."""

import bisect
import hashlib
import json
import logging
import os
import re
from typing import (
    Any,
    TYPE_CHECKING,
    Union,
)

from galaxy import model
from galaxy.job_execution.output_collect import default_exit_code_file
from galaxy.jobs.runners import (
    AsynchronousJobRunner,
    AsynchronousJobState,
    JobState,
)
from galaxy.util import (
    smart_str,
    unicodify,
)

if TYPE_CHECKING:
    from galaxy.jobs import MinimalJobWrapper

BOTO3_IMPORT_MSG = (
    "The Python 'boto3' package is required to use "
    "this feature, please install it or correct the "
    "following error:\nImportError {msg!s}"
)

try:
    import boto3
except ImportError as e:
    boto3 = None  # type: ignore[assignment]
    BOTO3_IMPORT_MSG.format(msg=unicodify(e))


__all__ = ("AWSBatchJobRunner",)
log = logging.getLogger(__name__)


class AWSBatchRunnerException(Exception):
    pass


def to_dict(segments, v):
    if len(segments) == 0:
        return v
    return {segments[0]: to_dict(segments[1:], v)}


def _write_logfile(logfile, msg):
    with open(logfile, "w") as fil:
        fil.write(msg)


def _add_galaxy_environment_variables(vcpu, memory):
    # Set:
    # GALAXY_SLOTS, round 0.25 vpuc to 1.
    # GALAXY_MEMORY_MB
    return [
        {"name": "GALAXY_SLOTS", "value": str(int(max(vcpu, 1)))},
        {"name": "GALAXY_MEMORY_MB", "value": str(memory)},
    ]


def _add_resource_requirements(destination_params):
    rval = [
        {"type": "VCPU", "value": str(destination_params.get("vcpu"))},
        {"type": "MEMORY", "value": str(destination_params.get("memory"))},
    ]
    if n_gpu := destination_params.get("gpu"):
        rval.append({"type": "GPU", "value": str(n_gpu)})
    return rval


class AWSBatchJobRunner(AsynchronousJobRunner[AsynchronousJobState]):
    """
    This runner uses container only. It requires that an AWS EFS is mounted as a local drive
    and all Galaxy job-related paths, such as objects, job_directory, tool_directory and so
    on, are placed in the EFS drive. Via this runner, Galaxy jobs are dispatched to AWS Batch
    for compute using the docker image specified by a Galaxy tool. As AWS EFS is designed to
    be able to mount at multiple places with read and write capabilities, Galaxy and Batch
    containers share the same EFS drive as a local device. Sample configurations can be found
    in `config/job_conf.sample.yml`.
    """

    runner_name = "AWSBatchRunner"
    RUNNER_PARAM_SPEC_KEY = "runner_param_specs"
    JOB_NAME_PREFIX = "galaxy-"
    # AWS Batch queries up to 100 jobs at once.
    MAX_JOBS_PER_QUERY = 100
    # Higher minimum interval as jobs are queried in batches.
    MIN_QUERY_INTERVAL = 10

    # fmt: off
    RUNNER_PARAM_SPEC = {
        "aws_access_key_id": {
            "map": str,
        },
        "aws_secret_access_key": {
            "map": str,
        },
        "region": {
            "map": str,
        },
    }

    # fmt: off
    DESTINATION_PARAMS_SPEC = {
        "vcpu": {
            "default": 1.0,
            "map": (lambda x: int(float(x)) if int(float(x)) == float(x) else float(x)),
        },
        "memory": {
            "default": 2048,
            "map": int,
        },
        "gpu": {
            "default": 0,
            "map": int,
        },
        "job_queue": {
            "default": "",
            "map": str,
            "required": True,
        },
        "job_role_arn": {
            "default": "",
            "map": str,
            "required": True,
        },
        "efs_filesystem_id": {
            "default": "",
            "map": str,
            "required": True,
        },
        "efs_mount_point": {
            "default": "",
            "map": str,
            "required": True,
        },
        "execute_role_arn": {
            "default": "",
            "map": str,
        },
        "fargate_version": {
            "default": "",
            "map": str,
        },
        "auto_platform": {
            "default": False,
            "map": lambda x: str(x).lower() == "true",
        },
        "ec2_host_volumes": {
            "default": "",
            "map": str,
        },
        "privileged": {
            "default": False,
            "map": lambda x: str(x).lower() == "true",
        },
        "retry_attempts": {
            "default": 1,
            "map": int,
        },
        "retry_on_exit_statusReason": {
            "default": "",
            "map": str,
        },
        "retry_on_exit_reason": {
            "default": "",
            "map": str,
        },
        "retry_on_exit_exitCode": {
            "default": "",
            "map": str,
        },
        "retry_on_exit_action": {
            "default": "RETRY",
            "map": str,
        },
    }

    # fmt: off
    FARGATE_RESOURCES = {
        0.25: [512, 1024, 2048],
        0.5: [1024, 2048, 3072, 4096],
        1: [2048, 3072, 4096, 5120, 6144, 7168, 8192],
        2: [4096, 5120, 6144, 7168, 8192, 9216, 10240, 11264, 12288, 13312, 14336, 15360, 16384],
        4: [8192, 9216, 10240, 11264, 12288, 13312, 14336, 15360, 16384, 17408, 18432,
            19456, 20480, 21504, 22528, 23552, 24576, 25600, 26624, 27648, 28672, 29696, 30720]
    }

    def __init__(self, app, nworkers, **kwargs):
        """Initialize this job runner and start the monitor thread"""
        assert boto3, BOTO3_IMPORT_MSG
        if self.RUNNER_PARAM_SPEC_KEY not in kwargs:
            kwargs[self.RUNNER_PARAM_SPEC_KEY] = {}
        kwargs[self.RUNNER_PARAM_SPEC_KEY].update(self.RUNNER_PARAM_SPEC)
        super().__init__(app, nworkers, **kwargs)
        session = boto3.Session(
            aws_access_key_id=self.runner_params.get("aws_access_key_id") or None,
            aws_secret_access_key=self.runner_params.get("aws_secret_access_key") or None,
            region_name=self.runner_params.get("region") or None,
        )
        self._batch_client = session.client("batch")

    @property
    def monitor_sleep_time(self):
        return max(self.app.config.job_runner_monitor_sleep, self.MIN_QUERY_INTERVAL)

    def queue_job(self, job_wrapper: "MinimalJobWrapper") -> None:
        log.debug(f"Starting queue_job for job {job_wrapper.get_id_tag()}")
        if not self.prepare_job(job_wrapper, include_metadata=False, modify_command_for_container=False):
            log.debug(f"Not ready {job_wrapper.get_id_tag()}")
            return

        job_destination = job_wrapper.job_destination
        destination_params = self.parse_destination_params(job_destination.params)
        job_def = self._get_job_definition(job_wrapper, destination_params)
        job_name, job_id = self._submit_job(job_def, job_wrapper, destination_params)
        job_wrapper.set_external_id(job_id)
        ajs = AsynchronousJobState(
            job_wrapper=job_wrapper,
            job_destination=job_destination,
            files_dir=job_wrapper.working_directory,
            job_name=job_name,
            job_id=job_id,
        )
        self.monitor_queue.put(ajs)

    def _get_job_definition(self, job_wrapper, destination_params):
        tool_id = job_wrapper.tool.id
        container_image = self._find_container(job_wrapper).container_id
        h = hashlib.new("sha256")
        h.update(smart_str(container_image))
        s = json.dumps(destination_params, sort_keys=True)
        h.update(smart_str(s))
        queue_name = destination_params.get("job_queue").rsplit("/", 1)[-1]

        jd_name = f"galaxy_tool__{tool_id}__{h.hexdigest()}__{queue_name}"
        jd_name = re.sub(r'[^-_A-Za-z0-9]', '-', jd_name)[:128]
        res = self._batch_client.describe_job_definitions(jobDefinitionName=jd_name, status="ACTIVE")
        if not res["jobDefinitions"]:
            jd_arn = self._register_job_definition(jd_name, container_image, destination_params)
        else:
            jd_arn = res["jobDefinitions"][0]["jobDefinitionArn"]
            log.debug(f"Found existing job definition: {jd_name}.")

        return jd_arn

    def _get_mount_volumes(self, destination_params):
        volumes, mount_points = [], []
        volumes.append(
            {
                "name": "efs_whole",
                "efsVolumeConfiguration": {
                    "fileSystemId": destination_params.get("efs_filesystem_id"),
                    "rootDirectory": "/",
                    "transitEncryption": "ENABLED",
                    "authorizationConfig": {"iam": "ENABLED"},
                },
            },
        )
        mount_points.append(
            {
                "containerPath": destination_params.get("efs_mount_point"),
                "readOnly": False,
                "sourceVolume": "efs_whole",
            },
        )
        if destination_params.get("platform") == 'Fargate':   # Fargate doesn't support host volumes
            return volumes, mount_points

        if (ec2_host_volumes := destination_params.get("ec2_host_volumes")):
            for ix, vol in enumerate(ec2_host_volumes.split(",")):
                vol = vol.strip()
                vol_name = "host_vol_" + str(ix)
                volumes.append(
                    {
                        "name": vol_name,
                        "host": {"sourcePath": vol},
                    },
                )
                mount_points.append(
                    {
                        "containerPath": vol,
                        "readOnly": False,
                        "sourceVolume": vol_name,
                    },
                )
        return volumes, mount_points

    def _get_retry_strategy(self, destination_params):
        """ Make a simple one-condition retry strategy
        """
        # TODO make multi-condition retry strategies
        attemps = destination_params.get("retry_attempts")
        status_reason = destination_params.get("retry_on_exit_statusReason")
        reason = destination_params.get("retry_on_exit_reason")
        exit_code = destination_params.get("retry_on_exit_exitCode")
        action = destination_params.get("retry_on_exit_action")

        if attemps <= 1:
            return

        strategy = {
            "attempts": attemps,
            "evaluateOnExit": [{"action": action}],
        }
        if status_reason:
            strategy["evaluateOnExit"][0]["onStatusReason"] = status_reason
        if reason:
            strategy["evaluateOnExit"][0]["onReason"] = reason
        if exit_code:
            strategy["evaluateOnExit"][0]["onExitCode"] = exit_code
        return strategy

    def _register_job_definition(self, jd_name, container_image, destination_params):
        log.debug(f"Registering a new job definition: {jd_name}.")
        platform = destination_params.get("platform")
        volumes, mount_points = self._get_mount_volumes(destination_params)

        # TODO: support multi-node
        containerProperties = {
            "image": container_image,
            "command": [
                "/bin/sh",
            ],
            "jobRoleArn": destination_params.get("job_role_arn"),
            "executionRoleArn": destination_params.get("execute_role_arn") or destination_params.get("job_role_arn"),
            "volumes": volumes,
            "mountPoints": mount_points,
            "resourceRequirements": _add_resource_requirements(destination_params),
            "environment": _add_galaxy_environment_variables(
                destination_params.get("vcpu"), destination_params.get("memory"),
            ),
            "user": f"{os.getuid()}:{os.getgid()}",
            "privileged": destination_params.get("privileged"),
            "logConfiguration": {"logDriver": "awslogs"},
        }
        if platform == "FARGATE":
            containerProperties.update(
                {
                    "networkConfiguration": {"assignPublicIp": "ENABLED"},
                    "fargatePlatformConfiguration": {"platformVersion": destination_params.get("fargate_version")},
                    "logConfiguration": {"logDriver": "awslogs"},
                }
            )
        other_kwargs = {}
        if (retry_strategy := self._get_retry_strategy(destination_params)):
            other_kwargs["retryStrategy"] = retry_strategy

        res = self._batch_client.register_job_definition(
            jobDefinitionName=jd_name,
            type="container",
            platformCapabilities=[platform],
            containerProperties=containerProperties,
            **other_kwargs,
        )

        assert res["ResponseMetadata"]["HTTPStatusCode"] == 200
        return res["jobDefinitionArn"]

    def _submit_job(self, job_def, job_wrapper, destination_params):
        job_name = self.JOB_NAME_PREFIX + job_wrapper.get_id_tag()
        command_script_path = self.write_command(job_wrapper)

        log.info(f"Submitting job {job_name} to AWS Batch.")
        res = self._batch_client.submit_job(
            jobName=job_name,
            jobQueue=destination_params.get("job_queue"),
            jobDefinition=job_def,
            containerOverrides={
                "command": [
                    "/bin/bash",
                    f"{command_script_path}",
                ]
            },
        )

        assert res["ResponseMetadata"]["HTTPStatusCode"] == 200
        return job_name, res["jobId"]

    def stop_job(self, job_wrapper):
        job = job_wrapper.get_job()
        external_id = job.get_job_runner_external_id()
        job_name = self.JOB_NAME_PREFIX + job_wrapper.get_id_tag()

        self._batch_client.terminate_job(jobId=external_id, reason="Terminated by Galaxy!")
        msg = "Job {name!r} is terminated"
        log.debug(msg.format(name=job_name))

    def recover(self, job: model.Job, job_wrapper: "MinimalJobWrapper") -> None:
        msg = "(name!r/runner!r) is still in {state!s} state, adding to the runner monitor queue"
        job_id = job.get_job_runner_external_id()
        job_name = self.JOB_NAME_PREFIX + job_wrapper.get_id_tag()
        ajs = AsynchronousJobState(
            job_wrapper=job_wrapper,
            job_destination=job_wrapper.job_destination,
            files_dir=job_wrapper.working_directory,
            job_id=str(job_id),
            job_name=job_name,
        )
        if job.state in (model.Job.states.RUNNING, model.Job.states.STOPPED):
            log.debug(msg.format(name=job.id, runner=job.job_runner_name, state=job.state))
            ajs.old_state = model.Job.states.RUNNING
            ajs.running = True
            self.monitor_queue.put(ajs)
        elif job.state == model.Job.states.QUEUED:
            log.debug(msg.format(name=job.id, runner=job.job_runner_name, state="queued"))
            ajs.old_state = model.Job.states.QUEUED
            ajs.running = False
            self.monitor_queue.put(ajs)

    def fail_job(self, job_state: JobState, exception: bool = False, message: str = "Job failed", full_status: Union[dict[str, Any], None] = None) -> None:
        job = job_state.job_wrapper.get_job()
        if job_state.stop_job and job.state != model.Job.states.NEW:
            self.stop_job(job_state.job_wrapper)
        job_state.job_wrapper.reclaim_ownership()
        self._handle_runner_state("failure", job_state)
        if not job_state.runner_state_handled:
            full_status = full_status or {}
            tool_stdout = full_status.get("stdout")
            tool_stderr = full_status.get("stderr")
            fail_message = getattr(job_state, "fail_message", message)
            job_state.job_wrapper.fail(
                fail_message, tool_stdout=tool_stdout, tool_stderr=tool_stderr, exception=exception
            )
            self._finish_or_resubmit_job(job_state, "", fail_message)
            if job_state.job_wrapper.cleanup_job == "always":
                job_state.cleanup()

    def check_watched_items(self) -> None:
        done: set[str] = set()
        self.check_watched_items_by_batch(0, len(self.watched), done)
        self.watched = [ajs for ajs in self.watched if ajs.job_id not in done]

    def check_watched_items_by_batch(self, start: int, end: int, done: set[str]) -> None:
        async_job_states = self.watched[start : start + self.MAX_JOBS_PER_QUERY]
        if not async_job_states:
            return

        jobs_dict = {ajs.job_id: ajs for ajs in async_job_states if ajs.job_id is not None}

        resp = self._batch_client.describe_jobs(jobs=list(jobs_dict.keys()))

        gotten = set()
        for job in resp["jobs"]:
            status = job["status"]
            job_id = job["jobId"]
            gotten.add(job_id)
            job_state = jobs_dict[job_id]

            if status == "SUCCEEDED":
                self._mark_as_successful(job_state)
                done.add(job_id)
            elif status == "FAILED":
                reason = job["statusReason"]
                self._mark_as_failed(job_state, reason)
                done.add(job_id)
            elif status in ("STARTING", "RUNNING"):
                self._mark_as_active(job_state)
            # remain queued for "SUBMITTED", "PENDING" and "RUNNABLE"
            # TODO else?

        for job_id, job_state in jobs_dict.items():
            if job_id in gotten:
                continue
            reason = f"The track of Job {job_state} was lost for unknown reason!"
            self._mark_as_failed(job_state, reason)
            done.add(job_id)

        self.check_watched_items_by_batch(start + self.MAX_JOBS_PER_QUERY, end, done)

    def _mark_as_successful(self, job_state: AsynchronousJobState) -> None:
        _write_logfile(job_state.output_file, "")
        _write_logfile(job_state.error_file, "")
        job_state.running = False
        self.mark_as_finished(job_state)

    def _mark_as_active(self, job_state: AsynchronousJobState) -> None:
        job_state.running = True
        job_state.job_wrapper.change_state(model.Job.states.RUNNING)

    def _mark_as_failed(self, job_state: AsynchronousJobState, reason: str) -> None:
        _write_logfile(job_state.error_file, reason)
        job_state.running = False
        job_state.stop_job = False
        job_state.job_wrapper.change_state(model.Job.states.ERROR)
        job_state.fail_message = reason
        self.mark_as_failed(job_state)

    def parse_destination_params(self, params):
        if not params.get("docker_enabled"):
            raise AWSBatchRunnerException("AWSBatchJobRunner needs 'docker_enabled' to be set as True!")

        check_required = []
        parsed_params = {}
        for k, spec in self.DESTINATION_PARAMS_SPEC.items():
            value = params.get(k, spec.get("default"))
            if spec.get("required") and not value:
                check_required.append(k)
            mapper = spec.get("map")
            parsed_params[k] = mapper(value)  # type: ignore[operator]
        if check_required:
            raise AWSBatchRunnerException(
                "AWSBatchJobRunner requires the following params to be provided: {}.".format(", ".join(check_required))
            )

        # parse Platform
        platform = "EC2"
        auto_platform = parsed_params.get("auto_platform")
        fargate_version = parsed_params.get("fargate_version")
        vcpu = parsed_params.get("vcpu")
        memory = parsed_params.get("memory")
        gpu = parsed_params.get("gpu")

        if auto_platform and not fargate_version:
            raise AWSBatchRunnerException("AWSBatchJobRunner needs 'farget_version' to be set to enable auto platform!")

        if gpu and (fargate_version or auto_platform):
            raise AWSBatchRunnerException(
                "GPU mode is not allowed when 'fargate_version' and/or 'auto_platform' are set!"
            )

        if fargate_version and not auto_platform:
            platform = "FARGATE"
        if auto_platform:
            fargate_vcpus = list(self.FARGATE_RESOURCES.keys())
            max_vcpu = fargate_vcpus[-1]
            max_memory = self.FARGATE_RESOURCES[max_vcpu][-1]
            if vcpu <= max_vcpu and memory <= max_memory:  # type: ignore[operator]
                c_ix = bisect.bisect_left(fargate_vcpus, vcpu)  # type: ignore[type-var]
                length = len(fargate_vcpus)
                while c_ix < length:
                    c = fargate_vcpus[c_ix]
                    m_ix = bisect.bisect_left(self.FARGATE_RESOURCES[c], memory)  # type: ignore[type-var]
                    if m_ix < len(self.FARGATE_RESOURCES[c]):
                        platform = "FARGATE"
                        parsed_params["vcpu"] = c
                        parsed_params["memory"] = self.FARGATE_RESOURCES[c][m_ix]
                        break
                    c_ix += 1
            # parse JOB QUEUE
            job_queues = parsed_params.get("job_queue").split(",")  # type: ignore[union-attr]
            if len(job_queues) < 2:
                raise AWSBatchRunnerException(
                    "AWSBatchJobRunner needs to set TWO job queues ('Farget Queue, EC2 Qeueue')"
                    " when 'auto_platform' is enabled!"
                )
            parsed_params["job_queue"] = job_queues[platform == "EC2"].strip()

        if platform == "Fargate" and parsed_params.get("ec2_host_volumes") and not auto_platform:
            raise AWSBatchRunnerException(
                "AWSBatchJobRunner: param 'ec2_host_volumes' only works for EC2 platform"
                " unless 'auto_platform' is enabled!"
            )

        parsed_params["platform"] = platform
        return parsed_params

    def write_command(self, job_wrapper: "MinimalJobWrapper") -> str:
        # Create command script instead passing it in the container
        # preventing wrong characters parsing.
        command_line = job_wrapper.runner_command_line
        job_id = job_wrapper.get_id_tag()
        job_file = JobState.default_job_file(job_wrapper.working_directory, job_id)
        exit_code_path = default_exit_code_file(job_wrapper.working_directory, job_id)
        job_script_props = {
            "command": command_line,
            "exit_code_path": exit_code_path,
            "working_directory": job_wrapper.working_directory,
            "shell": job_wrapper.shell,
            "galaxy_virtual_env": None,
        }
        job_file_contents = self.get_job_file(job_wrapper, **job_script_props)
        self.write_executable_script(job_file, job_file_contents, job_io=job_wrapper.job_io)
        return job_file
