import dataclasses
import functools
from collections.abc import Callable
from typing import TYPE_CHECKING, Any

import numpy
import numpy.typing as npt
import stevedore

if TYPE_CHECKING:  # pragma: no cover
    from cogent3.core.annotation_db import AnnotationDbLoaderBase
    from cogent3.core.seq_storage import AlignedSeqsDataABC, SeqsDataABC
    from cogent3.core.tree import PhyloNode
    from cogent3.evolve.fast_distance import DistanceMatrix
    from cogent3.format.sequence import SequenceWriterBase
    from cogent3.parse.sequence import SequenceParserBase

NumpyIntArrayType = npt.NDArray[numpy.integer]

# Entry point for plugins to register themselves as hooks
HOOK_ENTRY_POINT = "cogent3.hook"


def get_quick_tree_hook(
    *,
    name: str | None = None,
) -> Callable[["DistanceMatrix"], "PhyloNode"] | None:
    """returns app instance registered for quick_tree

    Parameters
    ----------
    name
        name of package to get the app from

    Notes
    -----
    The app must take a DistanceMatrix as input and return a PhyloNode
    """
    import cogent3

    mgr = stevedore.hook.HookManager(
        namespace=HOOK_ENTRY_POINT,
        name="quick_tree",
        invoke_on_load=False,
    )
    if name != "cogent3":
        for extension in mgr.extensions:
            if name is None or extension.module_name.startswith(name):
                return extension.plugin()

        if name:
            msg = f"Could not find quick_tree plugin for {name!r}"
            raise ValueError(msg)

    return cogent3.get_app("quick_tree")


def get_count_kmers_hook(
    *,
    name: str | None = None,
    **kwargs: Any,
) -> Callable[..., NumpyIntArrayType] | None:
    """returns app instance registered for count_kmers matching name

    Parameters
    ----------
    name
        name of package to get the app from
    **kwargs
        additional arguments to pass to the app

    Notes
    -----
    The app constructor must take keyword arguments and the
    instance must take a list of strings as input and return
    a numpy array of integers.

    Returns
    -------
    Returns the first found app instance, or one matching name
    or None if no app found.
    """
    mgr = stevedore.hook.HookManager(
        namespace=HOOK_ENTRY_POINT,
        name="count_kmers",
        invoke_on_load=False,
    )
    if name != "cogent3":
        for extension in mgr.extensions:
            if name and extension.module_name.startswith(name):
                return extension.plugin(**kwargs)
            # return first one if no name specified
            return extension.plugin(**kwargs)

        if name:
            msg = f"Could not find count_kmers plugin for {name!r}"
            raise ValueError(msg)

    # we don't have a cogent3 app for this, we just use the builtin function
    return None


# registry for parsing sequence file formats
SEQ_PARSER_ENTRY_POINT = "cogent3.parse.sequence"


@functools.cache
def get_seq_format_parser_plugin(
    *,
    format_name: str | None = None,
    file_suffix: str | None = None,
    unaligned_seqs: bool = True,
) -> "SequenceParserBase":
    """returns sequence format parser plugin

    Parameters
    ----------
    format_name
        name of sequence format
    file_suffix
        suffix of file to parse
    unaligned_seqs
        whether parser is for unaligned sequences

    Notes
    -----
    We default to third-party plugins if they are available, otherwise we
    use a built-in parser.
    """

    mgr = stevedore.ExtensionManager(
        namespace=SEQ_PARSER_ENTRY_POINT,
        invoke_on_load=True,
    )
    built_in = None
    for ext in mgr.extensions:
        plugin = ext.plugin()
        if file_suffix in plugin.supported_suffixes or plugin.name == format_name:
            if ext.module_name.startswith("cogent3."):
                built_in = plugin
                continue

            supports_role = (plugin.supports_unaligned and unaligned_seqs) or (
                plugin.supports_aligned and not unaligned_seqs
            )

            if not supports_role:
                out_type = "unaligned seqs" if unaligned_seqs else "aligned seqs"
                msg = f"{plugin.name} does not support parsing {out_type}"
                raise ValueError(msg)
            return plugin

    if built_in:
        # if we have a built-in plugin, return it
        return built_in

    msg = f"Unknown parser for format {format_name!r} or file suffix {file_suffix!r}"
    raise ValueError(msg)


