from __future__ import annotations

import datetime
from collections.abc import Sequence
from typing import TYPE_CHECKING, cast

from cognite.client import data_modeling as dm, CogniteClient

from {{ top_level_package }}.data_classes import (
    DomainModelCore,
    {{ data_class.read_name }},
    {% for class_ in unique_edge_data_classes %}
    {{ class_.read_name }},
    {% endfor %}
    {% for field in data_class.one_to_one_direct_relations_with_source %}
    {{ field.destination_class.read_name }},
    {% endfor %}
)
{% for edge_api in edge_apis %}
from {{ top_level_package }}.data_classes.{{ edge_api.end_class.file_name }} import (
    {{ edge_api.end_class.read_name }},
    {{ edge_api.end_class.filter_name }},
)
{% endfor %}
from {{ top_level_package }}._api._core import (
    DEFAULT_QUERY_LIMIT,
    EdgeQueryStep,
    NodeQueryStep,
    DataClassQueryBuilder,
    QueryAPI,
    T_DomainModelList,
    {% if not query_api.is_only_one_to_many_edges  %}
    _create_edge_filter,
    {% endif %}
)

{% for class_ in unique_edge_data_classes %}
from {{ top_level_package }}.data_classes.{{ class_.file_name }} import (
    {{ class_.filter_name }},
)
{% endfor %}
{% if has_edge_api_dependencies %}
if TYPE_CHECKING:
    {% for edge_api in edge_apis %}
    {% if edge_api.end_view_id != api_class.view_id %}
    from {{ top_level_package }}._api.{{ edge_api.query.file_name}} import {{ edge_api.query.name }}
    {% endif %}
    {% endfor %}
{% endif %}


class {{ query_api.name }}(QueryAPI[T_DomainModelList]):
    _view_id = dm.ViewId("{{ data_class.view_id.space }}", "{{ data_class.view_id.external_id }}", "{{ data_class.view_id.version }}")

    def __init__(
        self,
        client: CogniteClient,
        builder: DataClassQueryBuilder[T_DomainModelList],
        filter_: dm.filters.Filter | None = None,
        limit: int = DEFAULT_QUERY_LIMIT,
    ):
        super().__init__(client, builder)
        from_ = self._builder.get_from()
        self._builder.append(
            NodeQueryStep(
                name=self._builder.create_name(from_),
                expression=dm.query.NodeResultSetExpression(
                    from_=from_,
                    filter=filter_,
                ),
                result_cls={{ data_class.read_name }},
                max_retrieve_limit=limit,
            )
        )
{% for edge_api in edge_apis %}
    def {{ edge_api.field_name }}(
        self,
        {% for parm in edge_api.filter_parameters(False, include_end_node=True, suffix_edge=True) %}
        {{ parm.name }}: {{ parm.annotation }} = {{ parm.default }},
        {% endfor %}
        filter: dm.Filter | None = None,
        limit: int = DEFAULT_QUERY_LIMIT,
        {% for field in data_class.one_to_one_direct_relations_with_source %}
        retrieve_{{ field.name }}: bool = False,
        {% endfor %}
    ) -> {{ edge_api.query.name }}[T_DomainModelList]:
        """Query along the {{ edge_api.doc_name }} edges of the {{ data_class.doc_name}}.

        Args:
            {% for parm in edge_api.filter_parameters(False, include_end_node=True, suffix_edge=True) %}
            {{ parm.name }}: {{ parm.description }}
            {% endfor %}
            filter: (Advanced) Filter applied to node. If the filtering available in the above is not sufficient, you can write your own filtering which will be ANDed with the filter above.
            limit: Maximum number of {{ edge_api.doc_name }} edges to return. Defaults to 3. Set to -1, float("inf") or None
                to return all items.
            {% for field in data_class.one_to_one_direct_relations_with_source %}
            retrieve_{{ field.name }}: Whether to retrieve the {{ field.doc_name }} for each {{ data_class.doc_name }} or not.
            {% endfor %}

        Returns:
            {{ edge_api.query.name }}: The query API for the {{ edge_api.end_class.doc_name }}.
        """
        from .{{ edge_api.query.file_name }} import {{ edge_api.query.name}}

        # from is a string as we added a node query step in the __init__ method
        from_ = cast(str, self._builder.get_from())
        {% if edge_api.has_edge_class %}
        edge_view = {{ edge_api.edge_class.read_name }}._view_id
        {% endif %}
        edge_filter = {% if edge_api.has_edge_class %}{{ edge_api.edge_class.filter_name }}{% else %}_create_edge_filter{% endif %}(
            dm.DirectRelationReference("{{ edge_api.type.space }}", "{{ edge_api.type.external_id }}"),
            {% if edge_api.has_edge_class %}
            edge_view,
            {% endif %}
            {% for parm in edge_api.filter_parameters(False, suffix_edge=True) %}
            {{ parm.name.removesuffix('_edge') }}={{ parm.name }},
            {% endfor %}
        )
        self._builder.append(
            EdgeQueryStep(
                name=self._builder.create_name(from_),
                expression=dm.query.EdgeResultSetExpression(
                    filter=edge_filter,
                    from_=from_,
                    direction="{{ edge_api.direction }}",
                ),
                {% if edge_api.has_edge_class %}
                result_cls= {{ edge_api.edge_class.read_name }},
                {% endif %}
                max_retrieve_limit=limit,
            )
        )

        view_id = {{ edge_api.query.name}}._view_id
        has_data = dm.filters.HasData(views=[view_id])
        node_filer = {{ edge_api.end_class.filter_name }}(
            view_id,
            {% for parm in edge_api.end_filter_method.parameters %}
            {{ parm.name }},
            {% endfor %}
            (filter and dm.filters.And(filter, has_data)) or has_data,
        )
        {% for field in data_class.one_to_one_direct_relations_with_source %}
        if retrieve_{{ field.name }}:
            self._query_append_{{ field.name }}(from_)
        {% endfor %}
        return {{ edge_api.query.name }}(self._client, self._builder, node_filer, limit)
        {% endfor %}

    def query(
        self,
        {% for field in data_class.one_to_one_direct_relations_with_source %}
        retrieve_{{ field.name }}: bool = False,
        {% endfor %}
    ) -> T_DomainModelList:
        """Execute query and return the result.

        {% if data_class.has_one_to_one_direct_relations_with_source %}
        Args:
            {% for field in data_class.one_to_one_direct_relations_with_source %}
            retrieve_{{ field.name }}: Whether to retrieve the {{ field.doc_name }} for each {{ data_class.doc_name }} or not.
            {% endfor %}

        {% endif %}
        Returns:
            The list of the source nodes of the query.

        """
        {% if data_class.has_one_to_one_direct_relations_with_source %}
        from_ = self._builder[-1].name
        {% endif %}
        {% for field in data_class.one_to_one_direct_relations_with_source %}
        if retrieve_{{ field.name }}:
            self._query_append_{{ field.name }}(from_)
        {% endfor %}
        return self._query()

    {% for field in data_class.one_to_one_direct_relations_with_source %}
    def _query_append_{{ field.name }}(self, from_: str) -> None:
        self._builder.append(
            NodeQueryStep(
                name=self._builder.create_name(from_),
                expression=dm.query.NodeResultSetExpression(
                    from_=from_,
                    through=self._view_id.as_property_ref("{{ field.prop_name }}"),
                    direction="outwards",
                    {% if field.destination_class.has_container_fields %}
                    filter=dm.filters.HasData(views=[{{ field.destination_class.read_name }}._view_id]),
                    {% endif %}
                ),
                result_cls={{ field.destination_class.read_name }},
            ),
        )
    {% endfor %}