import os
from datetime import (
    datetime,
    timedelta,
)
from functools import partial
from typing import (
    Optional,
    TYPE_CHECKING,
    Union,
)

from sqlalchemy.orm import Session

import galaxy.workflow.schedulers
from galaxy import model
from galaxy.exceptions import HandlerAssignmentError
from galaxy.jobs.handler import InvocationGrabber
from galaxy.model.base import check_database_connection
from galaxy.schema.invocation import (
    FailureReason,
    InvocationFailureDatasetFailed,
    InvocationState,
    InvocationUnexpectedFailure,
)
from galaxy.schema.tasks import (
    MaterializeDatasetInstanceTaskRequest,
    RequestUser,
)
from galaxy.util import (
    plugin_config,
    unicodify,
)
from galaxy.util.custom_logging import get_logger
from galaxy.util.monitors import Monitors
from galaxy.util.xml_macros import load
from galaxy.web_stack.handlers import ConfiguresHandlers
from galaxy.web_stack.message import WorkflowSchedulingMessage

if TYPE_CHECKING:
    from galaxy.structured_app import MinimalManagerApp
    from galaxy.util import Element
    from galaxy.workflow.schedulers import (
        ActiveWorkflowSchedulingPlugin,
        WorkflowSchedulingPlugin,
    )

log = get_logger(__name__)

DEFAULT_SCHEDULER_ID = "default"  # well actually this should be called DEFAULT_DEFAULT_SCHEDULER_ID...
DEFAULT_SCHEDULER_PLUGIN_TYPE = "core"
DEFAULT_SCHEDULER_BACKFILL_SECONDS = int(os.getenv("GALAXY_SCHEDULER_BACKFILL_SECONDS", 300))

EXCEPTION_MESSAGE_SHUTDOWN = "Exception raised while attempting to shutdown workflow scheduler."
EXCEPTION_MESSAGE_NO_SCHEDULERS = "Failed to defined workflow schedulers - no workflow schedulers defined."
EXCEPTION_MESSAGE_NO_DEFAULT_SCHEDULER = (
    "Failed to defined workflow schedulers - no workflow scheduler found for default id '%s'."
)
EXCEPTION_MESSAGE_DUPLICATE_SCHEDULERS = (
    "Failed to defined workflow schedulers - workflow scheduling plugin id '%s' duplicated."
)
EXCEPTION_MESSAGE_SERIALIZE = "Parallelization is not desired but handler assignment methods are non-deterministic. Set DB_PREASSIGN in workflow_schedulers_conf.xml."


