from __future__ import annotations

import datetime
import math
import time
import warnings
from abc import ABC
from collections import defaultdict
from collections.abc import Collection, MutableSequence, Iterable, Sequence
from contextlib import suppress
from dataclasses import dataclass, field
from typing import (
    cast,
    ClassVar,
    Generic,
    Any,
    Iterator,
    TypeVar,
    overload,
    Union,
    SupportsIndex,
    Literal,
)

from cognite.client import CogniteClient
from cognite.client import data_modeling as dm
from cognite.client.data_classes._base import CogniteObject
from cognite.client.data_classes.aggregations import Count
from cognite.client.data_classes.data_modeling.instances import Instance
from cognite.client.exceptions import CogniteAPIError

from {{ top_level_package }}.data_classes._core.base import (
    DomainModelList,
    T_DomainList,
    DomainRelationList,
    DomainModelCore,
    T_DomainModelList,
    DomainRelation,
    DomainModel,
)
from {{ top_level_package }}.data_classes._core.constants import (
    _NotSetSentinel,
    DEFAULT_QUERY_LIMIT,
    {% if has_default_instance_space %}
    DEFAULT_INSTANCE_SPACE,
    {% endif %}
    ACTUAL_INSTANCE_QUERY_LIMIT,
    INSTANCE_QUERY_LIMIT,
    IN_FILTER_CHUNK_SIZE,
    MINIMUM_ESTIMATED_SECONDS_BEFORE_PRINT_PROGRESS,
    PRINT_PROGRESS_PER_N_NODES,
    SEARCH_LIMIT,
)
from {{ top_level_package }}.data_classes._core.helpers import as_node_id


T_DomainListEnd = TypeVar("T_DomainListEnd", bound=Union[DomainModelList, DomainRelationList], covariant=True)


class QueryCore(Generic[T_DomainList, T_DomainListEnd]):
    _view_id: ClassVar[dm.ViewId]
    _result_list_cls_end: type[T_DomainListEnd]
    _result_cls: ClassVar[type[DomainModelCore]]

    def __init__(
        self,
        created_types: set[type],
        creation_path: "list[QueryCore]",
        client: CogniteClient,
        result_list_cls: type[T_DomainList],
        expression: dm.query.ResultSetExpression | None = None,
        view_filter: dm.filters.Filter | None = None,
        connection_name: str | None = None,
        connection_type: Literal["reverse-list"] | None = None,
        reverse_expression: dm.query.ResultSetExpression | None = None,
    ):
        created_types.add(type(self))
        self._creation_path = creation_path[:] + [self]
        self._client = client
        self._result_list_cls = result_list_cls
        self._view_filter = view_filter
        self._expression = expression or dm.query.NodeResultSetExpression()
        self._reverse_expression = reverse_expression
        self._connection_name = connection_name
        self._connection_type = connection_type
        self._filter_classes: list[Filtering] = []

    @property
    def _connection_names(self) -> set[str]:
        return {step._connection_name for step in self._creation_path if step._connection_name}

    @property
    def _is_reverseable(self) -> bool:
        return self._reverse_expression is not None

    def __getattr__(self, item: str) -> Any:
        if item in self._connection_names:
            nodes = [step._result_cls.__name__ for step in self._creation_path]
            raise ValueError(f"Circular reference detected. Cannot query a circular reference: {nodes}")
        elif self._connection_type == "reverse-list":
            raise ValueError(f"Cannot query across a reverse-list connection.")
        raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{item}'")

    def _assemble_filter(self) -> dm.filters.Filter | None:
        filters: list[dm.filters.Filter] = [self._view_filter] if self._view_filter else []
        for filter_cls in self._filter_classes:
            if item := filter_cls._as_filter():
                filters.append(item)
        return dm.filters.And(*filters) if filters else None

    def _create_sort(self) -> list[dm.InstanceSort] | None:
        filters: list[tuple[dm.InstanceSort, int]] = []
        for filter_cls in self._filter_classes:
            item, priority = filter_cls._as_sort()
            if item:
                filters.append((item, priority))
        return [item for item, _ in sorted(filters, key=lambda x: x[1])] if filters else None

    def _has_limit_1(self) -> bool:
        return any(filter_cls._has_limit_1 for filter_cls in self._filter_classes)

    def _repr_html_(self) -> str:
        nodes = [step._result_cls.__name__ for step in self._creation_path]
        edges = [step._connection_name or "missing" for step in self._creation_path[1:]]
        last_connection_name = self._connection_name or "missing"
        w = 120
        h = 40
        circles = "    \n".join(f'<circle cx="{i * w + 40}" cy="{h}" r="2" />' for i in range(len(nodes)))
        circle_text = "    \n".join(
            f'<text x="{i * w + 40}" y="{h}" dy="-10">{node}</text>' for i, node in enumerate(nodes)
        )
        arrows = "    \n".join(
            f'<path id="arrow-line"  marker-end="url(#head)" stroke-width="2" fill="none" stroke="black" d="M{i*w+40},{h}, {i*w + 150} {h}" />'
            for i in range(len(edges))
        )
        arrow_text = "    \n".join(
            f'<text x="{i*w+40+120/2}" y="{h}" dy="-5">{edge}</text>' for i, edge in enumerate(edges)
        )

        return f"""<h5>Query</h5>
<div>
<svg height="50" xmlns="http://www.w3.org/2000/svg">
  <defs>
    <marker
      id='head'
      orient="auto"
      markerWidth='3'
      markerHeight='4'
      refX='0.1'
      refY='2'
    >
      <path d='M0,0 V4 L2,2 Z' fill="black" />
    </marker>
  </defs>

    {arrows}

<g stroke="black" stroke-width="3" fill="black">
    {circles}
</g>
<g font-size="10" font-family="sans-serif" text-anchor="middle">
    {arrow_text}
</g>
<g font-size="10" font-family="sans-serif" text-anchor="middle">
    {circle_text}
</g>
</svg>
</div>
<p>Call <em>.list_full()</em> to return a list of {nodes[0].title()} and
<em>.list_{last_connection_name}()</em> to return a list of {nodes[-1].title()}.</p>
"""


