from abc import ABC, abstractmethod
from collections import defaultdict
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Optional, Union

from bb_integrations_lib.gravitate.rita_api import GravitateRitaAPI
from bb_integrations_lib.models.rita.mapping import (
    Children,
    CompositeMapKey,
    Map,
    MappingType,
)
from bb_integrations_lib.secrets import RITACredential
from bb_integrations_lib.util.cache.custom_ttl_cache import CustomAsyncTTLCache


class MappingsNotLoadedException(Exception):
    pass


class RitaMapperKeyError(KeyError):
    def __init__(self, *args, from_key: str | CompositeMapKey):
        super().__init__(*args)

        self.from_key = from_key


class FastMapDirection(Enum):
    source_to_grav = 1
    grav_to_source = 2


# noinspection PyShadowingBuiltins
class FastMapKey:
    """
    A hashable key that can be used to identify mappings quickly if enough properties are known.
    Supports both simple string IDs and composite keys for entity identity (from_id).

    Note: parent_id is always a simple string. Composite keys are for entity identity resolution,
    not for hierarchical relationships. Children mappings only support string identifiers.
    """

    from_id: Union[str, CompositeMapKey]
    type: Optional[MappingType] = None
    parent_id: Optional[str] = None
    direction: FastMapDirection

    def __init__(
        self,
        from_id: Union[str, CompositeMapKey],
        type: Optional[MappingType],
        parent_id: Optional[str],
        direction: FastMapDirection,
    ):
        self.from_id = from_id
        self.type = type
        # parent_id is always a simple string (hierarchy uses string IDs)
        self.parent_id = parent_id
        self.direction = direction

        # Cache normalized values for hash/equality
        self._from_id_normalized = self._normalize_from_id(self.from_id)
        self._parent_id_normalized = self.parent_id or ""

    @staticmethod
    def _normalize_from_id(value: Union[str, CompositeMapKey]) -> str:
        """Convert from_id to hashable string representation."""
        if isinstance(value, CompositeMapKey):
            return value.to_cache_key()
        return value

    def __hash__(self):
        return hash(
            (
                self._from_id_normalized,
                self.type,
                self._parent_id_normalized,
                self.direction.value,
            )
        )

    def __eq__(self, other):
        if not isinstance(other, FastMapKey):
            return False
        return (
            self._from_id_normalized == other._from_id_normalized
            and self.type == other.type
            and self._parent_id_normalized == other._parent_id_normalized
            and self.direction == other.direction
        )

    def __repr__(self):
        type_name = self.type.name if self.type else None
        parent_id_str = self._parent_id_normalized or None
        return f"FastMap<from_id={self._from_id_normalized}, type_name={type_name}, parent_id={parent_id_str}, direction={self.direction.name}>"


class MappingProvider(ABC):
    """An abstract base class for implementing synchronous mapping providers for the Mapper class."""

    def __init__(self):
        pass

    @abstractmethod
    def get_mappings_by_source_system(self, source_system: str) -> list[Map]:
        pass


class AsyncMappingProvider(ABC):
    """An abstract base class for implementing asynchronous mapping providers for the Mapper class."""

    def __init__(self):
        pass

    @abstractmethod
    async def get_mappings_by_source_system(self, source_system: str) -> list[Map]:
        pass


class MockMappingProvider(MappingProvider):
    """A mapping provider that uses a constant list of mappings, intended for testing."""

    def __init__(self, mappings: list[Map]):
        super().__init__()
        self.mappings = mappings

    def get_mappings_by_source_system(self, source_system: str) -> list[Map]:
        """Returns mappings originally passed to init, filtered by source system."""
        return list(filter(lambda m: m.source_system == source_system, self.mappings))


class MockAsyncMappingProvider(AsyncMappingProvider):
    """A mapping provider that uses a constant list of mappings, intended for testing."""

    def __init__(self, mappings: list[Map]):
        super().__init__()
        self.mappings = mappings

    async def get_mappings_by_source_system(self, source_system: str) -> list[Map]:
        """Returns mappings originally passed to init, filtered by source system."""
        return list(filter(lambda m: m.source_system == source_system, self.mappings))