class WorkflowSchedulingManager(ConfiguresHandlers):
    """A workflow scheduling manager based loosely on pattern established by
    ``galaxy.manager.JobManager``. Only schedules workflows on handler
    processes.
    """

    DEFAULT_BASE_HANDLER_POOLS = ("workflow-schedulers", "job-handlers")

    def __init__(self, app: "MinimalManagerApp") -> None:
        super().__init__(app)
        self.__handlers_configured: bool = False
        self.workflow_schedulers: dict[str, WorkflowSchedulingPlugin] = {}
        self.active_workflow_schedulers: dict[str, ActiveWorkflowSchedulingPlugin] = {}
        # Passive workflow schedulers won't need to be monitored I guess.

        self.request_monitor = None

        self.default_handler_id = None

        self.__plugin_classes = self.__plugins_dict()
        self.__init_schedulers()

        if self._is_workflow_handler():
            log.debug("Starting workflow schedulers")
            self.__start_schedulers()
            if self.active_workflow_schedulers:
                self.__start_request_monitor()

        # When assinging handlers to workflows being queued - use job_conf
        # if not explicit workflow scheduling handlers have be specified or
        # else use those explicit workflow scheduling handlers (on self).
        self.__handlers_config = self if self.__handlers_configured else app.job_config

        if not self._is_workflow_handler():
            # Process should not schedule workflows but should check for any unassigned to handlers
            self.__startup_recovery()

    def __startup_recovery(self) -> None:
        sa_session = self.app.model.context
        for invocation_id in model.WorkflowInvocation.poll_unhandled_workflow_ids(sa_session):
            log.info(
                "(%s) Handler unassigned at startup, resubmitting workflow invocation for assignment", invocation_id
            )
            workflow_invocation = sa_session.get(model.WorkflowInvocation, invocation_id)
            assert workflow_invocation is not None
            self._assign_handler(workflow_invocation)

    def _handle_setup_msg(self, workflow_invocation_id=None):
        sa_session = self.app.model.context
        workflow_invocation = sa_session.get(model.WorkflowInvocation, workflow_invocation_id)
        if workflow_invocation.handler is None:
            workflow_invocation.handler = self.app.config.server_name
            sa_session.add(workflow_invocation)
            sa_session.commit()
        else:
            log.warning(
                "(%s) Handler '%s' received setup message for workflow invocation but handler '%s' is"
                " already assigned, ignoring",
                workflow_invocation.id,
                self.app.config.server_name,
                workflow_invocation.handler,
            )

    def _is_workflow_handler(self):
        # If we have explicitly configured handlers, check them.
        # Else just make sure we are a job handler.
        if self.__handlers_configured:
            is_handler = self.is_handler
        else:
            is_handler = self.app.is_job_handler
        return is_handler

    def _queue_callback(self, workflow_invocation):
        # There isn't an in-memory queue for workflow schedulers, so if MEM_SELF is used just assign with the DB
        workflow_invocation.handler = self.app.config.server_name
        sa_session = self.app.model.context
        sa_session.add(workflow_invocation)
        sa_session.commit()

    def _message_callback(self, workflow_invocation):
        return WorkflowSchedulingMessage(task="setup", workflow_invocation_id=workflow_invocation.id)

    def _assign_handler(self, workflow_invocation: model.WorkflowInvocation, flush: bool = True) -> str:
        # Use random-ish integer history_id to produce a consistent index to pick
        # job handler with.
        random_index = workflow_invocation.history.id
        queue_callback = partial(self._queue_callback, workflow_invocation)
        message_callback = partial(self._message_callback, workflow_invocation)
        if self.app.config.parallelize_workflow_scheduling_within_histories:
            random_index = None
        return self.__handlers_config.assign_handler(
            workflow_invocation,
            configured=None,
            flush=flush,
            index=random_index,
            queue_callback=queue_callback,
            message_callback=message_callback,
        )

    def shutdown(self):
        exception = None
        for workflow_scheduler in self.workflow_schedulers.values():
            try:
                workflow_scheduler.shutdown()
            except Exception as e:
                exception = exception or e
                log.exception(EXCEPTION_MESSAGE_SHUTDOWN)
        if self.request_monitor:
            try:
                self.request_monitor.shutdown()
            except Exception as e:
                exception = exception or e
                log.exception("Failed to shutdown workflow request monitor.")

        if exception:
            raise exception

    def queue(
        self,
        workflow_invocation: model.WorkflowInvocation,
        request_params,
        flush: bool = True,
        initial_state: Optional[InvocationState] = None,
    ):
        initial_state = initial_state or model.WorkflowInvocation.states.NEW
        workflow_invocation.set_state(initial_state)
        workflow_invocation.scheduler = request_params.get("scheduler", None) or self.default_scheduler_id
        sa_session = self.app.model.context
        sa_session.add(workflow_invocation)

        # Assign handler
        try:
            self._assign_handler(workflow_invocation, flush=flush)
        except HandlerAssignmentError:
            raise RuntimeError(f"Unable to set a handler for workflow invocation '{workflow_invocation.id}'")

        return workflow_invocation

    def __start_schedulers(self):
        for workflow_scheduler in self.workflow_schedulers.values():
            workflow_scheduler.startup(self.app)

    def __plugins_dict(self) -> dict[str, type["WorkflowSchedulingPlugin"]]:
        return plugin_config.plugins_dict(galaxy.workflow.schedulers, "plugin_type")

    @property
    def __stack_has_pool(self):
        # TODO: In the future it should be possible to map workflows to handlers based on workflow params. When that
        # happens, we'll need to defer pool checks until execution time.
        return any(map(self.app.application_stack.has_pool, self.DEFAULT_BASE_HANDLER_POOLS))

    def __init_schedulers(self):
        config_file = self.app.config.workflow_schedulers_config_file
        use_default_scheduler = False
        if not config_file or (
            not os.path.exists(config_file) and not self.app.config.is_set("workflow_schedulers_config_file")
        ):
            log.info("No workflow schedulers plugin config file defined, using default scheduler.")
            use_default_scheduler = True
        elif not os.path.exists(config_file):
            log.info(f"Cannot find workflow schedulers plugin config file '{config_file}', using default scheduler.")
            use_default_scheduler = True

        if use_default_scheduler:
            self.__init_default_scheduler()
        else:
            self.DEFAULT_BASE_HANDLER_POOLS = ("workflow-schedulers",)
            plugins_element = load(config_file).getroot()
            self.__init_schedulers_for_element(plugins_element)

        if not self.__handlers_configured and self.__stack_has_pool:
            # Stack has a pool for us so override inherited config and use the pool
            self.__init_handlers()
            self.__handlers_configured = True
        elif use_default_scheduler:
            self._set_default_handler_assignment_methods()

    def __init_default_scheduler(self):
        self.default_scheduler_id = DEFAULT_SCHEDULER_ID
        self.__init_plugin(DEFAULT_SCHEDULER_PLUGIN_TYPE)

    def __init_schedulers_for_element(self, plugins_element: "Element") -> None:
        plugins_kwds = {unicodify(k): unicodify(v) for k, v in plugins_element.items()}
        self.default_scheduler_id = plugins_kwds.get("default", DEFAULT_SCHEDULER_ID)
        for config_element in plugins_element:
            config_element_tag = config_element.tag
            if config_element_tag == "handlers":
                self.__init_handlers(config_element)

                # Determine the default handler(s)
                self.default_handler_id = self._get_default(
                    self.app.config, config_element, list(self.handlers.keys()), required=False
                )
            else:
                plugin_type = config_element_tag
                plugin_element = config_element
                # Configuring a scheduling plugin...
                plugin_kwds = {unicodify(k): unicodify(v) for k, v in plugin_element.items()}
                workflow_scheduler_id = plugin_kwds.get("id", None)
                self.__init_plugin(plugin_type, workflow_scheduler_id, **plugin_kwds)

        if not self.workflow_schedulers:
            raise Exception(EXCEPTION_MESSAGE_NO_SCHEDULERS)
        if self.default_scheduler_id not in self.workflow_schedulers:
            raise Exception(EXCEPTION_MESSAGE_NO_DEFAULT_SCHEDULER % self.default_scheduler_id)
        if (
            self.app.config.parallelize_workflow_scheduling_within_histories
            and not self.deterministic_handler_assignment
        ):
            if self.__handlers_configured:
                # There's an explicit configuration, the admin should fix it.
                raise Exception(EXCEPTION_MESSAGE_SERIALIZE)
            else:
                log.warning(EXCEPTION_MESSAGE_SERIALIZE)

    def __init_handlers(self, config_element=None):
        assert not self.__handlers_configured
        handling_config_dict = ConfiguresHandlers.xml_to_dict(self.app.config, config_element)
        self._init_handler_assignment_methods(handling_config_dict)
        self._init_handlers(handling_config_dict)
        if not self.handler_assignment_methods_configured:
            self._set_default_handler_assignment_methods()
        else:
            self.app.application_stack.init_job_handling(self)
        log.info("Workflow scheduling handler assignment method(s): %s", ", ".join(self.handler_assignment_methods))
        for tag, handlers in [(t, h) for t, h in self.handlers.items() if isinstance(h, list)]:
            log.info("Tag [%s] handlers: %s", tag, ", ".join(handlers))
        self.__handlers_configured = True

    def __init_plugin(self, plugin_type: str, workflow_scheduler_id: Union[str, None] = None, **kwds) -> None:
        workflow_scheduler_id = workflow_scheduler_id or self.default_scheduler_id

        if workflow_scheduler_id in self.workflow_schedulers:
            raise Exception(EXCEPTION_MESSAGE_DUPLICATE_SCHEDULERS % workflow_scheduler_id)

        workflow_scheduler = self.__plugin_classes[plugin_type](**kwds)
        self.workflow_schedulers[workflow_scheduler_id] = workflow_scheduler
        if isinstance(workflow_scheduler, galaxy.workflow.schedulers.ActiveWorkflowSchedulingPlugin):
            self.active_workflow_schedulers[workflow_scheduler_id] = workflow_scheduler

    def __start_request_monitor(self):
        self.request_monitor = WorkflowRequestMonitor(self.app, self)
        self.app.application_stack.register_postfork_function(self.request_monitor.start)


