from abc import ABC, abstractmethod
from collections import defaultdict
from collections.abc import Sequence
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from functools import cached_property, lru_cache, partial
from typing import ClassVar, Literal, TypeVar, get_args, overload

import questionary
from cognite.client.data_classes import (
    Asset,
    UserProfileList,
    filters,
)
from cognite.client.data_classes.aggregations import Count
from cognite.client.data_classes.data_modeling import ViewId
from cognite.client.data_classes.data_modeling.statistics import SpaceStatistics
from cognite.client.utils import ms_to_datetime
from questionary import Choice
from rich.console import Console

from cognite_toolkit._cdf_tk.client import ToolkitClient
from cognite_toolkit._cdf_tk.client.request_classes.filters import DataModelFilter, ViewFilter
from cognite_toolkit._cdf_tk.client.resource_classes.apm_config_v1 import APMConfigResponse
from cognite_toolkit._cdf_tk.client.resource_classes.chart import ChartResponse, Visibility
from cognite_toolkit._cdf_tk.client.resource_classes.data_modeling import (
    ContainerReference,
    DataModelReference,
    DataModelResponse,
    DataModelResponseWithViews,
    SpaceResponse,
    ViewReference,
    ViewResponse,
)
from cognite_toolkit._cdf_tk.client.resource_classes.dataset import DataSetResponse
from cognite_toolkit._cdf_tk.client.resource_classes.identifiers import ExternalId, RawTableId
from cognite_toolkit._cdf_tk.client.resource_classes.legacy.canvas import Canvas
from cognite_toolkit._cdf_tk.client.resource_classes.resource_view_mapping import ResourceViewMappingResponse
from cognite_toolkit._cdf_tk.client.resource_classes.three_d import ThreeDModelClassicResponse
from cognite_toolkit._cdf_tk.exceptions import ToolkitMissingResourceError, ToolkitValueError

from . import humanize_collection
from .aggregators import (
    AssetAggregator,
    AssetCentricAggregator,
    EventAggregator,
    FileAggregator,
    TimeSeriesAggregator,
)
from .useful_types import AssetCentricDestinationType

T_Type = TypeVar("T_Type", bound=Asset | DataSetResponse)