# registry for writing sequence file formats
SEQ_FORMAT_ENTRY_POINT = "cogent3.format.sequence"


@functools.cache
def get_seq_format_writer_plugin(
    *,
    format_name: str | None = None,
    file_suffix: str | None = None,
    unaligned_seqs: bool = True,
) -> "SequenceWriterBase":
    """returns sequence format writer

    Parameters
    ----------
    format_name
        name of sequence format
    file_suffix
        suffix of file to parse
    unaligned_seqs
        whether format is for unaligned sequences

    Notes
    -----
    We default to third-party plugins if they are available, otherwise we
    use a built-in parser.
    """

    mgr = stevedore.ExtensionManager(
        namespace=SEQ_FORMAT_ENTRY_POINT,
        invoke_on_load=True,
    )
    built_in = None
    for ext in mgr.extensions:
        plugin = ext.plugin()
        if file_suffix in plugin.supported_suffixes or plugin.name == format_name:
            if ext.module_name.startswith("cogent3."):
                built_in = plugin
                continue
            supports_role = (plugin.supports_unaligned and unaligned_seqs) or (
                plugin.supports_aligned and not unaligned_seqs
            )

            if not supports_role:
                out_type = "unaligned seqs" if unaligned_seqs else "aligned seqs"
                msg = f"{plugin.name} does not support writing {out_type}"
                raise ValueError(msg)
            return plugin

    if built_in:
        # if we have a built-in plugin, return it
        return built_in

    msg = f"Unknown writer for format {format_name!r} or file suffix {file_suffix!r}"
    raise ValueError(msg)


# sequence storage drivers
UNALIGNED_SEQ_STORAGE_ENTRY_POINT = "cogent3.storage.unaligned_seqs"


def _get_driver(namespace: str, storage_backend: str) -> stevedore.driver.DriverManager:
    try:
        mgr = stevedore.driver.DriverManager(
            namespace=namespace,
            name=storage_backend,
            invoke_on_load=False,
        )
    except stevedore.exception.NoMatches as err:
        msg = f"Invalid storage backend {storage_backend!r}"
        raise ValueError(msg) from err
    return mgr


def get_unaligned_storage_driver(
    storage_backend: str | None,
) -> type["SeqsDataABC"]:
    """returns unaligned sequence storage driver

    Parameters
    ----------
    storage_backend
        name of storage plugin to use
    """
    if not storage_backend:
        return _STORAGE_DEFAULT.unaligned

    mgr = _get_driver(UNALIGNED_SEQ_STORAGE_ENTRY_POINT, storage_backend)
    return mgr.extensions[0].plugin if mgr.extensions else _STORAGE_DEFAULT.unaligned


ALIGNED_SEQ_STORAGE_ENTRY_POINT = "cogent3.storage.aligned_seqs"


def get_aligned_storage_driver(
    storage_backend: str | None,
) -> type["AlignedSeqsDataABC"]:
    """returns aligned sequence storage driver

    Parameters
    ----------
    storage_backend
        name of storage plugin to use
    """
    if not storage_backend:
        return _STORAGE_DEFAULT.aligned

    mgr = _get_driver(ALIGNED_SEQ_STORAGE_ENTRY_POINT, storage_backend)
    return mgr.extensions[0].plugin if mgr.extensions else _STORAGE_DEFAULT.aligned


@dataclasses.dataclass
class DefaultStorageDrivers:
    _unaligned: type | None = dataclasses.field(init=False, default=None)
    _aligned: type | None = dataclasses.field(init=False, default=None)

    @property
    def unaligned(self) -> type["SeqsDataABC"]:
        if self._unaligned is None:
            from cogent3.core.seq_storage import SeqsData

            self._unaligned = SeqsData

        return self._unaligned

    @unaligned.setter
    def unaligned(self, driver: type["SeqsDataABC"] | None) -> None:
        self._unaligned = driver

    @property
    def aligned(self) -> type["AlignedSeqsDataABC"]:
        if self._aligned is None:
            from cogent3.core.alignment import AlignedSeqsData

            self._aligned = AlignedSeqsData

        return self._aligned

    @aligned.setter
    def aligned(self, driver: type["AlignedSeqsDataABC"] | None) -> None:
        self._aligned = driver


