# Copyright (c) 2017-2021 Digital Asset (Switzerland) GmbH and/or its affiliates. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

from collections import defaultdict
from dataclasses import dataclass
from functools import reduce
from operator import add
from threading import RLock
from types import MappingProxyType
from typing import Any, Collection, Dict, Generic, Iterable, List, Mapping, Optional, TypeVar, Union
import warnings

from ..damlast.daml_lf_1 import Archive, Expr, Package, PackageRef, TypeConName, ValName, _Name
from ..damlast.lookup import validate_template
from ..util.typing import safe_cast, safe_dict_cast

K = TypeVar("K", bound=_Name)
T = TypeVar("T")

warnings.warn(
    "The types of dazl.model.types_store are deprecated", DeprecationWarning, stacklevel=2
)

with warnings.catch_warnings():
    warnings.simplefilter("ignore", DeprecationWarning)

    from .types import (
        ConcreteType,
        PackageIdSet,
        Template,
        TemplateChoice,
        Type,
        TypeReference,
        UnresolvedTypeReference,
    )


class PackageStoreBuilder:
    """
    Convenience class for building up a :class:`PackageStore`.
    """

    def __init__(self):
        warnings.warn(
            "PackageStoreBuilder is deprecated; there is no direct replacement.",
            DeprecationWarning,
            stacklevel=2,
        )

        self._archives = list()  # type: List[Archive]
        self._value_types = dict()  # type: Dict[ValName, Expr]
        self._data_types = dict()  # type: Dict[TypeConName, Type]
        self._templates = dict()  # type: Dict[TypeConName, Template]
        self._expected_package_ids = None  # type: Optional[Collection[str]]

    def add_archive(self, archive: "Archive") -> None:
        self._archives.append(archive)

    def add_type(self, name: "TypeConName", data_type: Type):
        safe_cast(TypeConName, name)
        self._data_types[name] = safe_cast(Type, data_type)

    def add_value(self, name: "ValName", value: "Expr"):
        safe_cast(ValName, name)
        self._value_types[name] = safe_cast(Expr, value)

    def add_template(self, template: Template):
        self._templates[template.data_type.name.con] = safe_cast(Template, template)

    def add_expected_package_ids(self, expected_package_ids: "Collection[str]"):
        self._expected_package_ids = expected_package_ids

    def get_type(self, name: TypeConName) -> "Optional[Type]":
        return self._data_types.get(name)

    def build(self) -> "PackageStore":
        return PackageStore(
            self._archives,
            self._value_types,
            self._data_types,
            self._templates,
            self._expected_package_ids,
        )