class AssetCentricInteractiveSelect(ABC):
    def __init__(self, client: ToolkitClient, operation: str) -> None:
        self.client = client
        self.operation = operation
        self._aggregator = self._get_aggregator(client)

    @abstractmethod
    def _get_aggregator(self, client: ToolkitClient) -> AssetCentricAggregator:
        raise NotImplementedError()

    @lru_cache
    def aggregate_count(self, hierarchies: tuple[str, ...], data_sets: tuple[str, ...]) -> int:
        return self._aggregate_count(list(hierarchies), list(data_sets))

    def _aggregate_count(self, hierarchies: list[str], data_sets: list[str]) -> int:
        return self._aggregator.count(hierarchies, data_sets)

    @lru_cache
    def _get_available_data_sets(self, hierarchy: str | None = None) -> list[DataSetResponse]:
        if hierarchy is not None:
            datasets = self._aggregator.used_data_sets(hierarchy)
            return self.client.tool.datasets.retrieve(ExternalId.from_external_ids(datasets), ignore_unknown_ids=True)
        else:
            return self.client.tool.datasets.list(limit=None)

    @lru_cache
    def _get_available_hierarchies(self, data_set: str | None = None) -> list[Asset]:
        data_set_external_ids = [data_set] if data_set else None
        return list(self.client.assets.list(root=True, limit=-1, data_set_external_ids=data_set_external_ids))

    def _create_choice(self, item: Asset | DataSetResponse) -> tuple[questionary.Choice, int]:
        """Create a questionary choice for the given item."""
        if item.external_id is None:
            item_count = -1  # No count available for DataSet/Assets without external_id
        elif isinstance(item, DataSetResponse):
            item_count = self.aggregate_count(tuple(), (item.external_id,))
        elif isinstance(item, Asset):
            item_count = self.aggregate_count((item.external_id,), tuple())
        else:
            raise ToolkitValueError(f"Unexpected item type: {type(item)}. Expected Asset or DataSet.")

        return questionary.Choice(
            title=f"{item.name} ({item.external_id}) [{item_count:,}]"
            if item.name != item.external_id
            else f"({item.external_id}) [{item_count:,}]",
            value=item.external_id,
        ), item_count

    @overload
    def select_hierarchy(self, allow_empty: Literal[False] = False) -> str: ...

    @overload
    def select_hierarchy(self, allow_empty: Literal[True]) -> str | None: ...

    def select_hierarchy(self, allow_empty: Literal[False, True] = False) -> str | None:
        """Select a hierarchy interactively."""
        options = self._get_available_hierarchies()
        return self._select("hierarchy", options, allow_empty, "single")

    def select_hierarchies(self) -> list[str]:
        options = self._get_available_hierarchies()
        return self._select("hierarchies", options, False, "multiple")

    @overload
    def select_data_set(self, allow_empty: Literal[False] = False) -> str: ...

    @overload
    def select_data_set(self, allow_empty: Literal[True]) -> str | None: ...

    def select_data_set(self, allow_empty: Literal[False, True] = False) -> str | None:
        """Select a data set interactively."""
        options = self._get_available_data_sets()
        return self._select("data set", options, allow_empty, "single")

    def select_data_sets(self) -> list[str]:
        """Select multiple data sets interactively."""
        options = self._get_available_data_sets()
        return self._select("data sets", options, False, "multiple")

    def select_hierarchies_and_data_sets(self) -> tuple[list[str], list[str]]:
        what = self.select_hierarchies_or_data_sets()
        if what == "Hierarchy":
            hierarchy = self._select("hierarchy", self._get_available_hierarchies(), False, "single")
            data_sets = self._get_available_data_sets(hierarchy=hierarchy)
            if not data_sets:
                return [hierarchy], []
            selected_data_sets = self._select(f"data sets in hierarchy {hierarchy!r}", data_sets, True, "multiple")
            return [hierarchy], selected_data_sets or []
        elif what == "Data Set":
            data_set = self._select("data Set", self._get_available_data_sets(), False, "single")
            hierarchies = self._get_available_hierarchies(data_set=data_set)
            if not hierarchies:
                return [], [data_set]
            selected_hierarchies = self._select(f"hierarchies in data set {data_set!r} ", hierarchies, True, "multiple")
            return selected_hierarchies or [], [data_set]

    def select_hierarchies_or_data_sets(self) -> Literal["Hierarchy", "Data Set"]:
        what = questionary.select(
            f"Do you want to {self.operation} a hierarchy or a data set?", choices=["Hierarchy", "Data Set"]
        ).unsafe_ask()
        if what not in ["Hierarchy", "Data Set"]:
            raise ToolkitValueError(f"Unexpected selection: {what}. Aborting.")
        return what

    @overload
    def _select(
        self,
        what: str,
        options: Sequence[T_Type],
        allow_empty: Literal[False],
        plurality: Literal["single"],
    ) -> str: ...

    @overload
    def _select(
        self,
        what: str,
        options: Sequence[T_Type],
        allow_empty: Literal[True],
        plurality: Literal["single"],
    ) -> str | None: ...

    @overload
    def _select(
        self,
        what: str,
        options: Sequence[T_Type],
        allow_empty: Literal[False],
        plurality: Literal["multiple"],
    ) -> list[str]: ...

    @overload
    def _select(
        self,
        what: str,
        options: Sequence[T_Type],
        allow_empty: Literal[True],
        plurality: Literal["multiple"],
    ) -> list[str] | None: ...

    def _select(
        self,
        what: str,
        options: Sequence[Asset | DataSetResponse],
        allow_empty: Literal[True, False],
        plurality: Literal["single", "multiple"],
    ) -> str | list[str] | None:
        """Select a single item interactively."""
        if not options and not allow_empty:
            raise ToolkitValueError(f"No {what} available to select.")

        choices: list[questionary.Choice] = []
        none_sentinel = object()
        if allow_empty:
            choices.append(questionary.Choice(title=f"All {what}", value=none_sentinel))

        for choice, count in sorted(
            (self._create_choice(item) for item in options), key=lambda x: (-x[1], x[0].title)
        ):  # cont, choice.title
            if count > 0:
                choices.append(choice)
        if not choices and not allow_empty:
            raise ToolkitValueError(f"No {what} available with data to select.")

        message = f"Select a {what} to {self.operation} listed as 'name (external_id) [count]'"
        if plurality == "multiple":
            selected = questionary.checkbox(
                message,
                choices=choices,
                validate=lambda choices: "You must select at least one." if not choices else True,
            ).unsafe_ask()
        else:
            selected = questionary.select(message, choices=choices).unsafe_ask()
        if selected is None:
            raise ToolkitValueError(f"No {what} selected. Aborting.")
        elif selected is none_sentinel or (isinstance(selected, list) and none_sentinel in selected):
            return None
        return selected


class AssetInteractiveSelect(AssetCentricInteractiveSelect):
    def _get_aggregator(self, client: ToolkitClient) -> AssetCentricAggregator:
        return AssetAggregator(self.client)


class FileMetadataInteractiveSelect(AssetCentricInteractiveSelect):
    def _get_aggregator(self, client: ToolkitClient) -> AssetCentricAggregator:
        return FileAggregator(self.client)


class TimeSeriesInteractiveSelect(AssetCentricInteractiveSelect):
    def _get_aggregator(self, client: ToolkitClient) -> AssetCentricAggregator:
        return TimeSeriesAggregator(self.client)


class EventInteractiveSelect(AssetCentricInteractiveSelect):
    def _get_aggregator(self, client: ToolkitClient) -> AssetCentricAggregator:
        return EventAggregator(self.client)