class RitaAPIMappingProvider(AsyncMappingProvider):
    """Accesses the RITA API to retrieve mappings, providing them to a Mapper."""

    def __init__(self, rita_client: GravitateRitaAPI):
        super().__init__()
        self.rita_client = rita_client

    async def from_credential(self, credential: RITACredential):
        self.rita_client = GravitateRitaAPI.from_credential(credential)

    async def get_mappings_by_source_system(self, source_system: str) -> list[Map]:
        return await self.rita_client.get_mappings_by_source_system(source_system)


class RitaAPICachedMappingProvider(RitaAPIMappingProvider):
    def __init__(self, rita_client: GravitateRitaAPI, ttl_seconds: int = 120):
        super().__init__(rita_client)
        self.ttl_seconds = ttl_seconds
        self.cache: CustomAsyncTTLCache = CustomAsyncTTLCache(verbose=False)

    async def get_mappings_by_source_system(self, source_system: str) -> list[Map]:
        return await self.cache.get_or_set(
            source_system,
            self.ttl_seconds,
            super().get_mappings_by_source_system,
            source_system,
        )


class RitaAPICrossrefProvider(RitaAPIMappingProvider):
    def __init__(self, rita_client: GravitateRitaAPI, data_type: str | None = None):
        """
        Downloads crossrefs async for a RITA tenant and converts them into Maps so they can be used with a RitaMapper.

        Specify data_type to match mappings of a specific crossref data type, or leave it None to match on any. This may
        create mapping conflicts that cannot be disambiguated, as the crossref data type cannot currently be tracked
        once it is converted to a Map.
        """
        super().__init__(rita_client)

        self.data_type = data_type

    async def get_mappings_by_source_system(self, source_system: str) -> list[Map]:

        cv2s = await self.rita_client.crossref_v2_get_many(
            source_system_id=source_system,
            # If data_type is specified, only mappings of that type.
            type_id=self.data_type
        )
        return [
            Map(
                gravitate_id=str(c.gravitate_id),
                source_id=str(c.source_id),
                source_system=source_system,
                type=MappingType.other,
            )
            for c in cv2s
        ]


@dataclass
class MappingLoadProblem:
    map: Map
    problem: str


@dataclass
class MappingLoadResult:
    problems: list[MappingLoadProblem] = field(default_factory=list)