class PackageStore:
    """
    A thread-safe store of type information.
    """

    @classmethod
    def empty(cls):
        """
        Create an empty store.
        """
        warnings.warn(
            "PackageStore is deprecated; use PackageLoader/SymbolLookup instead",
            DeprecationWarning,
            stacklevel=2,
        )

        return cls([], {}, {}, {}, None)

    def __init__(
        self,
        archives: "Collection[Archive]",
        value_types: "Dict[ValName, Expr]",
        data_types: "Dict[TypeConName, Type]",
        templates: "Dict[TypeConName, Template]",
        expected_package_ids: "Optional[Collection[str]]" = None,
    ):
        warnings.warn(
            "PackageStore is deprecated; use PackageLoader/SymbolLookup instead",
            DeprecationWarning,
            stacklevel=2,
        )
        self._lock = RLock()
        self._archives = list(archives)
        self._cache = PackageStoreCache(EMPTY_TYPE_CACHE, EMPTY_TYPE_CACHE, EMPTY_TYPE_CACHE)
        self._value_types = safe_dict_cast(ValName, Expr, value_types)
        self._data_types = safe_dict_cast(TypeConName, Type, data_types)
        self._templates = safe_dict_cast(TypeConName, Template, templates)
        self._expected_package_ids = expected_package_ids

    def archives(self) -> "Collection[Archive]":
        """
        Return a copy of the collection of the set of loaded :class:`Archive`s.
        """
        warnings.warn(
            "PackageStore.archives() is deprecated; use SymbolLookup.archives() instead",
            DeprecationWarning,
            stacklevel=2,
        )

        with self._lock:
            return list(self._archives)

    def packages(self) -> "Collection[Package]":
        """
        Return a copy of the collection of the set of loaded :class:`Package`s.
        """
        warnings.warn(
            "PackageStore.packages() is deprecated; use SymbolLookup.archives() instead",
            DeprecationWarning,
            stacklevel=2,
        )

        with self._lock:
            return [a.package for a in self._archives]

    def package_ids(self) -> "Collection[PackageRef]":
        """
        Return a copy of the collection of the set of loaded :class:`Package`s.
        """
        warnings.warn(
            "PackageStore.package_ids() is deprecated; use SymbolLookup.package_ids() instead",
            DeprecationWarning,
            stacklevel=2,
        )

        with self._lock:
            return [a.hash for a in self._archives]

    def expected_package_ids(self) -> "Collection[str]":
        """
        Return package IDs that are expected to be found on the ledger.
        """
        warnings.warn(
            "PackageStore.expected_package_ids() is deprecated; there is no replacement.",
            DeprecationWarning,
            stacklevel=2,
        )
        return self._expected_package_ids

    def register_all(self, other_store: "PackageStore") -> "PackageStore":
        """
        Register all types declared in the other :class:`PackageStore`.

        :param other_store: A package store to copy types, templates, and choices from.
        :return: A reference to this object.
        """
        warnings.warn(
            "PackageStore.register_all() is deprecated; use PackageLoader instead.",
            DeprecationWarning,
            stacklevel=2,
        )

        if self is not other_store:
            with self._lock:
                self._archives.extend(other_store._archives)
                self._value_types.update(other_store._value_types)
                self._data_types.update(other_store._data_types)
                self._templates.update(other_store._templates)
                self._cache = PackageStoreCache(
                    TypeCache.build(self._value_types),
                    TypeCache.build(self._data_types),
                    TypeCache.build(self._templates),
                )
        return self

    def resolve_value_reference(self, value_ref: "ValName") -> "Expr":
        warnings.warn(
            "PackageStore.resolve_value_reference() is deprecated; "
            "use SymbolLookup.value() instead.",
            DeprecationWarning,
            stacklevel=2,
        )

        with self._lock:
            return self._value_types[value_ref]

    def resolve_type_reference(self, template_ref: "Union[TypeConName, TypeReference]") -> Type:
        """
        Resolve a type based on its reference.

        :param template_ref:
            A :class:`TypeReference` that refers to a type.
        :return:
            The :class:`Type` that is referred to by this type reference.
        :raise KeyError:
            If the :class:`TypeReference` does not have a corresponding value in this store.
        """
        warnings.warn(
            "PackageStore.resolve_value_reference() is deprecated; "
            "use SymbolLookup.data_type() instead.",
            DeprecationWarning,
            stacklevel=2,
        )

        if isinstance(template_ref, TypeReference):
            con = template_ref.con
        elif isinstance(template_ref, TypeConName):
            con = template_ref
        else:
            raise ValueError("template_ref must either by a TypeConName or TypeReference")

        with self._lock:
            return self._data_types[con]

    def types(self) -> "Mapping[TypeConName, ConcreteType]":
        warnings.warn(
            'PackageStore.types() is deprecated; use SymbolLookup.data_type("*") instead.',
            DeprecationWarning,
            stacklevel=2,
        )

        with self._lock:
            return MappingProxyType(dict(self._data_types))

    def get_templates_for_packages(
        self, package_ids: "Iterable[PackageRef]"
    ) -> "Collection[Template]":
        """
        Return a collection of :class:`Template` instances from the given set of package IDs.

        :param package_ids:
            The set of package IDs to restrict templates to.
        :return:
            A collection of matching templates, or an empty collection if none match. This method
            never returns ``None``.
        """
        warnings.warn(
            "PackageStore.get_templates_for_packages() is deprecated; "
            'use [SymbolLookup.templates(f"{pkg}:*") for pkg in package_ids] instead.',
            DeprecationWarning,
            stacklevel=2,
        )

        match = []
        for pkg_id in package_ids:
            match.extend(self._cache.templates.lookup(pkg_id, "*"))
        return match

    def resolve_template(
        self, template: "Union[None, str, TypeReference, UnresolvedTypeReference, TypeConName]"
    ) -> Collection[Template]:
        """
        Return a collection of :class:`Template` instances that match the specified template name.

        Some special parameters:
         * If ``"*"`` is passed in, all templates are returned.
         * If ``None`` is passed in, an empty collection is returned.

        :param template:
            A template name or a :class:`TypeReference`.
        :return:
            A collection of matching templates, or an empty collection if none match. This method
            never returns ``None``.
        """
        warnings.warn(
            "PackageStore.resolve_template() is deprecated; use SymbolLookup.template() instead.",
            DeprecationWarning,
            stacklevel=2,
        )

        if isinstance(template, Template):
            # if we were given a Template for some strange reason, just simply return a single-item
            # tuple of that given Template
            return [template]

        package_id, template_name = validate_template(template, allow_deprecated_identifiers=True)
        return self._cache.templates.lookup(package_id, template_name)

    def resolve_template_type(
        self, template: "Union[None, str, TypeReference, UnresolvedTypeReference, TypeConName]"
    ) -> Dict[TypeReference, Type]:
        """
        Return a collection of types that match for the template.

        :param template:
            A template name or a :class:`TypeReference`.
        :return:
            A dictionary of possible matches, or empty if there are no matches. This method never
            returns ``None``.
        """
        warnings.warn(
            "PackageStore.resolve_template() is deprecated; use SymbolLookup.template() instead.",
            DeprecationWarning,
            stacklevel=2,
        )

        with warnings.catch_warnings():
            warnings.simplefilter("ignore", DeprecationWarning)

            return {
                template.data_type.name: template.data_type
                for template in self.resolve_template(template)
            }

    def resolve_choice(self, template: Any, choice: str) -> Dict[TypeReference, TemplateChoice]:
        """
        Return all possible choices for the combination of template identifier and choice name.
        If template is ``'*'`` or ``None``, all choices with the specified name are returned.
        """
        warnings.warn(
            'PackageStore.resolve_choice() is deprecated; use SymbolLookup.template("*") instead, '
            "and lookup the choice within the returned Template.",
            DeprecationWarning,
            stacklevel=2,
        )

        with warnings.catch_warnings():
            warnings.simplefilter("ignore", DeprecationWarning)

            matches = dict()
            for t in self.resolve_template(template):
                for c in t.choices:
                    if c.name == choice:
                        matches[t.data_type.name] = c
            return matches


