from __future__ import annotations
{% if data_class.has_time_field or data_class.has_single_timeseries_fields %}
import datetime{% endif %}{% if data_class.has_single_timeseries_fields %}
import warnings{% endif %}
{% if data_class.has_one_to_many_edges %}from collections import defaultdict
from typing import Dict, List, Sequence, Tuple, overload{% else %}from typing import Sequence, overload{% endif %}{% if data_class.has_single_timeseries_fields %}, Literal{% endif %}
{% if data_class.has_single_timeseries_fields %}
import pandas as pd{% endif %}
from cognite.client import CogniteClient{% if data_class.has_single_timeseries_fields %}
from cognite.client.data_classes import TimeSeriesList, DatapointsList, Datapoints, DatapointsArrayList
from cognite.client.data_classes.datapoints import Aggregate{% endif %}
from cognite.client import data_modeling as dm

from ._core import DEFAULT_LIMIT_READ, TypeAPI{% if data_class.has_single_timeseries_fields %}, INSTANCE_QUERY_LIMIT{% endif %}
from {{ top_level_package }}.data_classes import {{ data_class.read_name }}, {{ data_class.write_name }}, {{ data_class.read_list_name }}, {{ data_class.write_list_name }}
{% if data_class.has_single_timeseries_fields %}{% include "api_class_timeseries.py.jinja" %}{% endif %}
{% for field in data_class.one_to_many_edges %}
class {{ field.edge_api_class }}:
    def __init__(self, client: CogniteClient):
        self._client = client

    def retrieve(self, external_id: str | Sequence[str]) -> dm.EdgeList:
        f = dm.filters
        is_edge_type = f.Equals(
            ["edge", "type"],
            {"space": "{{ field.prop.type.space }}", "externalId": "{{ field.prop.type.external_id }}"},
        )
        if isinstance(external_id, str):
            is_{{ data_class.variable }} = f.Equals(
                ["edge", "startNode"],
                {"space": "{{ field.prop.source.space }}", "externalId": external_id},
            )
            return self._client.data_modeling.instances.list("edge", limit=-1, filter=f.And(is_edge_type, is_{{ data_class.variable }}))

        else:
            is_{{ data_class.variable_list }} = f.In(
                ["edge", "startNode"],
                [{"space": "{{ field.prop.source.space }}", "externalId": ext_id} for ext_id in external_id],
            )
            return self._client.data_modeling.instances.list("edge", limit=-1, filter=f.And(is_edge_type, is_{{ data_class.variable_list }}))

    def list(self, {{ data_class.variable }}_id: str | list[str] | None = None, limit=DEFAULT_LIMIT_READ) -> dm.EdgeList:
        f = dm.filters
        filters = []
        is_edge_type = f.Equals(
            ["edge", "type"],
            {"space": "{{ field.prop.type.space }}", "externalId": "{{ field.prop.type.external_id }}"},
        )
        filters.append(is_edge_type)
        if {{ data_class.variable }}_id:
            {{ data_class.variable }}_ids = [{{ data_class.variable }}_id] if isinstance({{ data_class.variable }}_id, str) else {{ data_class.variable }}_id
            is_{{ data_class.variable_list }} = f.In(
                ["edge", "startNode"],
                [{"space": "{{ field.prop.source.space }}", "externalId": ext_id} for ext_id in {{ data_class.variable }}_ids],
            )
            filters.append(is_{{ data_class.variable_list }})

        return self._client.data_modeling.instances.list("edge", limit=limit, filter=f.And(*filters))