class RawTableInteractiveSelect:
    def __init__(self, client: ToolkitClient, operation: str) -> None:
        self.client = client
        self.operation = operation

    def _available_databases(self) -> list[str]:
        databases = self.client.tool.raw.databases.list(limit=None)
        return [db.name for db in databases]

    def _available_tables(self, database: str) -> list[RawTableId]:
        tables = self.client.tool.raw.tables.list(db_name=database, limit=None)
        return [table.as_id() for table in tables]

    def select_tables(self, database: str | None = None) -> list[RawTableId]:
        """Interactively select raw tables."""
        databases = self._available_databases()
        if not databases:
            raise ToolkitValueError("No raw databases available. Aborting.")
        if database and database not in databases:
            raise ToolkitValueError(
                f"Database '{database}' not found in available raw databases: {databases}. Aborting."
            )
        elif database:
            selected_database = database
        else:
            selected_database = questionary.select(
                f"Select a Raw Database to {self.operation}",
                choices=[questionary.Choice(title=db, value=db) for db in databases],
            ).unsafe_ask()
        available_tables = self._available_tables(selected_database)
        if not available_tables:
            raise ToolkitValueError(f"No raw tables available in database '{selected_database}'. Aborting.")

        selected_tables = questionary.checkbox(
            f"Select Raw Tables in {selected_database} to {self.operation}",
            choices=[questionary.Choice(title=f"{table.name}", value=table) for table in available_tables],
            validate=lambda choices: True if choices else "You must select at least one table.",
        ).unsafe_ask()

        return selected_tables


@dataclass
class CanvasFilter:
    visibility: Literal["public", "private"] | None = None
    created_by: Literal["user"] | None = None
    select_all: bool = False

    def as_dms_filter(self) -> filters.Filter:
        canvas_id = Canvas.get_source()
        leaf_filters: list[filters.Filter] = [
            filters.Not(filters.Equals(canvas_id.as_property_ref("isArchived"), True)),
            # When sourceCanvasId is not set, we get the newest version of the canvas
            filters.Not(filters.Exists(canvas_id.as_property_ref("sourceCanvasId"))),
        ]
        if self.visibility is not None:
            leaf_filters.append(filters.Equals(canvas_id.as_property_ref("visibility"), self.visibility))

        return filters.And(*leaf_filters)


class InteractiveCanvasSelect:
    opening_choices: ClassVar[list[questionary.Choice]] = [
        questionary.Choice(title="All public Canvases", value=CanvasFilter(visibility="public", select_all=True)),
        questionary.Choice(title="Selected public Canvases", value=CanvasFilter(visibility="public", select_all=False)),
        questionary.Choice(title="All by given user", value=CanvasFilter(created_by="user", select_all=True)),
        questionary.Choice(title="Selected by given user", value=CanvasFilter(created_by="user", select_all=False)),
        questionary.Choice(title="All Canvases", value=CanvasFilter(visibility=None, select_all=True)),
    ]

    def __init__(self, client: ToolkitClient) -> None:
        self.client = client

    def select_external_ids(self) -> list[str]:
        select_filter = self._select_filter()

        return self._select_external_ids(select_filter)

    @classmethod
    def _select_filter(cls) -> CanvasFilter:
        user_response = questionary.select(
            "Which Canvases do you want to select?",
            choices=cls.opening_choices,
        ).unsafe_ask()
        return user_response

    def _select_external_ids(self, select_filter: CanvasFilter) -> list[str]:
        available_canvases = self.client.canvas.list(filter=select_filter.as_dms_filter(), limit=-1)
        if select_filter.select_all and select_filter.created_by is None:
            return [canvas.external_id for canvas in available_canvases]
        users = self.client.iam.user_profiles.list(limit=-1)
        display_name_by_user_identifier = {user.user_identifier: user.display_name or "missing" for user in users}
        if select_filter.created_by == "user":
            canvas_by_user: dict[str, list[Canvas]] = defaultdict(list)
            for canvas in available_canvases:
                canvas_by_user[canvas.created_by].append(canvas)

            user_response = questionary.select(
                "Which user do you want to select Canvases for?",
                choices=[
                    questionary.Choice(
                        title=f"{user.display_name} ({user.user_identifier})",
                        value=canvas_by_user[user.user_identifier],
                    )
                    for user in users
                    if user.user_identifier in canvas_by_user
                ],
            ).unsafe_ask()
            available_canvases = user_response

        if select_filter.select_all:
            return [canvas.external_id for canvas in available_canvases]

        if not available_canvases:
            raise ToolkitValueError("No Canvases available to select. Aborting.")

        selected_canvases = questionary.checkbox(
            "Select Canvases",
            choices=[
                questionary.Choice(
                    title=f"{canvas.name} (Created by {display_name_by_user_identifier[canvas.created_by]!r}, last updated {canvas.updated_at!r})",
                    value=canvas.external_id,
                )
                for canvas in available_canvases
            ],
            validate=lambda selected: "You must select at least one Canvas." if not selected else True,
        ).unsafe_ask()

        return selected_canvases or []