class PackageProvider:
    """
    Interface to an object that can provide package information.
    """

    def __init__(self):
        warnings.warn(
            "PackageProvider is deprecated; use PackageLookup instead.",
            DeprecationWarning,
            stacklevel=2,
        )

    def get_package_ids(self) -> "PackageIdSet":
        """
        Return the current universe of package IDs.
        """
        raise NotImplementedError

    def fetch_package(self, package_id: "PackageRef") -> bytes:
        """
        Retrieve the bytes that correspond to a package.
        """
        raise NotImplementedError

    def get_all_packages(self) -> "Mapping[PackageRef, bytes]":
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", DeprecationWarning)
            return {pkg_id: self.fetch_package(pkg_id) for pkg_id in self.get_package_ids()}


class MemoryPackageProvider(PackageProvider):
    def __init__(self, mapping: "Mapping[PackageRef, bytes]"):
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", DeprecationWarning)
            super().__init__()
        warnings.warn(
            "MemoryPackageProvider is deprecated; there is no replacement.",
            DeprecationWarning,
            stacklevel=2,
        )
        self.mapping = mapping

    def get_package_ids(self) -> "PackageIdSet":
        return frozenset(self.mapping.keys())

    def fetch_package(self, package_id: "PackageRef") -> bytes:
        return self.mapping.get(package_id)