{% endfor %}
class {{ api_class.name }}(TypeAPI[{{ data_class.read_name }}, {{ data_class.write_name }}, {{ data_class.read_list_name }}]):
    def __init__(self, client: CogniteClient, view_id: dm.ViewId):
        super().__init__(
            client=client,
            sources=view_id,
            class_type={{ data_class.read_name }},
            class_apply_type={{ data_class.write_name }},
            class_list={{ data_class.read_list_name }},
        )
        self.view_id = view_id{% for field in data_class.one_to_many_edges %}
        self.{{ field.edge_api_attribute }} = {{ field.edge_api_class }}(client){% endfor %}{% for timeseries in data_class.single_timeseries_fields %}
        self.{{ timeseries.edge_api_attribute }} = {{ timeseries.edge_api_class }}API(client, view_id){% endfor %}

    def apply(self, {{ data_class.variable }}: {{ data_class.write_name }} | Sequence[{{ data_class.write_name }}], replace: bool = False) -> dm.InstancesApplyResult:
        if isinstance({{ data_class.variable }}, {{ data_class.write_name }}):
            instances = {{ data_class.variable }}.to_instances_apply()
        else:
            instances = {{ data_class.write_list_name }}({{ data_class.variable }}).to_instances_apply()
        {% if data_class.has_only_one_to_many_edges %}return self._client.data_modeling.instances.apply(
            nodes=instances.nodes,
            edges=instances.edges,
            auto_create_start_nodes=True,
            auto_create_end_nodes=True,
            replace=replace,
        ){% else %}return self._client.data_modeling.instances.apply(nodes=instances.nodes, edges=instances.edges, replace=replace){% endif %}

    def delete(self, external_id: str | Sequence[str]) -> dm.InstancesDeleteResult:
        if isinstance(external_id, str):
            return self._client.data_modeling.instances.delete(nodes=({{ data_class.write_name }}.space, external_id))
        else:
            return self._client.data_modeling.instances.delete(
                nodes=[({{ data_class.write_name }}.space, id) for id in external_id],
            )

    @overload
    def retrieve(self, external_id: str) -> {{ data_class.read_name }}:
        ...

    @overload
    def retrieve(self, external_id: Sequence[str]) -> {{ data_class.read_list_name }}:
        ...

    def retrieve(self, external_id: str | Sequence[str]) -> {{ data_class.read_name }} | {{ data_class.read_list_name }}:
        if isinstance(external_id, str):
            {% if data_class.has_one_to_many_edges %}{{ data_class.variable }} = self._retrieve((self.sources.space, external_id))
{% for field in data_class.one_to_many_edges %}
            {{ field.variable }}_edges = self.{{ field.edge_api_attribute }}.retrieve(external_id)
            {{ data_class.variable }}.{{ field.name }} = [edge.end_node.external_id for edge in {{ field.variable }}_edges]{% endfor %}

            return {{ data_class.variable }}{% else %}return self._retrieve((self.sources.space, external_id)){% endif %}
        else:
            {% if data_class.has_one_to_many_edges %}{{ data_class.variable_list }} = self._retrieve([(self.sources.space, ext_id) for ext_id in external_id])
{% for field in data_class.one_to_many_edges %}
            {{ field.variable }}_edges = self.{{ field.edge_api_attribute }}.retrieve(external_id)
            self._set_{{ field.name }}({{ data_class.variable_list }}, {{ field.variable }}_edges){% endfor %}

            return {{ data_class.variable_list }}{% else %}return self._retrieve([(self.sources.space, ext_id) for ext_id in external_id]){% endif %}

    def list(
        self,{% for parm in list_method.parameters %}
        {{ parm.name }}: {{ parm.annotation }} = {{ parm.default }},{% endfor %}
        limit: int = DEFAULT_LIMIT_READ,
        filter: dm.Filter | None = None,{% if data_class.has_one_to_many_edges %}
        retrieve_edges: bool = True,{% endif %}
    ) -> {{ data_class.read_list_name }}:
        filter_ = _create_filter(
            self.view_id,{% for parm in list_method.parameters %}
            {{ parm.name }},{% endfor %}
            filter,
        )

        {% if data_class.has_one_to_many_edges %}{{ data_class.variable_list }} = self._list(limit=limit, filter=filter_)

        if retrieve_edges:{% for field in data_class.one_to_many_edges %}
            {{ field.variable }}_edges = self.{{ field.edge_api_attribute }}.list({{ data_class.variable_list }}.as_external_ids(), limit=-1)
            self._set_{{ field.name }}({{ data_class.variable_list }}, {{ field.variable }}_edges){% endfor %}

        return {{ data_class.variable_list }}{% else %}return self._list(limit=limit, filter=filter_){% endif %}
{% for field in data_class.one_to_many_edges %}
    @staticmethod
    def _set_{{ field.name }}({{ data_class.variable_list }}: Sequence[{{ data_class.read_name }}], {{ field.variable }}_edges: Sequence[dm.Edge]):
        edges_by_start_node: Dict[Tuple, List] = defaultdict(list)
        for edge in {{ field.variable }}_edges:
            edges_by_start_node[edge.start_node.as_tuple()].append(edge)

        for {{ data_class.variable }} in {{ data_class.variable_list }}:
            node_id = {{ data_class.variable }}.id_tuple()
            if node_id in edges_by_start_node:
                {{ data_class.variable }}.{{ field.name }} = [edge.end_node.external_id for edge in edges_by_start_node[node_id]]
{% endfor %}

def _create_filter(
    view_id: dm.ViewId,{% for parm in list_method.parameters %}
    {{ parm.name }}: {{ parm.annotation }} = {{ parm.default }},{% endfor %}
    filter: dm.Filter | None = None,
) -> dm.Filter | None:
    filters = []{% for filter in list_method.filters %}
    if {{ filter.condition }}:
        filters.append({{ filter.filter_call }}({{ filter.arguments }})){% endfor %}
    if filter:
        filters.append(filter)
    return dm.filters.And(*filters) if filters else None
