# SPDX-FileCopyrightText: 2025-2026 Global Computing Lab
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

from __future__ import annotations

import sys

# Although importlib.metadata existed in versions 3.8 and 3.9,
# it was only provisional. The API of importlib.metadata was not
# standardized and made non-provisional until 3.10. So, we
# will depend on importlib_metadata for all versions of Python
# before 3.10 to ensure they all behave the same.
if sys.version_info < (3, 10):
    from importlib_metadata import entry_points
else:
    from importlib.metadata import entry_points

from typing import TYPE_CHECKING

from a4x.orchestration.command import Command
from a4x.orchestration.file import File
from a4x.orchestration.resources import JobspecSettings, SchedulableWork
from a4x.orchestration.site import (
    Directory,
    PersistencyType,
    Scheduler,
    Site,
    StorageType,
    get_default_site,
    set_default_site,
)
from a4x.orchestration.task import Task
from a4x.orchestration.workflow import Workflow

try:
    from a4x.orchestration.version import __version__
except ImportError:
    # Since version.py is built dynamically by hatch-vcs, it may
    # not yet exist for local development. Fall back to a dummy
    # version in this case to avoid an ImportError
    __version__ = "0.0.0-dev"

if TYPE_CHECKING:
    from a4x.orchestration.plugin import Plugin


def _build_factory_for_builtin_plugins() -> dict[str, type[Plugin]]:
    from a4x.orchestration.builtin_plugins.flux import FluxPlugin

    return {"flux": FluxPlugin}


builtin_plugins = _build_factory_for_builtin_plugins()

__wms_plugins = {ep.name: ep for ep in entry_points(group="a4x_wms_plugins")}


def get_plugin_names() -> list[str]:
    """Get the names of all registered plugins.

    :return: a list of the plugin names
    """
    plugin_names = list(builtin_plugins.keys())
    plugin_names.extend(list(__wms_plugins.keys()))
    return plugin_names


def convert_to_plugin(wflow: Workflow, plugin_key: str, ignore_builtin_plugins: bool = False, **kwargs) -> Plugin:
    """Convert the provided workflow to a WMS-specific plugin.

    :param wflow: the workflow to convert
    :param plugin_key: the identifier of the plugin to generate. Must be an element of the output of :code:`get_plugin_names`
    :param ignore_builtin_plugins: if true, do not consider plugins built into A4X-Orchestration directly
    :param kwargs: remaining keyword arguments that are passed as-is to the plugin
    :return: the WMS-specific plugin for the workflow
    """  # noqa: E501
    wms_entrypoint = None
    if not ignore_builtin_plugins and plugin_key in builtin_plugins:
        wms_entrypoint = builtin_plugins[plugin_key]
    else:
        if plugin_key not in __wms_plugins:
            raise KeyError(
                "Cannot convert to a WMS-specific type because there is no plugin called {}".format(plugin_key)
            )
        wms_entrypoint = __wms_plugins[plugin_key].load()
    if get_default_site() is None:
        entrypoint_default_site = wms_entrypoint.get_default_site()
        if entrypoint_default_site is not None:
            set_default_site(entrypoint_default_site)
    wflow.resolve()
    return wms_entrypoint(wflow, **kwargs)


__all__ = [
    "Command",
    "Directory",
    "File",
    "JobspecSettings",
    "PersistencyType",
    "SchedulableWork",
    "Scheduler",
    "Site",
    "StorageType",
    "Task",
    "Workflow",
    "__version__",
    "convert_to_plugin",
    "get_default_site",
    "get_wms_plugin_names",
    "set_default_site",
]
