from collections.abc import Sequence
from itertools import groupby
from typing import Any, Literal, TypeVar, cast, overload

from cognite.client._api.data_modeling import InstancesAPI
from cognite.client._constants import DEFAULT_LIMIT_READ
from cognite.client.data_classes.data_modeling import (
    NodeList,
    filters,
    query,
)
from cognite.client.utils.useful_types import SequenceNotStr

from cognite_toolkit._cdf_tk.client.api.instances import WrappedInstancesAPI
from cognite_toolkit._cdf_tk.client.cdf_client import PagedResponse, ResponseItems
from cognite_toolkit._cdf_tk.client.http_client import HTTPClient, ItemsSuccessResponse, SuccessResponse
from cognite_toolkit._cdf_tk.client.resource_classes.data_modeling import NodeReference, ViewReference
from cognite_toolkit._cdf_tk.client.resource_classes.instance_api import TypedNodeIdentifier
from cognite_toolkit._cdf_tk.client.resource_classes.legacy.migration import (
    AssetCentricId,
    CreatedSourceSystem,
    InstanceSource,
    SpaceSource,
)
from cognite_toolkit._cdf_tk.client.resource_classes.resource_view_mapping import (
    RESOURCE_VIEW_MAPPING_SPACE,
    ResourceViewMappingRequest,
    ResourceViewMappingResponse,
)
from cognite_toolkit._cdf_tk.utils.collection import chunker_sequence
from cognite_toolkit._cdf_tk.utils.useful_types import AssetCentricType


class InstanceSourceAPI:
    def __init__(self, instance_api: InstancesAPI) -> None:
        self._instance_api = instance_api
        self._RETRIEVE_LIMIT = 1000
        self._view_id = InstanceSource.get_source()

    def retrieve(self, ids: Sequence[AssetCentricId]) -> NodeList[InstanceSource]:
        """Retrieve a list of instance sources by their IDs.

        Args:
            ids (Sequence[AssetCentricId]): A sequence of AssetCentricId objects representing the IDs of the instance sources to retrieve.
        """
        results: NodeList[InstanceSource] = NodeList[InstanceSource]([])
        for chunk in chunker_sequence(ids, self._RETRIEVE_LIMIT):
            retrieve_query = query.Query(
                with_={
                    "instanceSource": query.NodeResultSetExpression(
                        filter=filters.And(filters.HasData(views=[self._view_id]), self._create_dms_filter(ids)),
                        limit=len(chunk),
                    ),
                },
                select={"instanceSource": query.Select([query.SourceSelector(self._view_id, ["*"])])},
            )
            chunk_response = self._instance_api.query(retrieve_query)
            results.extend([InstanceSource._load(item.dump()) for item in chunk_response.get("instanceSource", [])])
        return results

    @staticmethod
    def _create_dms_filter(ids: Sequence[AssetCentricId]) -> filters.Filter:
        """Create a filter that matches all the AssetCentricIds in the list."""
        if not ids:
            raise ValueError("Cannot create a filter from an empty AssetCentricIdList.")
        to_or_filters: list[filters.Filter] = []
        instance_source_view = InstanceSource.get_source()
        for resource_type, resource_ids in groupby(
            sorted(ids, key=lambda x: x.resource_type), key=lambda x: x.resource_type
        ):
            is_resource = filters.Equals(instance_source_view.as_property_ref("resourceType"), resource_type)
            is_id = filters.In(
                instance_source_view.as_property_ref("id"), [resource_id.id_ for resource_id in resource_ids]
            )
            to_or_filters.append(filters.And(is_resource, is_id))
        return filters.Or(*to_or_filters)

    def list(self, resource_type: AssetCentricType, limit: int | None = DEFAULT_LIMIT_READ) -> NodeList[InstanceSource]:
        """List instance sources filtered by resource type"""
        is_selected = filters.Equals(self._view_id.as_property_ref("resourceType"), resource_type)

        nodes = self._instance_api.list(instance_type="node", filter=is_selected, limit=limit, sources=self._view_id)
        return NodeList[InstanceSource]([InstanceSource._load(node.dump()) for node in nodes])