_STORAGE_DEFAULT = DefaultStorageDrivers()


def set_storage_defaults(
    *,
    unaligned_seqs: str | None = None,
    aligned_seqs: str | None = None,
    reset: bool = False,
) -> None:
    """set default values for storage of unaligned and aligned seqs data

    Parameters
    ----------
    unaligned_seqs
        name of storage backend for unaligned sequences
    aligned_seqs
        name of storage backend for aligned sequences
    reset
        resets defaults to cogent3 objects
    """
    if reset:
        _STORAGE_DEFAULT.unaligned = None
        _STORAGE_DEFAULT.aligned = None
        return

    if not any([unaligned_seqs, aligned_seqs]):
        return

    if unaligned_seqs:
        _STORAGE_DEFAULT.unaligned = get_unaligned_storage_driver(
            storage_backend=unaligned_seqs,
        )

    if aligned_seqs:
        _STORAGE_DEFAULT.aligned = get_aligned_storage_driver(
            storage_backend=aligned_seqs,
        )


# Entry point for plugins to register themselves as apps
APP_ENTRY_POINT = "cogent3.app"

# private global to hold an ExtensionManager instance for apps
# Note that this is relied on for the tests
__apps = None


def get_app_manager() -> stevedore.ExtensionManager:
    """
    Lazy load a stevedore ExtensionManager to collect apps.
    """
    global __apps  # noqa: PLW0603
    if not __apps:
        __apps = stevedore.ExtensionManager(
            namespace=APP_ENTRY_POINT,
            invoke_on_load=False,
        )

    return __apps


# Entry point for annotation database loaders
ANNOTATION_LOADER_ENTRY_POINT = "cogent3.load.annotations"


@functools.cache
def get_annotation_loader_plugin(
    *,
    storage_backend: str | None = None,
    file_suffix: str,
) -> "AnnotationDbLoaderBase":
    """Returns annotation database loader plugin.

    Parameters
    ----------
    storage_backend
        Unique identifier for the storage backend plugin. If None, selects
        first compatible third-party plugin, falling back to cogent3 built-in.
    file_suffix
        File suffix to match against plugin's supported formats

    Returns
    -------
    AnnotationDbLoaderBase instance

    Raises
    ------
    ValueError
        If storage_backend is specified but doesn't identify a plugin, if the
        selected plugin doesn't support the file suffix, or if no compatible
        plugin is found.

    Notes
    -----
    When storage_backend is specified, it uniquely identifies a loader plugin.
    When storage_backend is None, third-party plugins are preferred over
    built-in cogent3 loaders.
    """
    if storage_backend:
        # Explicit backend specified - use DriverManager
        try:
            mgr = stevedore.driver.DriverManager(
                namespace=ANNOTATION_LOADER_ENTRY_POINT,
                name=storage_backend,
                invoke_on_load=True,
            )
        except stevedore.exception.NoMatches as err:
            msg = f"Invalid storage backend {storage_backend!r}"
            raise ValueError(msg) from err

        plugin = mgr.extensions[0].plugin()

        # Validate that plugin supports the file format
        if file_suffix not in plugin.supported_suffixes:
            msg = (
                f"Storage backend {storage_backend!r} does not support file "
                f"format {file_suffix!r}. Supported formats: "
                f"{', '.join(sorted(plugin.supported_suffixes))}"
            )
            raise ValueError(msg)

        return plugin

    # No backend specified - find first compatible plugin, prefer third-party
    mgr = stevedore.ExtensionManager(
        namespace=ANNOTATION_LOADER_ENTRY_POINT,
        invoke_on_load=True,
    )

    built_in = None
    for ext in mgr.extensions:
        plugin = ext.plugin()

        if file_suffix in plugin.supported_suffixes:
            # Found a match - defer built-in, return third-party immediately
            if ext.module_name.startswith("cogent3."):
                built_in = plugin
                continue
            return plugin

    # Fall back to built-in if found
    if built_in:
        return built_in

    # No compatible plugin found
    msg = f"No loader found for file suffix {file_suffix!r}"
    raise ValueError(msg)