class WorkflowRequestMonitor(Monitors):

    def __init__(self, app: "MinimalManagerApp", workflow_scheduling_manager: WorkflowSchedulingManager) -> None:
        self.app = app
        self.workflow_scheduling_manager = workflow_scheduling_manager
        self._init_monitor_thread(
            name="WorkflowRequestMonitor.monitor_thread", target=self.__monitor, config=app.config
        )
        self.invocation_grabber = None
        self.update_time_tracking_dict: dict[int, datetime] = {}
        backfill_seconds = (
            min(app.config.maximum_workflow_invocation_duration, DEFAULT_SCHEDULER_BACKFILL_SECONDS)
            if app.config.maximum_workflow_invocation_duration > 0
            else DEFAULT_SCHEDULER_BACKFILL_SECONDS
        )
        self.timedelta = timedelta(seconds=backfill_seconds)
        self_handler_tags = set(self.app.job_config.self_handler_tags)
        self_handler_tags.add(self.workflow_scheduling_manager.default_handler_id)
        handler_assignment_method = InvocationGrabber.get_grabbable_handler_assignment_method(
            self.workflow_scheduling_manager.handler_assignment_methods
        )
        if handler_assignment_method:
            self.invocation_grabber = InvocationGrabber(
                app=app,
                handler_assignment_method=handler_assignment_method,
                max_grab=self.workflow_scheduling_manager.handler_max_grab,
                self_handler_tags=self_handler_tags,
                handler_tags=self_handler_tags,
            )

    def ready_to_schedule_more(self, invocation: model.WorkflowInvocation):
        # Improve reactivity of scheduling using the history update_time as a heuristic.
        # If there wasn't a change in the history we're unlikely to be able to make more progress.
        if invocation.id not in self.update_time_tracking_dict:
            return True
        else:
            last_schedule_time = self.update_time_tracking_dict[invocation.id]
            last_history_update_time = invocation.history.update_time
            do_schedule = last_history_update_time > last_schedule_time
            if not do_schedule and (
                invocation_step_update_time := invocation.get_last_workflow_invocation_step_update_time()
            ):
                do_schedule = invocation_step_update_time > last_schedule_time
            if not do_schedule and (datetime.now() - last_schedule_time) > self.timedelta:
                # If we haven't scheduled in a while, schedule anyway.
                log.debug(
                    "Scheduling workflow invocation [%s] after %s seconds without scheduling.",
                    invocation.id,
                    (datetime.now() - last_schedule_time).total_seconds(),
                )
                do_schedule = True
            return do_schedule

    def __monitor(self):
        to_monitor = self.workflow_scheduling_manager.active_workflow_schedulers
        while self.monitor_running:
            try:
                if self.invocation_grabber:
                    self.invocation_grabber.grab_unhandled_items()

                monitor_step_timer = self.app.execution_timer_factory.get_timer(
                    "internal.galaxy.workflows.scheduling_manager.monitor_step",
                    "Workflow scheduling manager monitor step complete.",
                )
                for workflow_scheduler_id, workflow_scheduler in to_monitor.items():
                    if not self.monitor_running:
                        return

                    self.__schedule(workflow_scheduler_id, workflow_scheduler)
                log.trace(monitor_step_timer.to_str())
            except Exception:
                log.exception("An exception occured scheduling while scheduling workflows")
            self._monitor_sleep(self.app.config.workflow_monitor_sleep)

    def __schedule(self, workflow_scheduler_id, workflow_scheduler):
        invocation_ids = self.__active_invocation_ids(workflow_scheduler_id)
        for invocation_id in invocation_ids:
            log.debug("Attempting to schedule workflow invocation [%s]", invocation_id)
            self.__attempt_schedule(invocation_id, workflow_scheduler)
            if not self.monitor_running:
                return

    def __attempt_materialize(self, workflow_invocation: model.WorkflowInvocation, session: Session) -> bool:
        try:
            inputs_to_materialize = workflow_invocation.inputs_requiring_materialization()
            for input_to_materialize in inputs_to_materialize:
                hda = input_to_materialize.hda
                user = RequestUser(user_id=workflow_invocation.history.user_id)
                task_request = MaterializeDatasetInstanceTaskRequest(
                    user=user,
                    history_id=workflow_invocation.history.id,
                    source="hda",
                    content=hda.id,
                )
                materialized_okay = self.app.hda_manager.materialize(task_request, session, in_place=True)
                if not materialized_okay:
                    workflow_invocation.fail()
                    assert input_to_materialize.input_dataset.workflow_step
                    workflow_invocation.add_message(
                        InvocationFailureDatasetFailed(
                            workflow_step_id=input_to_materialize.input_dataset.workflow_step.id,
                            reason=FailureReason.dataset_failed,
                            hda_id=hda.id,
                        )
                    )
                    session.add(workflow_invocation)
                    session.commit()
                    return False

            # place back into ready and let it proceed normally on next iteration?
            workflow_invocation.set_state(model.WorkflowInvocation.states.READY)
            session.add(workflow_invocation)
            session.commit()
            return True
        except Exception as e:
            log.exception(f"Failed to materialize dataset for workflow {workflow_invocation.id} - {e}")
            workflow_invocation.fail()
            failure = InvocationUnexpectedFailure(reason=FailureReason.unexpected_failure, details=str(e))
            workflow_invocation.add_message(failure)
            session.add(workflow_invocation)
            session.commit()
        return False

    def __attempt_schedule(self, invocation_id, workflow_scheduler):
        with self.app.model.context() as session:
            check_database_connection(session)
            workflow_invocation = session.get(model.WorkflowInvocation, invocation_id)
            if workflow_invocation.state == workflow_invocation.states.REQUIRES_MATERIALIZATION:
                if not self.__attempt_materialize(workflow_invocation, session):
                    return None
                if self.app.config.workflow_scheduling_separate_materialization_iteration:
                    return None
            try:
                if workflow_invocation.state == workflow_invocation.states.CANCELLING:
                    workflow_invocation.cancel_invocation_steps()
                    workflow_invocation.mark_cancelled()
                    session.commit()
                    self.update_time_tracking_dict.pop(invocation_id, None)
                    return False

                if not workflow_invocation or not workflow_invocation.active:
                    self.update_time_tracking_dict.pop(invocation_id, None)
                    return False

                # This ensures we're only ever working on the 'first' active
                # workflow invocation in a given history, to force sequential
                # activation.
                if self.app.config.history_local_serial_workflow_scheduling:
                    for i in workflow_invocation.history.workflow_invocations:
                        if i.active and i.id < workflow_invocation.id:
                            return False
                if self.ready_to_schedule_more(workflow_invocation):
                    self.update_time_tracking_dict[invocation_id] = datetime.now()
                    workflow_scheduler.schedule(workflow_invocation)
                    log.debug("Workflow invocation [%s] scheduled", invocation_id)
            except Exception:
                self.update_time_tracking_dict.pop(invocation_id, None)
                # TODO: eventually fail this - or fail it right away?
                log.exception("Exception raised while attempting to schedule workflow request.")
                return False

        # A workflow was obtained and scheduled...
        return True

    def __active_invocation_ids(self, scheduler_id):
        handler = self.app.config.server_name
        return model.WorkflowInvocation.poll_active_workflow_ids(
            self.app.model.engine,
            scheduler=scheduler_id,
            handler=handler,
        )

    def start(self):
        self.monitor_thread.start()

    def shutdown(self):
        self.shutdown_monitor()