class ResourceViewMappingsAPI(WrappedInstancesAPI[TypedNodeIdentifier, ResourceViewMappingResponse]):
    def __init__(self, http_client: HTTPClient) -> None:
        super().__init__(http_client, ResourceViewMappingRequest.VIEW_ID)

    def _validate_response(self, response: SuccessResponse) -> ResponseItems[TypedNodeIdentifier]:
        return ResponseItems[TypedNodeIdentifier].model_validate_json(response.body)

    def _validate_page_response(
        self, response: SuccessResponse | ItemsSuccessResponse
    ) -> PagedResponse[ResourceViewMappingResponse]:
        return PagedResponse[ResourceViewMappingResponse].model_validate_json(response.body)

    def list(self, resource_type: str | None = None, limit: int | None = 100) -> list[ResourceViewMappingResponse]:
        filter_: dict[str, Any] = {
            "equals": {
                "property": ["node", "space"],
                "value": RESOURCE_VIEW_MAPPING_SPACE,
            }
        }
        if resource_type:
            filter_ = {
                "and": [
                    filter_,
                    {
                        "equals": {
                            "property": ResourceViewMappingRequest.VIEW_ID.as_property_reference("resourceType"),
                            "value": resource_type,
                        }
                    },
                ]
            }
        return super()._list_instances(filter=filter_, instance_type="node", limit=limit)


class CreatedSourceSystemAPI:
    def __init__(self, instance_api: InstancesAPI) -> None:
        self._instance_api = instance_api
        self._RETRIEVE_LIMIT = 1000
        self._view_id = CreatedSourceSystem.get_source()

    def retrieve(self, source: SequenceNotStr[str]) -> NodeList[CreatedSourceSystem]:
        """Retrieve one or more view sources by their external IDs."""
        results: NodeList[CreatedSourceSystem] = NodeList[CreatedSourceSystem]([])
        # MyPy does not understand that SequenceNotStr is a sequence.
        for chunk in chunker_sequence(source, self._RETRIEVE_LIMIT):  # type: ignore[type-var]
            retrieve_query = query.Query(
                with_={
                    "sourceSystem": query.NodeResultSetExpression(
                        filter=filters.And(filters.HasData(views=[self._view_id]), self._create_dms_filter(chunk)),
                        limit=len(chunk),
                    ),
                },
                select={"sourceSystem": query.Select([query.SourceSelector(self._view_id, ["*"])])},
            )
            chunk_response = self._instance_api.query(retrieve_query)
            results.extend([CreatedSourceSystem._load(item.dump()) for item in chunk_response.get("sourceSystem", [])])
        return results

    def _create_dms_filter(self, source: SequenceNotStr[str]) -> filters.Filter:
        """Create a filter that matches all CreatedSourceSystem with given source in the list."""
        if not source:
            raise ValueError("Cannot create a filter from an empty source list.")
        return filters.In(self._view_id.as_property_ref("source"), list(source))

    def list(self, limit: int = -1) -> NodeList[CreatedSourceSystem]:
        """Lists all created source systems."""
        return self._instance_api.list(instance_type=CreatedSourceSystem, limit=limit)


_T = TypeVar("_T", bound=int | str)