class NodeQueryCore(QueryCore[T_DomainModelList, T_DomainListEnd]):
    _result_cls: ClassVar[type[DomainModel]]

    def list_full(self, limit: int = DEFAULT_QUERY_LIMIT) -> T_DomainModelList:
        builder = self._create_query(limit, self._result_list_cls, return_step="first", try_reverse=True)
        builder.execute_query(self._client, remove_not_connected=True)
        return builder.unpack()

    def _list(self, limit: int = DEFAULT_QUERY_LIMIT) -> T_DomainListEnd:
        builder = self._create_query(limit, cast(type[DomainModelList], self._result_list_cls_end), return_step="last")
        set_limit = builder[0].max_retrieve_limit
        for step in builder[:-1]:
            step.select = None
        # We cannot use the limit above, as the _create_query will set limit=1 if latest/earliest is used
        builder[-1].max_retrieve_limit = set_limit
        builder.execute_query(self._client, remove_not_connected=False)
        return builder.unpack()

    def _dump_yaml(self) -> str:
        return self._create_query(DEFAULT_QUERY_LIMIT, self._result_list_cls)._dump_yaml()

    def _create_query(
        self,
        limit: int,
        result_list_cls: type[DomainModelList],
        return_step: Literal["first", "last"] | None = None,
        try_reverse: bool = False,
    ) -> DataClassQueryBuilder:
        builder = DataClassQueryBuilder(result_list_cls, return_step=return_step)
        from_: str | None = None
        first: bool = True
        is_last_reverse_list = False
        for item in self._creation_path:
            if is_last_reverse_list:
                raise ValueError(
                    "Cannot traverse past reverse direct relation of list. "
                    "This is a limitation with the modeling implementation of your data model."
                    "To do this query, you need to reimplement the data model and use an edge to "
                    "implement this connection instead of a reverse direct relation"
                )
            name = builder.create_name(from_)
            if first and item._has_limit_1():
                if limit != DEFAULT_QUERY_LIMIT:
                    warnings.warn(
                        "When selecting earliest and latest, the limit is ignored.", UserWarning, stacklevel=2
                    )
                max_retrieve_limit = 1
            elif first:
                max_retrieve_limit = limit
            else:
                max_retrieve_limit = -1

            step: QueryStep
            if isinstance(item, NodeQueryCore) and isinstance(item._expression, dm.query.NodeResultSetExpression):
                step = NodeQueryStep(
                    name=name,
                    expression=item._expression,
                    result_cls=item._result_cls,
                    max_retrieve_limit=max_retrieve_limit,
                    connection_type=item._connection_type,
                )
                step.expression.from_ = from_
                step.expression.filter = item._assemble_filter()
                step.expression.sort = item._create_sort()
                builder.append(step)
            elif isinstance(item, NodeQueryCore) and isinstance(item._expression, dm.query.EdgeResultSetExpression):
                edge_name = name
                step = EdgeQueryStep(name=edge_name, expression=item._expression, max_retrieve_limit=max_retrieve_limit)
                step.expression.from_ = from_
                builder.append(step)

                name = builder.create_name(edge_name)
                node_step = NodeQueryStep(
                    name=name,
                    expression=dm.query.NodeResultSetExpression(
                        from_=edge_name,
                        filter=item._assemble_filter(),
                        sort=item._create_sort(),
                    ),
                    result_cls=item._result_cls,
                )
                builder.append(node_step)
            elif isinstance(item, EdgeQueryCore):
                step = EdgeQueryStep(
                    name=name,
                    expression=cast(dm.query.EdgeResultSetExpression, item._expression),
                    result_cls=item._result_cls,
                )
                step.expression.from_ = from_
                step.expression.filter = item._assemble_filter()
                step.expression.sort = item._create_sort()
                builder.append(step)
            else:
                raise TypeError(f"Unsupported query step type: {type(item._expression)}")

            is_last_reverse_list = item._connection_type == "reverse-list"
            first = False
            from_ = name
        return builder


