"""
Optimade support.
"""

import logging
import sys
from collections import namedtuple
from os.path import join
from typing import Dict, List, Optional, Union, Set
from urllib.parse import urlparse

import requests
from tqdm import tqdm

# from retrying import retry

from pymatgen.core.periodic_table import DummySpecies
from pymatgen.core.structure import Structure
from pymatgen.util.provenance import StructureNL


# TODO: importing optimade-python-tool's data structures will make more sense
Provider = namedtuple("Provider", ["name", "base_url", "description", "homepage", "prefix"])

_logger = logging.getLogger(__name__)
_handler = logging.StreamHandler(sys.stdout)
_logger.addHandler(_handler)
_logger.setLevel(logging.WARNING)


class OptimadeRester:
    """
    Class to call OPTIMADE-compliant APIs, see https://optimade.org and [1].

    This class is ready to use but considered in-development and subject to change.

    [1] Andersen, C.W., *et al*.
        OPTIMADE, an API for exchanging materials data.
        Sci Data 8, 217 (2021). https://doi.org/10.1038/s41597-021-00974-z

    """

    # regenerate on-demand from official providers.json using OptimadeRester.refresh_aliases()
    # these aliases are provided as a convenient shortcut for users of the OptimadeRester class
    aliases: Dict[str, str] = {
        "aflow": "http://aflow.org/API/optimade/",
        "cod": "https://www.crystallography.net/cod/optimade",
        "mcloud.2dstructures": "https://aiida.materialscloud.org/2dstructures/optimade",
        "mcloud.2dtopo": "https://aiida.materialscloud.org/2dtopo/optimade",
        "mcloud.curated-cofs": "https://aiida.materialscloud.org/curated-cofs/optimade",
        "mcloud.li-ion-conductors": "https://aiida.materialscloud.org/li-ion-conductors/optimade",
        "mcloud.optimade-sample": "https://aiida.materialscloud.org/optimade-sample/optimade",
        "mcloud.pyrene-mofs": "https://aiida.materialscloud.org/pyrene-mofs/optimade",
        "mcloud.scdm": "https://aiida.materialscloud.org/autowannier/optimade",
        "mcloud.sssp": "https://aiida.materialscloud.org/sssplibrary/optimade",
        "mcloud.stoceriaitf": "https://aiida.materialscloud.org/stoceriaitf/optimade",
        "mcloud.tc-applicability": "https://aiida.materialscloud.org/tc-applicability/optimade",
        "mcloud.threedd": "https://aiida.materialscloud.org/3dd/optimade",
        "mp": "https://optimade.materialsproject.org",
        "mpds": "https://api.mpds.io",
        "nmd": "https://nomad-lab.eu/prod/rae/optimade/",
        "odbx": "https://optimade.odbx.science",
        "omdb.omdb_production": "http://optimade.openmaterialsdb.se",
        "oqmd": "http://oqmd.org/optimade/",
        "tcod": "https://www.crystallography.net/tcod/optimade",
    }

    # The set of OPTIMADE fields that are required to define a `pymatgen.core.Structure`
    mandatory_response_fields: Set[str] = {"lattice_vectors", "cartesian_site_positions", "species", "species_at_sites"}

    def __init__(self, aliases_or_resource_urls: Optional[Union[str, List[str]]] = None, timeout: int = 5):
        """
        OPTIMADE is an effort to provide a standardized interface to retrieve information
        from many different materials science databases.

        This is a client to retrieve structures from OPTIMADE v1 compliant endpoints. It
        does not yet support all features of the OPTIMADE v1 specification but is intended
        as a way to quickly search an endpoint in a way familiar to users of pymatgen without
        needing to know the full OPTIMADE specification.

        For advanced usage, please see the OPTIMADE documentation at optimade.org and
        consider calling the APIs directly.

        For convenience, known OPTIMADE endpoints have been given aliases in pymatgen to save
        typing the full URL. The current list of aliases is:

        aflow, cod, mcloud.sssp, mcloud.2dstructures, mcloud.2dtopo, mcloud.tc-applicability,
        mcloud.threedd, mcloud.scdm, mcloud.curated-cofs, mcloud.optimade-sample, mcloud.stoceriaitf,
        mcloud.pyrene-mofs, mcloud.li-ion-conductors, mp, odbx, omdb.omdb_production, oqmd, tcod

        To refresh this list of aliases, generated from the current list of OPTIMADE providers
        at optimade.org, call the refresh_aliases() method.

        Args:
            aliases_or_resource_urls: the alias or structure resource URL or a list of
            aliases or resource URLs, if providing the resource URL directly it should not
            be an index, this interface can only currently access the "v1/structures"
            information from the specified resource URL
            timeout: number of seconds before an attempted request is abandoned, a good
            timeout is useful when querying many providers, some of which may be offline
        """

        # TODO: maybe we should use the nice pydantic models from optimade-python-tools
        #  for response validation, and use the Lark parser for filter validation
        self.session = requests.Session()
        self._timeout = timeout  # seconds

        if isinstance(aliases_or_resource_urls, str):
            aliases_or_resource_urls = [aliases_or_resource_urls]

        # this stores a dictionary with keys provider id (in the same format as the aliases)
        # and values as the corresponding URL
        self.resources = {}

        if not aliases_or_resource_urls:
            aliases_or_resource_urls = list(self.aliases.keys())
            _logger.warning(
                "Connecting to all known OPTIMADE providers, this will be slow. Please connect to only the "
                f"OPTIMADE providers you want to query. Choose from: {', '.join(self.aliases.keys())}"
            )

        for alias_or_resource_url in aliases_or_resource_urls:

            if alias_or_resource_url in self.aliases:
                self.resources[alias_or_resource_url] = self.aliases[alias_or_resource_url]

            elif self._validate_provider(alias_or_resource_url):

                # TODO: unclear what the key should be here, the "prefix" is for the root provider,
                # may need to walk back to the index for the given provider to find the correct identifier

                self.resources[alias_or_resource_url] = alias_or_resource_url

            else:
                _logger.error(f"The following is not a known alias or a valid url: {alias_or_resource_url}")

        self._providers = {url: self._validate_provider(provider_url=url) for url in self.resources.values()}

    def __repr__(self):
        return f"OptimadeRester connected to: {', '.join(self.resources.values())}"

    def __str__(self):
        return self.describe()

    def describe(self):
        """
        Provides human-readable information about the resources being searched by the OptimadeRester.
        """
        provider_text = "\n".join(map(str, (provider for provider in self._providers.values() if provider)))
        description = f"OptimadeRester connected to:\n{provider_text}"
        return description

    # @retry(stop_max_attempt_number=3, wait_random_min=1000, wait_random_max=2000)
    def _get_json(self, url):
        """
        Retrieves JSON, will attempt to (politely) try again on failure subject to a
        random delay and a maximum number of attempts.
        """
        return self.session.get(url, timeout=self._timeout).json()

    @staticmethod
    def _build_filter(
        elements: Union[str, List[str]] = None,
        nelements: int = None,
        nsites: int = None,
        chemical_formula_anonymous: str = None,
        chemical_formula_hill: str = None,
    ):
        """
        Convenience method to build an OPTIMADE filter.
        """

        filters = []

        if elements:
            if isinstance(elements, str):
                elements = [elements]
            elements_str = ", ".join([f'"{el}"' for el in elements])
            filters.append(f"(elements HAS ALL {elements_str})")

        if nsites:
            if isinstance(nsites, (list, tuple)):
                filters.append(f"(nsites>={min(nsites)} AND nsites<={max(nsites)})")
            else:
                filters.append(f"(nsites={int(nsites)})")

        if nelements:
            if isinstance(nelements, (list, tuple)):
                filters.append(f"(nelements>={min(nelements)} AND nelements<={max(nelements)})")
            else:
                filters.append(f"(nelements={int(nelements)})")

        if chemical_formula_anonymous:
            filters.append(f'(chemical_formula_anonymous="{chemical_formula_anonymous}")')

        if chemical_formula_hill:
            filters.append(f'(chemical_formula_hill="{chemical_formula_anonymous}")')

        return " AND ".join(filters)

    def get_structures(
        self,
        elements: Union[List[str], str] = None,
        nelements: int = None,
        nsites: int = None,
        chemical_formula_anonymous: str = None,
        chemical_formula_hill: str = None,
    ) -> Dict[str, Dict[str, Structure]]:
        """
        Retrieve Structures from OPTIMADE providers.

        Not all functionality of OPTIMADE is currently exposed in this convenience method. To
        use a custom filter, call get_structures_with_filter().

        Args:
            elements: List of elements
            nelements: Number of elements, e.g. 4 or [2, 5] for the range >=2 and <=5
            nsites: Number of sites, e.g. 4 or [2, 5] for the range >=2 and <=5
            chemical_formula_anonymous: Anonymous chemical formula
            chemical_formula_hill: Chemical formula following Hill convention

        Returns: Dict of (Dict Structures keyed by that database's id system) keyed by provider
        """

        optimade_filter = self._build_filter(
            elements=elements,
            nelements=nelements,
            nsites=nsites,
            chemical_formula_anonymous=chemical_formula_anonymous,
            chemical_formula_hill=chemical_formula_hill,
        )

        return self.get_structures_with_filter(optimade_filter)

    def get_snls(
        self,
        elements: Union[List[str], str] = None,
        nelements: int = None,
        nsites: int = None,
        chemical_formula_anonymous: str = None,
        chemical_formula_hill: str = None,
        additional_response_fields: Union[str, List[str], Set[str]] = None,
    ) -> Dict[str, Dict[str, StructureNL]]:
        """
        Retrieve StructureNL from OPTIMADE providers.

        A StructureNL is an object provided by pymatgen which combines Structure with
        associated metadata, such as the URL is was downloaded from and any additional namespaced
        data.

        Not all functionality of OPTIMADE is currently exposed in this convenience method. To
        use a custom filter, call get_structures_with_filter().

        Args:
            elements: List of elements
            nelements: Number of elements, e.g. 4 or [2, 5] for the range >=2 and <=5
            nsites: Number of sites, e.g. 4 or [2, 5] for the range >=2 and <=5
            chemical_formula_anonymous: Anonymous chemical formula
            chemical_formula_hill: Chemical formula following Hill convention
            additional_response_fields: Any additional fields desired from the OPTIMADE API,
            these will be stored under the `'_optimade'` key in each `StructureNL.data` dictionary.

        Returns: Dict of (Dict of StructureNLs keyed by that database's id system) keyed by provider
        """

        optimade_filter = self._build_filter(
            elements=elements,
            nelements=nelements,
            nsites=nsites,
            chemical_formula_anonymous=chemical_formula_anonymous,
            chemical_formula_hill=chemical_formula_hill,
        )

        return self.get_snls_with_filter(optimade_filter, additional_response_fields=additional_response_fields)

    def get_structures_with_filter(self, optimade_filter: str) -> Dict[str, Dict[str, Structure]]:
        """
        Get structures satisfying a given OPTIMADE filter.

        Args:
            filter: An OPTIMADE-compliant filter

        Returns: Dict of Structures keyed by that database's id system
        """

        all_snls = self.get_snls_with_filter(optimade_filter)
        all_structures = {}

        for identifier, snls_dict in all_snls.items():
            all_structures[identifier] = {k: snl.structure for k, snl in snls_dict.items()}

        return all_structures

    def get_snls_with_filter(
        self,
        optimade_filter: str,
        additional_response_fields: Union[str, List[str], Set[str]] = None,
    ) -> Dict[str, Dict[str, StructureNL]]:
        """
        Get structures satisfying a given OPTIMADE filter.

        Args:
            filter: An OPTIMADE-compliant filter

        Returns: Dict of Structures keyed by that database's id system
        """

        all_snls = {}

        fields = self._handle_response_fields(additional_response_fields)

        for identifier, resource in self.resources.items():

            url = join(resource, f"v1/structures?filter={optimade_filter}&response_fields={fields}")

            try:

                json = self._get_json(url)

                structures = self._get_snls_from_resource(json, url, identifier)

                pbar = tqdm(total=json["meta"].get("data_returned", 0), desc=identifier, initial=len(structures))

                # TODO: check spec for `more_data_available` boolean, may simplify this conditional
                if ("links" in json) and ("next" in json["links"]) and (json["links"]["next"]):
                    while "next" in json["links"] and json["links"]["next"]:
                        next_link = json["links"]["next"]
                        if isinstance(next_link, dict) and "href" in next_link: