from __future__ import annotations

import datetime
from typing import TYPE_CHECKING

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.fields_of_type(ft.EdgeOneToOne) %}
    {{ field.data_class.read_name }},{% endfor %}
)
from ._core import DEFAULT_QUERY_LIMIT, QueryBuilder, QueryStep, 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 edge_apis %}
if TYPE_CHECKING:{% for edge_api in edge_apis %}
    from .{{ edge_api.query.file_name}} import {{ edge_api.query.name }}{% endfor %}
{% endif %}


class {{ query_api.name }}(QueryAPI[T_DomainModelList]):
    def __init__(
        self,
        client: CogniteClient,
        builder: QueryBuilder[T_DomainModelList],
        view_by_read_class: dict[type[DomainModelCore], dm.ViewId],
        filter_: dm.filters.Filter | None = None,
        limit: int = DEFAULT_QUERY_LIMIT,
    ):
        super().__init__(client, builder, view_by_read_class)

        self._builder.append(
            QueryStep(
                name=self._builder.next_name("{{ data_class.variable }}"),
                expression=dm.query.NodeResultSetExpression(
                    from_=self._builder[-1].name if self._builder else None,
                    filter=filter_,
                ),
                select=dm.query.Select([dm.query.SourceSelector(self._view_by_read_class[{{ data_class.read_name }}], ["*"])]),
                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) %}
        {{ parm.name }}: {{ parm.annotation }} = {{ parm.default }},{% endfor %}
        limit: int | None = DEFAULT_QUERY_LIMIT,{% for field in data_class.fields_of_type(ft.EdgeOneToOne) %}
        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) %}
            {{ parm.name }}: {{ parm.description }}{% endfor %}
            limit: Maximum number of {{ edge_api.doc_name }} edges to return. Defaults to 25. Set to -1, float("inf") or None
                to return all items.{% for field in data_class.fields_of_type(ft.EdgeOneToOne) %}
            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_ = self._builder[-1].name
{% if edge_api.has_edge_class %}
        edge_view = self._view_by_read_class[{{ edge_api.edge_class.read_name }}]{% 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) %}
            {{ parm.name }}={{ parm.name }},{% endfor %}
        )
        self._builder.append(
            QueryStep(
                name=self._builder.next_name("{{ edge_api.field_name }}"),
                expression=dm.query.EdgeResultSetExpression(
                    filter=edge_filter,
                    from_=from_,
                    direction="{{ edge_api.direction }}",
                ),{% if edge_api.has_edge_class %}
                select=dm.query.Select(
                    [dm.query.SourceSelector(edge_view, ["*"])],
                ),
                result_cls= {{ edge_api.edge_class.read_name }},{% else %}
                select=dm.query.Select(),{% endif %}
                max_retrieve_limit=limit,
            )
        ){% for field in data_class.fields_of_type(ft.EdgeOneToOne) %}
        if retrieve_{{ field.name }}:
            self._query_append_{{ field.name }}(from_){% endfor %}
        return {{ edge_api.query.name }}(self._client, self._builder, self._view_by_read_class, None, limit)
    {% endfor %}
    def query(
        self,{% for field in data_class.fields_of_type(ft.EdgeOneToOne) %}
        retrieve_{{ field.name }}: bool = False,{% endfor %}
    ) -> T_DomainModelList:
        """Execute query and return the result.
{% if data_class.has_field_of_type(ft.EdgeOneToOne) %}
        Args:{% for field in data_class.fields_of_type(ft.EdgeOneToOne) %}
            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_field_of_type(ft.EdgeOneToOne) %}
        from_ = self._builder[-1].name{% endif %}{% for field in data_class.fields_of_type(ft.EdgeOneToOne) %}
        if retrieve_{{ field.name }}:
            self._query_append_{{ field.name }}(from_){% endfor %}
        return self._query()
{% for field in data_class.fields_of_type(ft.EdgeOneToOne) %}
    def _query_append_{{ field.name }}(self, from_: str) -> None:
        view_id = self._view_by_read_class[{{ field.data_class.read_name }}]
        self._builder.append(
            QueryStep(
                name=self._builder.next_name("{{ field.name }}"),
                expression=dm.query.NodeResultSetExpression(
                    filter=dm.filters.HasData(views=[view_id]),
                    from_=from_,
                    through=self._view_by_read_class[{{ data_class.read_name }}].as_property_ref("{{ field.prop_name }}"),
                    direction="outwards",
                ),
                select=dm.query.Select([dm.query.SourceSelector(view_id, ["*"])]),
                max_retrieve_limit=-1,
                result_cls={{ field.data_class.read_name }},
            ),
        )
{% endfor %}