class EdgeQueryCore(QueryCore[T_DomainList, T_DomainListEnd]):
    _result_cls: ClassVar[type[DomainRelation]]


{{ query_builder }}


class NodeQueryStep(QueryStep):
    def __init__(
        self,
        name: str,
        expression: dm.query.NodeResultSetExpression,
        result_cls: type[DomainModel],
        max_retrieve_limit: int = -1,
        select: dm.query.Select | None | type[_NotSetSentinel] = _NotSetSentinel,
        raw_filter: dm.Filter | None = None,
        connection_type: Literal["reverse-list"] | None = None,
    ):
        self.result_cls = result_cls
        super().__init__(name, expression, result_cls._view_id, max_retrieve_limit, select, raw_filter, connection_type)

    def unpack(self) -> dict[dm.NodeId{% if has_default_instance_space %} | str{% endif %}, DomainModel]:
        return {{'{'}}
            {% if has_default_instance_space %}
            (
                instance.as_id() if instance.space != DEFAULT_INSTANCE_SPACE else instance.external_id
            ){% else %}
            instance.as_id(){% endif %}: self.result_cls.from_instance(instance){% if has_default_instance_space +%}
           {%+ else %}
            {% endif %} for instance in cast(list[dm.Node], self.results)
        {{'}'}}

    @property
    def node_results(self) -> list[dm.Node]:
        return cast(list[dm.Node], self.results)

    @property
    def node_expression(self) -> dm.query.NodeResultSetExpression:
        return cast(dm.query.NodeResultSetExpression, self.expression)


class EdgeQueryStep(QueryStep):
    def __init__(
        self,
        name: str,
        expression: dm.query.EdgeResultSetExpression,
        result_cls: type[DomainRelation] | None = None,
        max_retrieve_limit: int = -1,
        select: dm.query.Select | None | type[_NotSetSentinel] = _NotSetSentinel,
        raw_filter: dm.Filter | None = None,
    ):
        self.result_cls = result_cls
        view_id = result_cls._view_id if result_cls is not None else None
        super().__init__(name, expression, view_id, max_retrieve_limit, select, raw_filter, None)

    def unpack(self) -> dict[dm.NodeId, list[dm.Edge | DomainRelation]]:
        output: dict[dm.NodeId, list[dm.Edge | DomainRelation]] = defaultdict(list)
        for edge in cast(list[dm.Edge], self.results):
            edge_source = edge.start_node if self.expression.direction == "outwards" else edge.end_node
            value = self.result_cls.from_instance(edge) if self.result_cls is not None else edge
            output[as_node_id(edge_source)].append(value)  # type: ignore[arg-type]
        return output

    @property
    def edge_results(self) -> list[dm.Edge]:
        return cast(list[dm.Edge], self.results)

    @property
    def edge_expression(self) -> dm.query.EdgeResultSetExpression:
        return cast(dm.query.EdgeResultSetExpression, self.expression)