# noinspection PyShadowingBuiltins
class RitaMapperCore(ABC):
    """
    Provides the core logic of mapping lookups. Entirely synchronous; the concrete implementations can use these
    methods directly, and should generally only implement loading mappings.
    """

    def __init__(self):
        self.loaded = False
        self.fast_maps = defaultdict(list)

    def _raise_if_not_loaded(self):
        if not self.loaded:
            raise MappingsNotLoadedException("Mappings used before getting loaded")

    def _load_mappings(self, mappings: list[Map]) -> MappingLoadResult:
        """
        Load the mappings from mappings into the mapper's internal data structures, making the get_ calls available
        for use. Supports both simple string IDs and composite keys for entity identity.

        Will raise a TypeError if you attempt to load any items where the source_id or gravitate_id is a CompositeMapKey
        but the mapping_type is not composite.

        Note: parent_id in FastMapKey is always a string. If a parent Map has a composite source_id,
        it's converted to string for hierarchical lookups.
        """

        self.fast_maps = defaultdict(list)

        # noinspection PyShadowingNames
        def _get_identifier(map_obj: Map | Children, direction: FastMapDirection):
            """Extract the appropriate identifier based on a map direction."""
            if direction == FastMapDirection.grav_to_source:
                return map_obj.gravitate_id
            return map_obj.source_id

        # noinspection PyShadowingNames
        def _get_parent_id_str(
            parent: Map | None, direction: FastMapDirection
        ) -> str | None:
            """Get parent identifier as string (hierarchy always uses string IDs)."""
            if parent is None:
                return None
            parent_id = _get_identifier(parent, direction)
            # Convert composite keys to string for parent_id (hierarchy uses strings)
            if isinstance(parent_id, CompositeMapKey):
                return parent_id.to_cache_key()
            return parent_id

        # noinspection PyShadowingNames
        def _reg_map(
            map: Map | Children, direction: FastMapDirection, parent: Map | None = None
        ):
            if map.type != MappingType.composite:
                if isinstance(map.source_id, CompositeMapKey):
                    raise TypeError(
                        f"CompositeMapKeys are only supported on mappings of type composite. "
                        f"(source_id was type {type(map.source_id)}, map.type={map.type}). Map: {map}"
                    )
                if isinstance(map.gravitate_id, CompositeMapKey):
                    raise TypeError(
                        f"CompositeMapKeys are only supported on mappings of type composite "
                        f"(gravitate_id of type {type(map.gravitate_id)}, map.type={map.type}). Map: {map}"
                    )
            from_id = _get_identifier(map, direction)
            parent_id = _get_parent_id_str(parent, direction)
            map_key = FastMapKey(from_id, map.type, parent_id, direction)
            self.fast_maps[map_key].append(map)

        problems: list[MappingLoadProblem] = []
        for direction in FastMapDirection:
            for m in mappings:
                try:
                    for child in m.children or []:
                        _reg_map(child, direction, m)
                    _reg_map(m, direction, None)
                except TypeError as e:
                    problems.append(MappingLoadProblem(m, str(e)))

        self.loaded = True

        return MappingLoadResult(problems)

    # noinspection PyUnreachableCode
    def get_mappings(
        self,
        from_id: str,
        direction: FastMapDirection,
        mapping_type: MappingType,
        parent_from_id: str | None = None,
    ) -> list[Map | Children]:
        """
        Retrieve mappings meeting the criteria. Base method for many of the get_gravitate/source* methods - prefer using
        those for a more expressive syntax.

        Does not support composite key lookups - use *_by_composite() methods for that.

        :arg str from_id: The original value we're looking to map. If direction is source_to_grav, this matches on the
          source_id field; if direction is grav_to_source, this matches on the gravitate_id field.
        :arg FastMapDirection direction: The direction of the mapping (gravitate->source or source->gravitate).
        :arg MappingType mapping_type: The type of the mapping.
        :arg str | None parent_from_id: The from_id of the child mapping's parent, if searching for a child. To get
          parent mappings instead, use None (the default).
        :raises MappingsNotLoadedException: If mappings have not been loaded yet.
        """
        self._raise_if_not_loaded()
        if mapping_type == MappingType.composite:
            raise TypeError("Use *_by_composite() methods for composite key lookups")
        if not isinstance(mapping_type, MappingType):
            raise TypeError(
                "mapping_type must be of type bb_integrations_lib.models.rita.mapping.MappingType "
                "(check whether you imported the right MappingType)"
            )

        return self.fast_maps.get(
            FastMapKey(from_id, mapping_type, parent_from_id, direction), []
        )

    def get_mappings_guaranteed(
        self,
        from_id: str,
        direction: FastMapDirection,
        mapping_type: MappingType,
        parent_from_id: str | None = None,
    ) -> list[Map | Children]:
        """Like get_mappings, but additionally raises a RitaMapperKeyError if no mappings are found."""
        self._raise_if_not_loaded()
        mappings = self.get_mappings(from_id, direction, mapping_type, parent_from_id)
        if not mappings:
            from_name = (
                "gravitate"
                if direction == FastMapDirection.grav_to_source
                else "source"
            )
            child_slug = (
                f", parent_{from_name}_id='{parent_from_id}'" if parent_from_id else ""
            )
            raise RitaMapperKeyError(
                f"No mappings found for {from_name}_id='{from_id}'{child_slug} with type='{mapping_type.name}'",
                from_key=from_id,
            )
        return mappings

    def get_single_mapping_guaranteed(
        self,
        from_id: str,
        direction: FastMapDirection,
        mapping_type: MappingType,
        parent_from_id: str | None = None,
    ) -> Map | Children:
        """
        Like get_mappings_guaranteed, but additionally raises a RitaMapperKeyError if more than 1 mapping is found.
        """
        mappings = self.get_mappings_guaranteed(
            from_id, direction, mapping_type, parent_from_id
        )
        if len(mappings) > 1:
            from_name = (
                "gravitate"
                if direction == FastMapDirection.grav_to_source
                else "source"
            )
            child_slug = (
                f", parent_{from_name}_id='{parent_from_id}'" if parent_from_id else ""
            )
            raise RitaMapperKeyError(
                f"Too many mappings ({len(mappings)}) for {from_name}_id='{from_id}'{child_slug} "
                f"with type='{mapping_type.name}', expected 1",
                from_key=from_id,
            )
        return mappings[0]

    @staticmethod
    def _ensure_str_or_none(val: Any) -> str | None:
        if val is not None and not isinstance(val, str):
            raise TypeError(f"Expected str, got {type(val)} instead")
        return val

    @staticmethod
    def _ensure_list_str_or_none(val: list[Any]) -> list[str | None]:
        for v in val:
            if v is not None and not isinstance(v, str):
                raise TypeError(f"Expected str, got {type(v)} instead")
        return val

    def get_gravitate_parent_id(
        self, source_id: str, mapping_type: MappingType
    ) -> str | None:
        return self._ensure_str_or_none(
            self.get_single_mapping_guaranteed(
                source_id, FastMapDirection.source_to_grav, mapping_type
            ).gravitate_id
        )

    def get_gravitate_child_id(
        self,
        source_parent_id: str,
        source_child_id: str,
        mapping_type: MappingType,
    ) -> str | None:
        """
        Get gravitate_id for a child mapping.

        Note: Child lookups use string IDs only. Children mappings don't support composite keys.
        For composite key lookups (without hierarchy), use get_gravitate_id_by_composite().
        """
        return self._ensure_str_or_none(
            self.get_single_mapping_guaranteed(
                source_child_id,
                FastMapDirection.source_to_grav,
                mapping_type,
                source_parent_id,
            ).gravitate_id
        )

    def get_source_parent_id(
        self, gravitate_id: str, mapping_type: MappingType
    ) -> str | None:
        return self._ensure_str_or_none(
            self.get_single_mapping_guaranteed(
                gravitate_id, FastMapDirection.grav_to_source, mapping_type
            ).source_id
        )

    def get_source_child_id(
        self,
        gravitate_parent_id: str,
        gravitate_child_id: str,
        mapping_type: MappingType,
    ) -> str | None:
        """
        Get source_id for a child mapping.

        Note: Child lookups use string IDs only. Children mappings don't support composite keys.
        For composite key lookups (without hierarchy), use get_source_id_by_composite().
        """
        return self._ensure_str_or_none(
            self.get_single_mapping_guaranteed(
                gravitate_child_id,
                FastMapDirection.grav_to_source,
                mapping_type,
                gravitate_parent_id,
            ).source_id
        )

    def get_gravitate_parent_ids(
        self, source_id: str, mapping_type: MappingType
    ) -> list[str | None]:
        return self._ensure_list_str_or_none(
            [
                x.gravitate_id
                for x in self.get_mappings_guaranteed(
                    source_id, FastMapDirection.source_to_grav, mapping_type
                )
            ]
        )

    def get_gravitate_child_ids(
        self,
        source_parent_id: str,
        source_child_id: str,
        mapping_type: MappingType,
    ) -> list[str | None]:
        """
        Get all gravitate_ids for child mappings matching the criteria.

        Note: Child lookups use string IDs only. Children mappings don't support composite keys.
        """
        return self._ensure_list_str_or_none(
            [
                x.gravitate_id
                for x in self.get_mappings_guaranteed(
                    source_child_id,
                    FastMapDirection.source_to_grav,
                    mapping_type,
                    source_parent_id,
                )
            ]
        )

    def get_source_parent_ids(
        self, gravitate_id: str, mapping_type: MappingType
    ) -> list[str | None]:
        return self._ensure_list_str_or_none(
            [
                x.source_id
                for x in self.get_mappings_guaranteed(
                    gravitate_id, FastMapDirection.grav_to_source, mapping_type
                )
            ]
        )

    def get_source_child_ids(
        self,
        gravitate_parent_id: str,
        gravitate_child_id: str,
        mapping_type: MappingType,
    ) -> list[str | None]:
        """
        Get all source_ids for child mappings matching the criteria.

        Note: Child lookups use string IDs only. Children mappings don't support composite keys.
        """
        return self._ensure_list_str_or_none(
            [
                x.source_id
                for x in self.get_mappings_guaranteed(
                    gravitate_child_id,
                    FastMapDirection.grav_to_source,
                    mapping_type,
                    gravitate_parent_id,
                )
            ]
        )

    # Composite key lookup methods
    def get_mappings_by_composite(
        self,
        from_key: Union[dict, CompositeMapKey],
        direction: FastMapDirection,
        mapping_type: MappingType | None,
        parent_from_id: Optional[str] = None,
    ) -> list[Map | Children]:
        """
        Retrieve mappings using a composite key for entity identity.

        :param from_key: Dict or CompositeMapKey for lookup (entity identity)
        :param direction: Mapping direction
        :param mapping_type: Type of mapping
        :param parent_from_id: Parent identifier (string only - hierarchy uses simple IDs)

        Note: Composite keys are for entity identity, not hierarchy. If looking up children,
        use string parent_from_id. Children mappings don't support composite keys.
        """
        self._raise_if_not_loaded()
        if isinstance(from_key, dict):
            from_key = CompositeMapKey(key=from_key)
        return self.fast_maps.get(
            FastMapKey(from_key, mapping_type, parent_from_id, direction), []
        )

    def get_mapping_by_composite(
        self,
        from_key: Union[dict, CompositeMapKey],
        direction: FastMapDirection,
        mapping_type: MappingType | None,
    ) -> Map | Children:
        """Get single mapping by composite key, raises if not found or multiple."""
        mappings = self.get_mappings_by_composite(from_key, direction, mapping_type)
        if isinstance(from_key, dict):
            from_key = CompositeMapKey(key=from_key)
        if not mappings:
            raise RitaMapperKeyError(
                f"No mapping found for composite key: {from_key}", from_key=from_key
            )
        if len(mappings) > 1:
            raise RitaMapperKeyError(
                f"Multiple mappings ({len(mappings)}) found for: {from_key}",
                from_key=from_key,
            )
        return mappings[0]

    def get_gravitate_id_by_composite(
        self, source_key: Union[dict, CompositeMapKey], mapping_type: MappingType | None
    ) -> str | CompositeMapKey | None:
        """Get gravitate_id using a composite source key."""
        return self.get_mapping_by_composite(
            source_key, FastMapDirection.source_to_grav, mapping_type
        ).gravitate_id

    def get_source_id_by_composite(
        self,
        gravitate_key: Union[dict, CompositeMapKey],
        mapping_type: MappingType | None,
    ) -> str | CompositeMapKey | None:
        """Get source_id using a composite gravitate key."""
        return self.get_mapping_by_composite(
            gravitate_key, FastMapDirection.grav_to_source, mapping_type
        ).source_id

    def find_mappings_by_partial_key(
        self,
        partial_key: dict[str, str],
        direction: FastMapDirection,
        mapping_type: MappingType,
    ) -> list[Map | Children]:
        """
        Find all mappings where the composite key contains the partial key fields.

        Example: partial_key={"product": "X"} matches {"product": "X", "terminal": "Y"}
        """
        self._raise_if_not_loaded()
        results = []
        for fast_key, mappings in self.fast_maps.items():
            if fast_key.direction != direction:
                continue
            if fast_key.type != mapping_type:
                continue
            # Check if from_id is a CompositeMapKey and matches partial
            if isinstance(fast_key.from_id, CompositeMapKey):
                if fast_key.from_id.matches(partial_key):
                    results.extend(mappings)
        return results


class RitaMapper(RitaMapperCore):
    def __init__(
        self, provider: MappingProvider | AsyncMappingProvider, source_system: str
    ):
        super().__init__()
        self.provider = provider
        self.source_system = source_system

    async def load_mappings_async(self) -> MappingLoadResult:
        """
        Load mappings from the provider asynchronously.

        :raises TypeError: If the provider is not async.

        Note that this method does not check whether mappings are already loaded; calling it will overwrite any stored
        mappings with a fresh set from the provider.
        """
        if isinstance(self.provider, AsyncMappingProvider):
            mappings = await self.provider.get_mappings_by_source_system(
                self.source_system
            )
        else:
            raise TypeError("Provider is not async")
        return super()._load_mappings(mappings)

    def load_mappings(self) -> MappingLoadResult:
        """
        Load mappings from the provider synchronously.

        :raises TypeError: If the provider is async.

        Note that this method does not check whether mappings are already loaded; calling it will overwrite any stored
        mappings with a fresh set from the provider.
        """
        if isinstance(self.provider, MappingProvider):
            mappings = self.provider.get_mappings_by_source_system(self.source_system)
        else:
            raise TypeError("Provider is not synchronous")
        return super()._load_mappings(mappings)
