import json
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Any
from xml.etree.ElementTree import ElementTree

import phylodata
from phylodata.data_types import ModelType
from phylodata.process.utils.beast2_xml_utils import get_attribute


class Beast2PackageParser(ABC):
    """An abstract base class for parsing BEAST 2 packages."""

    def is_used(self, beast2_xml: ElementTree) -> bool:
        """Returns true if the package is detected in the given BEAST2 xml.
        Checks the top-level namespaces and the spec attribute of all elements."""
        root = beast2_xml.getroot()
        if root is None:
            return False

        # check namespaces

        namespaces = get_attribute(root, "namespace", "")
        namespaces = namespaces.split(":")

        package_namespace = self.get_namespace()

        if namespaces and any(
            namespace == package_namespace
            or namespace.startswith(f"{package_namespace}.")
            for namespace in namespaces
        ):
            return True

        # check maps

        elements_with_map = root.findall(".//map")
        if any(
            elem.text and elem.text.startswith(package_namespace)
            for elem in elements_with_map
        ):
            return True

        # check specs

        elements_with_spec = root.findall(".//*[@spec]")
        spec_values = [get_attribute(elem, "spec") for elem in elements_with_spec]

        for spec in spec_values:
            if not spec:
                continue

            # for the following cases: assume package_namespace = "babel.distribution"

            if spec == package_namespace:
                return True

            if spec.startswith(f"{package_namespace}."):
                # e.g. spec = "babel.distribution.Normal"
                return True

            if any(
                f"{namespace}.{spec}" == package_namespace for namespace in namespaces
            ):
                # e.g. namespace = "babel" and spec = "distribution"
                return True

            if any(
                f"{namespace}.{spec}".startswith(f"{package_namespace}.")
                for namespace in namespaces
            ):
                # e.g. namespace = "babel" and spec = "distribution.Normal"
                return True

        return False

    @abstractmethod
    def get_name(self) -> str:
        """Returns a human-readable name of the BEAST 2 package."""
        raise NotImplementedError

    @abstractmethod
    def get_type(self) -> ModelType:
        """Returns the type of the BEAST 2 package."""
        raise NotImplementedError

    @abstractmethod
    def get_namespace(self) -> str:
        """Returns the namespace of the BEAST 2 package."""
        raise NotImplementedError

    def get_parameters(self, beast2_xml: ElementTree) -> dict[str, Any]:
        """Extracts potential package parameters from the BEAST 2 xml configuration."""
        return {}


class AutoGeneratedBeast2PackageParser(Beast2PackageParser):
    def __init__(self, name: str, type: str, namespace: str):
        self.name = name
        self.type = type
        self.namespace = namespace

    def get_name(self) -> str:
        return self.name

    def get_type(self) -> ModelType:
        return next(t for t in ModelType if t.value == self.type)

    def get_namespace(self) -> str:
        return self.namespace


def load_package_parsers_from_json() -> list[Beast2PackageParser]:
    """Dynamically loads BEAST 2 package parsers using the models JSON file.
    This is useful to sync the models detected by the PhyloData python library
    with the models presented on the PhyloData website."""

    with open(
        Path(phylodata.__path__[0]) / "process" / "evolutionary_model" / "models.json",
        "r",
    ) as handle:  # type: ignore
        json_models = json.loads(handle.read())

        return [
            AutoGeneratedBeast2PackageParser(
                name=model["name"],
                type=model["model"],
                namespace=model["namespaces"][0],
            )
            for model in json_models
        ]

        return []


BEAST2_PACKAGE_PARSERS: list[Beast2PackageParser] = load_package_parsers_from_json()
BEAST2_PACKAGE_PARSERS_PER_NAME: dict[str, Beast2PackageParser] = {
    parser.get_name(): parser for parser in BEAST2_PACKAGE_PARSERS
}