class DataClassQueryBuilder(QueryBuilder, Generic[T_DomainModelList]):
    """This is a helper class to build and execute a query. It is responsible for
    doing the paging of the query and keeping track of the results."""

    def __init__(
        self,
        result_cls: type[T_DomainModelList] | None,
        steps: Collection[QueryStep] | None = None,
        return_step: Literal["first", "last"] | None = None,
    ):
        super().__init__(steps or [])
        self._result_list_cls = result_cls
        self._return_step: Literal["first", "last"] | None = return_step

    def unpack(self) -> T_DomainModelList:
        if self._result_list_cls is None:
            raise ValueError("No result class set, unable to unpack results")
        selected = [step for step in self if step.select is not None]
        if len(selected) == 0:
            return self._result_list_cls([])
        elif len(selected) == 1:
            # Validated in the append method
            if self._return_step == "first":
                selected_step = cast(NodeQueryStep, self[0])
            elif self._return_step == "last":
                selected_step = cast(NodeQueryStep, self[-1])
            else:
                raise ValueError(f"Invalid return_step: {self._return_step}")
            return self._result_list_cls(selected_step.unpack().values())
        # More than one step, we need to unpack the nodes and edges
        nodes_by_from: dict[str | None, dict[dm.NodeId | str, DomainModel]] = defaultdict(dict)
        edges_by_from: dict[str, dict[dm.NodeId, list[dm.Edge | DomainRelation]]] = defaultdict(dict)
        for step in reversed(self):
            # Validated in the append method
            from_ = cast(str, step.from_)
            if isinstance(step, EdgeQueryStep):
                edges_by_from[from_].update(step.unpack())
                if step.name in nodes_by_from:
                    nodes_by_from[from_].update(nodes_by_from[step.name])
                    del nodes_by_from[step.name]
            elif isinstance(step, NodeQueryStep):
                unpacked = step.unpack()
                nodes_by_from[from_].update(unpacked)  # type: ignore[arg-type]
                if step.name in nodes_by_from or step.name in edges_by_from:
                    step.result_cls._update_connections(
                        unpacked,  # type: ignore[arg-type]
                        nodes_by_from.get(step.name, {}),  # type: ignore[arg-type]
                        edges_by_from.get(step.name, {}),
                    )
        if self._return_step == "first":
            return self._result_list_cls(nodes_by_from[None].values())
        elif self._return_step == "last" and self[-1].from_ in nodes_by_from:
            return self._result_list_cls(nodes_by_from[self[-1].from_].values())
        elif self._return_step == "last":
            raise ValueError("Cannot return the last step when the last step is an edge query")
        else:
            raise ValueError(f"Invalid return_step: {self._return_step}")

    def append(self, __object: QueryStep, /) -> None:
        # Extra validation to ensure all assumptions are met
        if len(self) == 0:
            if __object.from_ is not None:
                raise ValueError("The first step should not have a 'from_' value")
            if self._result_list_cls is None:
                if self._return_step is None:
                    self._return_step = "first"
            else:
                if not isinstance(__object, NodeQueryStep):
                    raise ValueError("The first step should be a NodeQueryStep")
                # If the first step is a NodeQueryStep, and matches the instance
                # in the result_list_cls we can return the result from the first step
                # Alternative is result_cls is not set, then we also assume that the first step
                if self._return_step is None:
                    if __object.result_cls is self._result_list_cls._INSTANCE:
                        self._return_step = "first"
                    else:
                        # If not, we assume that the last step is the one we want to return
                        self._return_step = "last"
        else:
            if __object.from_ is None:
                raise ValueError("The 'from_' value should be set")
        super().append(__object)

    def extend(self, __iterable: Iterable[QueryStep], /) -> None:
        for item in __iterable:
            self.append(item)

    # The implementations below are to get proper type hints
    def __iter__(self) -> Iterator[QueryStep]:
        return super().__iter__()

    @overload
    def __getitem__(self, item: SupportsIndex) -> QueryStep: ...

    @overload
    def __getitem__(self, item: slice) -> DataClassQueryBuilder[T_DomainModelList]: ...

    def __getitem__(self, item: SupportsIndex | slice) -> QueryStep | DataClassQueryBuilder[T_DomainModelList]:
        value = super().__getitem__(item)
        if isinstance(item, slice):
            return DataClassQueryBuilder(self._result_list_cls, value)  # type: ignore[arg-type]
        return cast(QueryStep, value)


T_QueryCore = TypeVar("T_QueryCore")