@dataclass
class ChartFilter:
    visibility: Visibility | None = None
    owned_by: Literal["user"] | None = None
    select_all: bool = False


class InteractiveChartSelect:
    opening_choices: ClassVar[list[questionary.Choice]] = [
        questionary.Choice(title="All public Charts", value=ChartFilter(visibility="PUBLIC", select_all=True)),
        questionary.Choice(title="Selected public Charts", value=ChartFilter(visibility="PUBLIC", select_all=False)),
        questionary.Choice(title="All owned by given user", value=ChartFilter(owned_by="user", select_all=True)),
        questionary.Choice(title="Selected owned by given user", value=ChartFilter(owned_by="user", select_all=False)),
    ]

    def __init__(self, client: ToolkitClient) -> None:
        self.client = client

    def select_external_ids(self) -> list[str]:
        select_filter = self._select_filter()
        return self._select_external_ids(select_filter)

    @classmethod
    def _select_filter(cls) -> ChartFilter:
        user_response = questionary.select(
            "Which Charts do you want to select?",
            choices=cls.opening_choices,
        ).unsafe_ask()
        return user_response

    def _select_external_ids(self, select_filter: ChartFilter) -> list[str]:
        available_charts = self.client.charts.list(visibility=(select_filter.visibility or "PUBLIC"))
        if select_filter.select_all and select_filter.owned_by is None:
            return [chart.external_id for chart in available_charts]

        users = self.client.iam.user_profiles.list(limit=-1)
        display_name_by_user_identifier = {
            user.user_identifier: user.display_name for user in users if user.display_name
        }
        if select_filter.owned_by == "user":
            available_charts = self._select_charts_by_user(available_charts, users)
        if select_filter.select_all:
            return [chart.external_id for chart in available_charts]
        selected_charts = questionary.checkbox(
            "Select Charts",
            choices=[
                questionary.Choice(
                    title=f"{chart.data.name} (Created by "
                    f"{display_name_by_user_identifier.get(chart.owner_id, 'missing')!r} "
                    f"- {ms_to_datetime(chart.last_updated_time)})",
                    value=chart.external_id,
                )
                for chart in available_charts
            ],
            validate=lambda selected: "You must select at least one Chart." if not selected else True,
        ).unsafe_ask()
        return selected_charts or []

    @classmethod
    def _select_charts_by_user(
        cls, available_charts: list[ChartResponse], users: UserProfileList
    ) -> list[ChartResponse]:
        chart_by_user: dict[str, list[ChartResponse]] = defaultdict(list)
        for chart in available_charts:
            chart_by_user[chart.owner_id].append(chart)
        user_response = questionary.select(
            "Which user do you want to select Charts for?",
            choices=[
                questionary.Choice(
                    title=f"{user.display_name} ({user.user_identifier})",
                    value=chart_by_user[user.user_identifier],
                )
                for user in users
                if user.user_identifier in chart_by_user
            ],
        ).unsafe_ask()
        return list(user_response)


class AssetCentricDestinationSelect:
    valid_destinations: tuple[str, ...] = get_args(AssetCentricDestinationType)

    @classmethod
    def validate(cls, destination_type: str) -> AssetCentricDestinationType:
        if destination_type not in cls.valid_destinations:
            raise ToolkitValueError(
                f"Invalid destination type: {destination_type!r}. Must be one of {humanize_collection(cls.valid_destinations)}."
            )
        # We validated the destination type above
        return destination_type  # type: ignore[return-value]

    @classmethod
    def select(cls) -> AssetCentricDestinationType:
        """Interactively select a destination type."""
        destination_type = questionary.select(
            "Select a destination type",
            choices=[questionary.Choice(title=dest.capitalize(), value=dest) for dest in cls.valid_destinations],
        ).unsafe_ask()
        # We only input valid destination types, so we can safely skip MyPy's type checking here
        return destination_type

    @classmethod
    def get(cls, destination_type: str | None = None) -> AssetCentricDestinationType:
        """Get the destination type, either from the input or interactively."""
        if destination_type is None:
            return cls.select()
        return cls.validate(destination_type)