@dataclass(frozen=True)
class PackageStoreCache:
    value_types: "TypeCache[ValName, Expr]"
    data_types: "TypeCache[TypeConName, Type]"
    templates: "TypeCache[TypeConName, Template]"


@dataclass(frozen=True)
class TypeCache(Generic[K, T]):
    """
    Immutable cache of type information.

    Instance attributes:

    .. attribute:: TypeCache.by_package_lookup

        A mapping from package ID strings to another mapping, where keys are fully-qualified type
        names for templates to a collection of matching Template instances. These collections are
        never non-empty; they will only ever contain multiple entries if there is a collision in
        dot-separated names.

    .. attribute:: TypeCache.by_name_lookup

        A mapping from fully-qualified type names (excluding package IDs) to ``Collection[T]``.
        These subcollections will only contain multiple values if names are not unique with respect
        to package IDs.
    """

    everything: "Collection[T]"
    by_package_lookup: "Mapping[str, Mapping[str, Collection[T]]]"
    by_name_lookup: "Mapping[str, Collection[T]]"

    @classmethod
    def build(cls, objects: "Mapping[K, T]") -> "TypeCache[K, T]":
        from ..damlast.util import package_local_name, package_ref

        everything = tuple(objects.values())
        by_package_lookup = defaultdict(lambda: defaultdict(list))
        by_name_lookup = defaultdict(list)

        for k, v in objects.items():
            package_id = package_ref(k)
            module_entity_name = package_local_name(k)
            # This is kept here for backwards compatibility, but its use should be discouraged
            module_entity_name_deprecated = module_entity_name.replace(":", ".")
            for valid_name in (module_entity_name, module_entity_name_deprecated):
                by_package_lookup[package_id][valid_name].append(v)
                by_name_lookup[valid_name].append(v)

        return TypeCache(
            everything,
            _immutable_mmc(by_package_lookup),
            MappingProxyType({k: tuple(v) for k, v in by_name_lookup.items()}),
        )

    def lookup(self, package_id: str, type_name: str) -> "Collection[T]":
        """
        Look up items based on package ID and type name. The values cannot be ``None``, but they
        can be the special-meaning ``'*'`` value.
        """
        safe_cast(str, package_id)
        safe_cast(str, type_name)

        if package_id == "*":
            if type_name == "*":
                # *:* means return everything
                return self.everything
            else:
                # *:Specific.Module:Entity means return all matches of the requested name, but
                # under any package
                return self.by_name_lookup.get(type_name, ())
        elif type_name == "*":
            # PKG_ID:* means return all matches in a specific package
            return reduce(add, self.by_package_lookup.get(package_id, {}).values(), ())
        else:
            # PKG_ID:Specific.Module:Entity; we are looking for a very specific type
            return self.by_package_lookup.get(package_id, {}).get(type_name, ())


def _immutable_mmc(
    mapping: "Mapping[str, Mapping[str, Collection[T]]]",
) -> "Mapping[str, Mapping[str, Collection[T]]]":
    """
    Create an immutable copy of :class:`TemplateStoreCache` data structures.
    """
    return MappingProxyType(
        {k1: MappingProxyType({k2: tuple(v) for k2, v in v1.items()}) for k1, v1 in mapping.items()}
    )


EMPTY_MAPPING = MappingProxyType({})
EMPTY_TYPE_CACHE = TypeCache((), EMPTY_MAPPING, EMPTY_MAPPING)