class Filtering(Generic[T_QueryCore], ABC):
    counter: ClassVar[int] = 0

    def __init__(self, query: T_QueryCore, prop_path: list[str] | tuple[str, ...]) -> None:
        self._query = query
        self._prop_path = prop_path
        self._filter: dm.Filter | None = None
        self._sort: dm.InstanceSort | None = None
        self._sort_priority: int | None = None
        # Used for earliest/latest
        self._limit: int | None = None

    def _raise_if_filter_set(self):
        if self._filter is not None:
            raise ValueError("Filter has already been set")

    def _raise_if_sort_set(self):
        if self._sort is not None:
            raise ValueError("Sort has already been set")

    @classmethod
    def _get_sort_priority(cls) -> int:
        # This is used in case of multiple sorts, to ensure that the order is correct
        Filtering.counter += 1
        return Filtering.counter

    def _as_filter(self) -> dm.Filter | None:
        return self._filter

    def _as_sort(self) -> tuple[dm.InstanceSort | None, int]:
        return self._sort, self._sort_priority or 0

    @property
    def _has_limit_1(self) -> bool:
        return self._limit == 1

    def sort_ascending(self) -> T_QueryCore:
        self._raise_if_sort_set()
        self._sort = dm.InstanceSort(self._prop_path, "ascending")
        self._sort_priority = self._get_sort_priority()
        return self._query

    def sort_descending(self) -> T_QueryCore:
        self._raise_if_sort_set()
        self._sort = dm.InstanceSort(self._prop_path, "descending")
        self._sort_priority = self._get_sort_priority()
        return self._query


class StringFilter(Filtering[T_QueryCore]):
    def equals(self, value: str) -> T_QueryCore:
        self._raise_if_filter_set()
        self._filter = dm.filters.Equals(self._prop_path, value)
        return self._query

    def prefix(self, prefix: str) -> T_QueryCore:
        self._raise_if_filter_set()
        self._filter = dm.filters.Prefix(self._prop_path, prefix)
        return self._query

    def in_(self, values: list[str]) -> T_QueryCore:
        self._raise_if_filter_set()
        self._filter = dm.filters.In(self._prop_path, values)
        return self._query


class BooleanFilter(Filtering[T_QueryCore]):
    def equals(self, value: bool) -> T_QueryCore:
        self._raise_if_filter_set()
        self._filter = dm.filters.Equals(self._prop_path, value)
        return self._query


class IntFilter(Filtering[T_QueryCore]):
    def range(self, gte: int | None, lte: int | None) -> T_QueryCore:
        self._raise_if_filter_set()
        self._filter = dm.filters.Range(self._prop_path, gte=gte, lte=lte)
        return self._query


class FloatFilter(Filtering[T_QueryCore]):
    def range(self, gte: float | None, lte: float | None) -> T_QueryCore:
        self._raise_if_filter_set()
        self._filter = dm.filters.Range(self._prop_path, gte=gte, lte=lte)
        return self._query


class TimestampFilter(Filtering[T_QueryCore]):
    def range(self, gte: datetime.datetime | None, lte: datetime.datetime | None) -> T_QueryCore:
        self._raise_if_filter_set()
        self._filter = dm.filters.Range(
            self._prop_path,
            gte=gte.isoformat(timespec="milliseconds") if gte else None,
            lte=lte.isoformat(timespec="milliseconds") if lte else None,
        )
        return self._query

    def earliest(self) -> T_QueryCore:
        self._raise_if_sort_set()
        self._sort = dm.InstanceSort(self._prop_path, "ascending")
        self._sort_priority = self._get_sort_priority()
        self._limit = 1
        return self._query

    def latest(self) -> T_QueryCore:
        self._raise_if_sort_set()
        self._sort = dm.InstanceSort(self._prop_path, "descending")
        self._sort_priority = self._get_sort_priority()
        self._limit = 1
        return self._query


class DateFilter(Filtering[T_QueryCore]):
    def range(self, gte: datetime.date | None, lte: datetime.date | None) -> T_QueryCore:
        self._raise_if_filter_set()
        self._filter = dm.filters.Range(
            self._prop_path,
            gte=gte.isoformat() if gte else None,
            lte=lte.isoformat() if lte else None,
        )
        return self._query

    def earliest(self) -> T_QueryCore:
        self._raise_if_sort_set()
        self._sort = dm.InstanceSort(self._prop_path, "ascending")
        self._sort_priority = self._get_sort_priority()
        self._limit = 1
        return self._query

    def latest(self) -> T_QueryCore:
        self._raise_if_sort_set()
        self._sort = dm.InstanceSort(self._prop_path, "descending")
        self._sort_priority = self._get_sort_priority()
        self._limit = 1
        return self._query