@dataclass
class ViewSelectFilter:
    """Filter for creating view options when selecting views for instance selection.

    This filter is used to narrow down the list of views presented to the user when they are asked to select a view.

    Args:
        strategy: The strategy used to create the view options.
            - "SchemaSpace": Ask the user to select a schema space first.
        include_global: Whether to include global views in the options.
        schema_space : The schema space to look for view in.
        instance_type: If set, only include views that can be used for the given instance type ("node", "edge", or "all").
        mapped_container: If set, only include views that have the given container in their mapped containers.

    """

    strategy: Literal["schemaSpace", "dataModel"] = "schemaSpace"
    include_global: bool | None = None
    schema_space: str | None = None
    instance_type: Literal["node", "edge", "all"] | None = None
    mapped_container: ContainerReference | None = None
    data_model: DataModelReference | None = None

    def __str__(self) -> str:
        message: list[str] = []
        if self.include_global:
            message.append("including global")
        if self.schema_space:
            message.append(f"in schema space {self.schema_space!r}")
        if self.instance_type and self.instance_type != "all":
            message.append(f"for instance type {self.instance_type!r}")
        if self.mapped_container:
            message.append(f"with mapped container {self.mapped_container!r}")
        return f"Views {humanize_collection(message)}"


class DataModelingSelect:
    """A utility class to select Data Modeling nodes interactively."""

    def __init__(self, client: ToolkitClient, operation: str, console: Console | None = None) -> None:
        self.client = client
        self.operation = operation
        self.console = console or Console()
        self._available_spaces: list[SpaceResponse] | None = None

    @cached_property
    def stats_by_space(self) -> dict[str, SpaceStatistics]:
        result = self.client.data_modeling.statistics.spaces.list()
        return {stat.space: stat for stat in result}

    @overload
    def select_view(
        self,
        multiselect: Literal[False] = False,
        message: str | None = None,
        filter: ViewSelectFilter | None = None,
    ) -> ViewResponse: ...

    @overload
    def select_view(
        self,
        multiselect: Literal[True],
        message: str | None = None,
        filter: ViewSelectFilter | None = None,
    ) -> list[ViewResponse]: ...

    def select_view(
        self,
        multiselect: bool = False,
        message: str | None = None,
        filter: ViewSelectFilter | None = None,
    ) -> ViewResponse | list[ViewResponse]:
        """Select one or more views interactively.

        Args:
            multiselect: Whether to allow selecting multiple views.
             message: The message to display when prompting for a view. If None, a default message
                will be used.
            filter: Optional filter to apply when listing views. If None, no additional filtering will be applied.
        Returns:
            The selected view(s).

        Raises:
            ToolkitValueError: If no view(s) are selected or if the selected view(s) are not valid.
            ToolkitMissingResourceError: If no views are found in the selected space.

        """
        views = self._get_view_options(filter or ViewSelectFilter())
        question = message or f"Which view do you want to use to select instances to {self.operation}?"
        choices = [
            Choice(title=f"{view.space}:{view.external_id}(version={view.version})", value=view) for view in views
        ]
        if multiselect:
            selected_views = questionary.checkbox(
                question,
                choices=choices,
                validate=lambda choises: True if choises else "You must select at least one view.",
            ).unsafe_ask()
            if not isinstance(selected_views, list) or not all(isinstance(v, ViewResponse) for v in selected_views):
                raise ToolkitValueError(f"Selected views is not a valid list of View objects: {selected_views!r}")
            return selected_views
        else:
            selected_views = questionary.select(question, choices=choices).unsafe_ask()
            if not isinstance(selected_views, ViewResponse):
                raise ToolkitValueError(f"Selected view is not a valid View object: {selected_views!r}")
            return selected_views

    def _get_view_options(self, filter: ViewSelectFilter) -> list[ViewResponse]:
        if filter.strategy == "schemaSpace":
            selected_space = (
                filter.schema_space
                or self.select_schema_space(
                    filter.include_global or False,
                    message=f"In which Spaces is the view you will use to select instances to {self.operation}?",
                ).space
            )
            views = self.client.tool.views.list(
                filter=ViewFilter(
                    space=selected_space,
                    include_inherited_properties=True,
                    include_global=filter.include_global,
                ),
                limit=None,
            )
        elif filter.strategy == "dataModel":
            if filter.data_model:
                data_models = self.client.tool.data_models.retrieve([filter.data_model], inline_views=True)
                if not data_models or not data_models[0].views:
                    raise ToolkitMissingResourceError(
                        f"No views found in data model {filter.data_model!r} for filter: {filter!s}"
                    )
                datamodel = data_models[0]
            else:
                datamodel = self.select_data_model(
                    inline_views=True,
                    message=f"Select the data model through which to {self.operation}:",
                    schema_space=filter.schema_space,
                    include_global=filter.include_global,
                )
            views = datamodel.views or []
            parents = {parent for view in views for parent in view.implements or []}
            # We only allow the user to select child views
            views = [view for view in views if view.as_id() not in parents]
        else:
            raise NotImplementedError(f"Strategy {filter.strategy} is not implemented.")
        views = [
            view
            for view in views
            if filter.instance_type in (None, "all", view.used_for)
            and (filter.mapped_container is None or filter.mapped_container in view.mapped_containers)
        ]
        if not views:
            raise ToolkitMissingResourceError(f"No views found that match the filter: {filter!s}")
        return views

    @overload
    def select_data_model(
        self,
        inline_views: Literal[False] = False,
        message: str | None = None,
        schema_space: str | None = None,
        include_global: bool | None = None,
    ) -> DataModelResponse: ...

    @overload
    def select_data_model(
        self,
        inline_views: Literal[True],
        message: str | None = None,
        schema_space: str | None = None,
        include_global: bool | None = None,
    ) -> DataModelResponseWithViews: ...

    def select_data_model(
        self,
        inline_views: Literal[True, False] = False,
        message: str | None = None,
        schema_space: str | None = None,
        include_global: bool | None = None,
    ) -> DataModelResponse | DataModelResponseWithViews:
        datamodels = self.client.tool.data_models.list(
            inline_views=inline_views,
            filter=DataModelFilter(space=schema_space, all_versions=False, include_global=include_global),
            limit=None,
        )
        if not datamodels:
            raise ToolkitMissingResourceError(
                f"No data models found {f'in schema space {schema_space!r}' if schema_space else ''}."
            )

        message = message or f"Select a data model to {self.operation}:"
        selected_datamodel = questionary.select(
            message,
            choices=[
                Choice(
                    title=f"{dm.space}:{dm.external_id} (version={dm.version})",
                    value=dm,
                )
                for dm in datamodels
            ],
        ).unsafe_ask()
        if not isinstance(selected_datamodel, DataModelResponse | DataModelResponseWithViews):
            raise ToolkitValueError(
                f"Selected data model is not a valid DataModelResponse object: {selected_datamodel!r}"
            )
        return selected_datamodel

    def select_schema_space(self, include_global: bool, message: str | None = None) -> SpaceResponse:
        message = message or f"Select the space to {self.operation}:"
        spaces = self._get_available_spaces(include_global)
        schema_spaces = {
            space
            for space, stats in self.stats_by_space.items()
            if (stats.containers + stats.views + stats.data_models) > 0
        }
        spaces = [space for space in spaces if space.space in schema_spaces or space.is_global]
        if not spaces:
            raise ToolkitMissingResourceError("No spaces with schema (containers, views, or data models) found.")
        selected_space = questionary.select(
            message,
            [Choice(title=space.space, value=space) for space in sorted(spaces, key=lambda s: s.space)],
        ).unsafe_ask()
        if not isinstance(selected_space, SpaceResponse):
            raise ToolkitValueError(f"Selected space is not a valid Space object: {selected_space!r}")
        return selected_space

    def select_instance_type(
        self, view_used_for: Literal["node", "edge", "all"] | None = None, message: str | None = None
    ) -> Literal["node", "edge"]:
        """Selects an instance type (node or edge) interactively.

        Args:
            view_used_for: If 'node' or 'edge', that type is returned directly.
                           If 'all' or None, the user is prompted to select.

        Returns:
            The selected instance type, either 'node' or 'edge'.
        """
        if view_used_for is not None and view_used_for != "all":
            return view_used_for
        selected_instance_type = questionary.select(
            message or f"What type of instances do you want to {self.operation}?",
            choices=[
                Choice(title="Nodes", value="node"),
                Choice(title="Edges", value="edge"),
            ],
        ).unsafe_ask()

        if selected_instance_type not in ("node", "edge"):
            raise ToolkitValueError(f"Selected instance type is not valid: {selected_instance_type!r}")
        return selected_instance_type

    def select_space_type(self) -> Literal["instance", "schema", "empty"]:
        selected_space_type = questionary.select(
            f"What type of spaces do you want to {self.operation}?",
            choices=[
                Choice(title="Instance spaces with instances", value="instance"),
                Choice(title="Schema spaces with schema", value="schema"),
                Choice(title="Empty spaces without schema or instances", value="empty"),
            ],
        ).unsafe_ask()
        if selected_space_type not in ["instance", "schema", "empty"]:
            raise ToolkitValueError(f"Selected space type is not valid: {selected_space_type!r}")
        return selected_space_type

    @overload
    def select_instance_space(
        self,
        multiselect: Literal[False],
        selected_view: ViewReference | None = None,
        instance_type: Literal["node", "edge"] | None = None,
        message: str | None = None,
        include_empty: bool = False,
    ) -> str: ...

    @overload
    def select_instance_space(
        self,
        multiselect: Literal[True] = True,
        selected_view: ViewReference | None = None,
        instance_type: Literal["node", "edge"] | None = None,
        message: str | None = None,
        include_empty: bool = False,
    ) -> list[str]: ...

    def select_instance_space(
        self,
        multiselect: bool = True,
        selected_view: ViewReference | None = None,
        instance_type: Literal["node", "edge"] | None = None,
        message: str | None = None,
        include_empty: bool = False,
    ) -> str | list[str]:
        if (selected_view is not None and instance_type is None) or (
            selected_view is None and instance_type is not None
        ):
            raise ToolkitValueError("Both selected_view and instance_type must be provided or both must be None.")
        if selected_view is not None and instance_type is not None:
            all_spaces = self._get_available_spaces(include_global=False)
            count_by_space = self._get_instance_count_in_view_by_space(all_spaces, selected_view, instance_type)
        else:
            count_by_space = {
                space: stats.nodes + stats.edges
                for space, stats in self.stats_by_space.items()
                if (stats.nodes + stats.edges) > 0
                or (
                    include_empty
                    and (stats.nodes + stats.edges + stats.containers + stats.views + stats.data_models) == 0
                )
            }

        if not count_by_space:
            raise ToolkitMissingResourceError(
                f"No instances found in any space for the view {selected_view!r} with instance type {instance_type!r}."
            )
        if len(count_by_space) == 1:
            selected_spaces = next(iter(count_by_space.keys()))
            self.console.print(f"Only one space with instances found: {selected_spaces!r}. Using this space.")
            return [selected_spaces] if multiselect else selected_spaces

        if not message:
            message = f"In which Space{'(s)' if multiselect else ''} do you want to {self.operation} instances?"
        choices = [
            Choice(title=f"{space} ({count:,} instances)", value=space)
            for space, count in sorted(count_by_space.items(), key=lambda item: item[1], reverse=True)
        ]
        if multiselect:
            selected_spaces = questionary.checkbox(
                message,
                choices=choices,
                validate=lambda choices: True if choices else "You must select at least one space.",
            ).unsafe_ask()
        else:
            selected_spaces = questionary.select(message, choices=choices).unsafe_ask()
        if selected_spaces is None or (isinstance(selected_spaces, list) and len(selected_spaces) == 0):
            raise ToolkitValueError(
                f"No instance space selected to {self.operation}. Use arrow keys to navigate and space key to select. Press enter to confirm."
            )
        return selected_spaces

    @overload
    def select_empty_spaces(self, multiselect: Literal[False]) -> str: ...

    @overload
    def select_empty_spaces(self, multiselect: Literal[True]) -> list[str]: ...

    def select_empty_spaces(self, multiselect: bool = True) -> str | list[str]:
        empty_spaces = [
            space
            for space, stats in self.stats_by_space.items()
            if (stats.nodes + stats.edges + stats.containers + stats.views + stats.data_models) == 0
        ]
        if not empty_spaces:
            raise ToolkitMissingResourceError("No empty spaces found.")
        if len(empty_spaces) == 1:
            selected_space = empty_spaces[0]
            self.console.print(f"Only one empty space found: {selected_space!r}. Using this space.")
            return [selected_space] if multiselect else selected_space

        message = f"In which empty Space{'(s)' if multiselect else ''} do you want to {self.operation}?"
        choices = [Choice(title=f"{space}", value=space) for space in sorted(empty_spaces)]
        if multiselect:
            selected_spaces = questionary.checkbox(
                message,
                choices=choices,
                validate=lambda choises: True if choises else "You must select at least one space.",
            ).unsafe_ask()
        else:
            selected_spaces = questionary.select(message, choices=choices).unsafe_ask()
        if selected_spaces is None or (isinstance(selected_spaces, list) and len(selected_spaces) == 0):
            raise ToolkitValueError(f"No empty space selected to {self.operation}. Aborting.")
        return selected_spaces

    def _get_instance_count_in_view_by_space(
        self, all_spaces: list[SpaceResponse], view_id: ViewReference, instance_type: Literal["node", "edge"]
    ) -> dict[str, float]:
        count_by_space: dict[str, float] = {}
        result = self.client.data_modeling.statistics.project()
        read_limit = result.concurrent_read_limit

        with ThreadPoolExecutor(max_workers=read_limit // 2) as executor:
            results = executor.map(
                partial(self._instance_count_space, view_id=view_id, instance_type=instance_type),
                (space.space for space in all_spaces),
            )
            for space_name, count in results:
                if count > 0:
                    count_by_space[space_name] = count

        return count_by_space

    @lru_cache
    def _instance_count_space(
        self, space: str, view_id: ViewReference, instance_type: Literal["node", "edge"]
    ) -> tuple[str, float]:
        """Get the count of instances in a specific space for a given view and instance type."""
        sdk_view_id = ViewId(space=view_id.space, external_id=view_id.external_id, version=view_id.version)
        return space, self.client.data_modeling.instances.aggregate(
            sdk_view_id, Count("externalId"), instance_type=instance_type, space=space
        ).value or 0.0

    def _get_available_spaces(self, include_global: bool = False) -> list[SpaceResponse]:
        if self._available_spaces is None:
            self._available_spaces = self.client.tool.spaces.list(include_global=True)
        if include_global:
            return self._available_spaces
        return [space for space in self._available_spaces if not space.is_global]


class ResourceViewMappingInteractiveSelect:
    def __init__(self, client: ToolkitClient, operation: str) -> None:
        self.client = client
        self.operation = operation

    def select_resource_view_mapping(self, resource_type: str) -> ResourceViewMappingResponse:
        """Select a Resource View Mapping interactively.

        Args:
            resource_type: The resource type to filter Resource View Mappings by.
        Returns:
            The selected Resource View Mapping.
        """
        mappings = self.client.migration.resource_view_mapping.list(resource_type=resource_type, limit=-1)
        if not mappings:
            raise ToolkitMissingResourceError(f"No Resource View Mappings found for resource type {resource_type!r}.")
        choices = [
            Choice(
                title=f"{mapping.external_id} ({mapping.view_id!r})",
                value=mapping,
            )
            for mapping in mappings
        ]
        selected_mapping = questionary.select(
            f"Which Resource View Mapping do you want to use to {self.operation}? [identifier] (ingestion view)",
            choices=choices,
        ).unsafe_ask()
        if not isinstance(selected_mapping, ResourceViewMappingResponse):
            raise ToolkitValueError(
                f"Selected Resource View Mapping is not a valid ResourceViewMappingResponse object: {selected_mapping!r}"
            )
        return selected_mapping


class ThreeDInteractiveSelect:
    def __init__(self, client: ToolkitClient, operation: str) -> None:
        self.client = client
        self.operation = operation

    def select_three_d_models(
        self, model_type: Literal["classic", "dm"] | None = None
    ) -> list[ThreeDModelClassicResponse]:
        """Select multiple 3D models interactively."""
        if model_type is None:
            model_type = questionary.select(
                f"What type of 3D models do you want to {self.operation}?",
                choices=[
                    Choice(title="Classic models", value="classic"),
                    Choice(title="Data modeling 3D", value="dm"),
                ],
            ).unsafe_ask()
        published = questionary.select(
            f"Do you want to {self.operation} published or unpublished 3D models?",
            choices=[
                Choice(title="Published models", value=True),
                Choice(title="Unpublished models", value=False),
                Choice(title="Both published and unpublished models", value=None),
            ],
        ).unsafe_ask()

        models = self.client.tool.three_d.models_classic.list(
            published=published, include_revision_info=True, limit=None
        )
        if model_type == "classic":
            models = [model for model in models if model.space is None]
        else:
            models = [model for model in models if model.space is not None]
        if not models:
            raise ToolkitMissingResourceError(
                f"No 3D models found for type {model_type!r} with published={published!r}."
            )

        choices = [Choice(title=f"{model.name} ({model.id})", value=model) for model in models]
        selected_models = questionary.checkbox(
            f"Select 3D models to {self.operation}:",
            choices=choices,
            validate=lambda choices: True if choices else "You must select at least one 3D model.",
        ).unsafe_ask()
        if selected_models is None or len(selected_models) == 0:
            raise ToolkitValueError("No 3D models selected.")
        return selected_models


class APMConfigInteractiveSelect:
    """Interactive select for APM Config (Infield V1) configurations."""

    def __init__(self, client: ToolkitClient, operation: str) -> None:
        self.client = client
        self.operation = operation

    def _get_choices(self) -> list[Choice]:
        """Fetch APM configs and create choices for selection.

        Returns:
            List of questionary Choice objects for the available APM configs.

        Raises:
            ToolkitMissingResourceError: If no APM configs are found.
        """
        configs = self.client.infield.apm_config.list(limit=None)
        if not configs:
            raise ToolkitMissingResourceError("No APM configs found.")

        return [
            Choice(
                title=f"{config.name} ({config.external_id})" if config.name else f"({config.external_id})",
                value=config,
            )
            for config in configs
        ]

    def select_apm_configs(self) -> list[APMConfigResponse]:
        """Select multiple APM configs interactively.

        Returns:
            List of selected APMConfigResponse objects.

        Raises:
            ToolkitMissingResourceError: If no APM configs are found.
            ToolkitValueError: If no APM configs are selected.
        """
        choices = self._get_choices()
        selected_configs = questionary.checkbox(
            f"Select APM configs to {self.operation}:",
            choices=choices,
            validate=lambda choices: True if choices else "You must select at least one APM config.",
        ).unsafe_ask()
        if selected_configs is None or len(selected_configs) == 0:
            raise ToolkitValueError("No APM configs selected.")
        return selected_configs

    def select_apm_config(self) -> APMConfigResponse:
        """Select a single APM config interactively.

        Returns:
            The selected APMConfigResponse object.

        Raises:
            ToolkitMissingResourceError: If no APM configs are found.
            ToolkitValueError: If no APM config is selected.
        """
        choices = self._get_choices()
        selected_config = questionary.select(
            f"Select an APM config to {self.operation}:",
            choices=choices,
        ).unsafe_ask()
        if selected_config is None:
            raise ToolkitValueError("No APM config selected.")
        if not isinstance(selected_config, APMConfigResponse):
            raise ToolkitValueError(f"Selected APM config is not a valid APMConfigResponse object: {selected_config!r}")
        return selected_config