class SpaceSourceAPI:
    def __init__(self, instance_api: InstancesAPI) -> None:
        self._instance_api = instance_api
        self._RETRIEVE_LIMIT = 1000
        self._view_id = SpaceSource.get_source()
        # Cache for space sources by data set ID and external ID
        self._cache_by_id: dict[int, SpaceSource] = {}
        self._cache_by_external_id: dict[str, SpaceSource] = {}

    @overload
    def retrieve(self, data_set_id: int) -> SpaceSource | None: ...

    @overload
    def retrieve(self, data_set_id: Sequence[int]) -> NodeList[SpaceSource]: ...

    @overload
    def retrieve(self, *, data_set_external_id: str) -> SpaceSource | None: ...

    @overload
    def retrieve(self, *, data_set_external_id: SequenceNotStr[str]) -> NodeList[SpaceSource]: ...

    def retrieve(
        self,
        data_set_id: int | Sequence[int] | None = None,
        data_set_external_id: str | SequenceNotStr[str] | None = None,
    ) -> SpaceSource | NodeList[SpaceSource] | None:
        """Retrieve a space source by data set ID or external ID.

        This method uses caching to avoid redundant API calls. If a space source
        has been retrieved before, it will be returned from cache.
        """
        if data_set_id is not None and data_set_external_id is None:
            return self._retrieve_with_cache(
                value=data_set_id,
                property_name="dataSetId",
                cache=self._cache_by_id,
                is_single=isinstance(data_set_id, int),
            )
        elif data_set_external_id is not None:
            return self._retrieve_with_cache(
                # MyPy is confused by SequenceNotStr
                value=data_set_external_id,  # type: ignore[arg-type]
                property_name="dataSetExternalId",
                cache=self._cache_by_external_id,
                is_single=isinstance(data_set_external_id, str),
            )
        else:
            raise ValueError("Either data_set_id or data_set_external_id must be provided, but not both.")

    def _retrieve_with_cache(
        self,
        value: _T | Sequence[_T],
        property_name: str,
        cache: dict[_T, SpaceSource],
        is_single: bool,
    ) -> SpaceSource | NodeList[SpaceSource] | None:
        """Retrieve space sources with caching support."""
        # MyPy does not understand that if is_single is True, value is _T, else Sequence[_T].
        values: list[_T] = [value] if is_single else list(value)  # type: ignore[arg-type, list-item]

        # Check cache for all requested values
        cached_results: list[SpaceSource] = []
        missing_values: list[_T] = []
        for val in values:
            if val in cache:
                cached_results.append(cache[val])
            else:
                missing_values.append(val)

        # Fetch missing values from API
        if missing_values:
            fetched = self._retrieve_by_property(property_name=property_name, values=missing_values)
            for item in fetched:
                self._add_to_cache(item)
                cached_results.append(item)

        if is_single:
            return cached_results[0] if cached_results else None
        return NodeList[SpaceSource](cached_results)

    def _add_to_cache(self, space_source: SpaceSource) -> None:
        """Add a space source to both caches."""
        if space_source.data_set_id:
            self._cache_by_id[space_source.data_set_id] = space_source
        if space_source.data_set_external_id is not None:
            self._cache_by_external_id[space_source.data_set_external_id] = space_source

    def _retrieve_by_property(
        self,
        property_name: str,
        values: Sequence[_T],
    ) -> NodeList[SpaceSource]:
        """Retrieve space sources by filtering on a specific property."""
        results: NodeList[SpaceSource] = NodeList[SpaceSource]([])
        for chunk in chunker_sequence(values, self._RETRIEVE_LIMIT):
            retrieve_query = query.Query(
                with_={
                    "spaceSource": query.NodeResultSetExpression(
                        filter=filters.And(
                            filters.HasData(views=[self._view_id]),
                            filters.In(self._view_id.as_property_ref(property_name), chunk),
                        ),
                        limit=len(chunk),
                    ),
                },
                select={"spaceSource": query.Select([query.SourceSelector(self._view_id, ["*"])])},
            )
            chunk_response = self._instance_api.query(retrieve_query)
            items = chunk_response.get("spaceSource", [])
            results.extend([SpaceSource._load(item.dump()) for item in items])

        return results

    def list(self, limit: int = -1) -> NodeList[SpaceSource]:
        """Lists all space sources and populates the cache.

        Note: This method always fetches from the API and does not use the cache,
        but it will populate the cache with all retrieved space sources.
        """
        results = self._instance_api.list(instance_type=SpaceSource, limit=limit)
        # Populate cache with all retrieved space sources
        for space_source in results:
            self._add_to_cache(space_source)
        return results


_T_Cached = TypeVar("_T_Cached", bound=NodeReference | ViewReference)


class LookupAPI:
    def __init__(self, instance_api: InstancesAPI, resource_type: AssetCentricType) -> None:
        self._instance_api = instance_api
        self._resource_type = resource_type
        self._view_id = InstanceSource.get_source()
        self._node_id_by_id: dict[int, NodeReference | None] = {}
        self._node_id_by_external_id: dict[str, NodeReference | None] = {}
        self._consumer_view_id_by_id: dict[int, ViewReference | None] = {}
        self._consumer_view_id_by_external_id: dict[str, ViewReference | None] = {}
        self._RETRIEVE_LIMIT = 1000

    @overload
    def __call__(self, id: int, external_id: None = None) -> NodeReference | None: ...

    @overload
    def __call__(self, id: SequenceNotStr[int], external_id: None = None) -> dict[int, NodeReference]: ...

    @overload
    def __call__(self, *, external_id: str) -> NodeReference | None: ...

    @overload
    def __call__(self, *, external_id: SequenceNotStr[str]) -> dict[str, NodeReference]: ...

    def __call__(
        self, id: int | SequenceNotStr[int] | None = None, external_id: str | SequenceNotStr[str] | None = None
    ) -> dict[int, NodeReference] | dict[str, NodeReference] | NodeReference | None:
        """Lookup NodeId by either internal ID or external ID.

        Args:
            id (int | Sequence[int] | None): The internal ID(s) to lookup.
            external_id (str | SequenceNotStr[str] | None): The external ID(s) to lookup.

        Returns:
            NodeReference | dict[int, NodeReference] | dict[str, NodeReference] | None: The corresponding NodeReference(s) if found, otherwise None.

        """
        if id is not None and external_id is None:
            return self._lookup(
                identifier=id,
                cache=self._node_id_by_id,
                property_name="id",
                return_type=NodeReference,
                input_type=int,
            )
        elif external_id is not None and id is None:
            return self._lookup(
                identifier=external_id,
                cache=self._node_id_by_external_id,
                property_name="classicExternalId",
                return_type=NodeReference,
                input_type=str,
            )
        else:
            raise TypeError("Either id or external_id must be provided, but not both.")

    @overload
    def consumer_view(self, id: int, external_id: None = None) -> ViewReference | None: ...

    @overload
    def consumer_view(self, id: SequenceNotStr[int], external_id: None = None) -> dict[int, ViewReference]: ...

    @overload
    def consumer_view(self, *, external_id: str) -> ViewReference | None: ...

    @overload
    def consumer_view(self, *, external_id: SequenceNotStr[str]) -> dict[str, ViewReference]: ...

    def consumer_view(
        self, id: int | SequenceNotStr[int] | None = None, external_id: str | SequenceNotStr[str] | None = None
    ) -> dict[int, ViewReference] | dict[str, ViewReference] | ViewReference | None:
        """Lookup Consumer ViewReference by either internal ID or external ID.

        Args:
            id (int | Sequence[int] | None): The internal ID(s) to lookup.
            external_id (str | SequenceNotStr[str] | None): The external ID(s) to lookup.
        Returns:
            ViewReference | dict[int, ViewReference] | dict[str, ViewReference] | None
        """
        if id is not None and external_id is None:
            return self._lookup(
                identifier=id,
                cache=self._consumer_view_id_by_id,
                property_name="id",
                return_type=ViewReference,
                input_type=int,
            )
        elif external_id is not None and id is None:
            return self._lookup(
                identifier=external_id,
                cache=self._consumer_view_id_by_external_id,
                property_name="classicExternalId",
                return_type=ViewReference,
                input_type=str,
            )
        else:
            raise TypeError("Either id or external_id must be provided, but not both.")

    def _lookup(
        self,
        identifier: _T | SequenceNotStr[_T],
        cache: dict[_T, _T_Cached | None],
        property_name: Literal["id", "classicExternalId"],
        return_type: type[_T_Cached],
        input_type: type[_T],
    ) -> dict[_T, _T_Cached] | _T_Cached | None:
        """Generic lookup method for both NodeId and ViewReference by id or external_id."""
        is_single = isinstance(identifier, input_type)
        # MyPy does not understand that if is_single is True, identifier is _T, else SequenceNotStr[_T].
        identifiers: list[_T] = [identifier] if is_single else list(identifier)  # type: ignore[arg-type, list-item]

        missing = [id_ for id_ in identifiers if id_ not in cache]
        if missing:
            self._fetch_and_cache(missing, by=property_name)

        if is_single:
            return cache.get(identifier)  # type: ignore[arg-type]

        return {id_: value for id_ in identifiers if isinstance(value := cache.get(id_), return_type)}

    def _fetch_and_cache(self, identifiers: Sequence[int | str], by: Literal["id", "classicExternalId"]) -> None:
        for chunk in chunker_sequence(identifiers, self._RETRIEVE_LIMIT):
            retrieve_query = query.Query(
                with_={
                    "instanceSource": query.NodeResultSetExpression(
                        filter=filters.And(
                            filters.HasData(views=[self._view_id]),
                            filters.Equals(self._view_id.as_property_ref("resourceType"), self._resource_type),
                            filters.In(self._view_id.as_property_ref(by), list(chunk)),
                        ),
                        limit=len(chunk),
                    ),
                },
                select={"instanceSource": query.Select([query.SourceSelector(self._view_id, ["*"])])},
            )
            chunk_response = self._instance_api.query(retrieve_query)
            items = chunk_response.get("instanceSource", [])
            for item in items:
                instance_source = InstanceSource._load(item.dump())
                node_id = NodeReference(space=instance_source.space, external_id=instance_source.external_id)
                self._node_id_by_id[instance_source.id_] = node_id
                self._consumer_view_id_by_id[instance_source.id_] = instance_source.consumer_view()
                if instance_source.classic_external_id:
                    self._node_id_by_external_id[instance_source.classic_external_id] = node_id
                    self._consumer_view_id_by_external_id[instance_source.classic_external_id] = (
                        instance_source.consumer_view()
                    )
            missing = set(chunk) - set(self._node_id_by_id.keys()) - set(self._node_id_by_external_id.keys())
            if by == "id":
                for missing_id in cast(set[int], missing):
                    if missing_id not in self._node_id_by_id:
                        self._node_id_by_id[missing_id] = None
                    if missing_id not in self._consumer_view_id_by_id:
                        self._consumer_view_id_by_id[missing_id] = None
            elif by == "classicExternalId":
                for missing_ext_id in cast(set[str], missing):
                    if missing_ext_id not in self._node_id_by_external_id:
                        self._node_id_by_external_id[missing_ext_id] = None
                    if missing_ext_id not in self._consumer_view_id_by_external_id:
                        self._consumer_view_id_by_external_id[missing_ext_id] = None


class MigrationLookupAPI:
    def __init__(self, instance_api: InstancesAPI) -> None:
        self.assets = LookupAPI(instance_api, "asset")
        self.events = LookupAPI(instance_api, "event")
        self.files = LookupAPI(instance_api, "file")
        self.time_series = LookupAPI(instance_api, "timeseries")


class MigrationAPI:
    def __init__(self, instance_api: InstancesAPI, http_client: HTTPClient) -> None:
        self.instance_source = InstanceSourceAPI(instance_api)
        self.resource_view_mapping = ResourceViewMappingsAPI(http_client)
        self.created_source_system = CreatedSourceSystemAPI(instance_api)
        self.space_source = SpaceSourceAPI(instance_api)
        self.lookup = MigrationLookupAPI(instance_api)
