from __future__ import annotations

{% if data_class.has_primitive_field_of_type((dm.Timestamp, dm.Date)) %}
import datetime
{% endif %}
import warnings
from collections.abc import Sequence
from typing import Any, ClassVar, Literal, Optional, TYPE_CHECKING, Union

from cognite.client import data_modeling as dm, CogniteClient
{% if data_class.use_pydantic_field %}
{{ data_class.import_pydantic_field }}
{% endif %}

from {{ top_level_package }}.data_classes._core import (
    {% if has_default_instance_space %}
    DEFAULT_INSTANCE_SPACE,
    {% endif %}
    DataRecord,
    DataRecordWrite,
    DomainModel,
    DomainModelCore,
    DomainModelWrite,
    DomainRelation,
    DomainRelationWrite,
    DomainRelationList,
    DomainRelationWriteList,
    GraphQLCore,
    ResourcesWrite,
    DomainModelList,
    T_DomainList,
    as_direct_relation_reference,
    as_instance_dict_id,
    as_node_id,
    as_read_args,
    as_write_args,
    as_pygen_node_id,
    is_tuple_id,
    EdgeQueryCore,
    NodeQueryCore,
    QueryCore,
    StringFilter,
    ViewPropertyId,
    {% if data_class.has_filtering_fields %}
    {{ data_class.filtering_import }}
    {% endif %}
)
{% if data_class.implements %}
{% for implements in data_class.implements %}
from {{ top_level_package }}.data_classes.{{ implements.file_name }} import {{ implements.read_name }}{% if data_class.is_writable or data_class.is_interface %}, {{ implements.write_name }}{% endif +%}
{% endfor %}
{% endif %}
{% for class_ in unique_start_classes %}
from {{ top_level_package }}.data_classes.{{class_.file_name }} import {{ class_.write_name }}
{% endfor %}
{% for classes in unique_end_classes %}
from {{ top_level_package }}.data_classes.{{classes.file_name }} import {{ classes.read_name }}, {{ classes.graphql_name }}, {{ classes.write_name }}
{% endfor %}
{% if data_class.has_dependencies_not_self %}
if TYPE_CHECKING:
{% for dependency_class in data_class.dependencies %}
{% if dependency_class.file_name != data_class.file_name %}
    from {{ top_level_package }}.data_classes.{{ dependency_class.file_name }} import {{ dependency_class.read_name }}, {{ dependency_class.graphql_name }}{% if dependency_class.is_writable or dependency_class.is_interface %}, {{ dependency_class.write_name }}
{% endif %}
{% endif %}
{% endfor %}
{% endif %}


__all__ = [
    "{{ data_class.read_name }}",{% if data_class.is_writable or data_class.is_interface %}
    "{{ data_class.write_name }}",
    "{{ data_class.read_name }}Apply",{% endif %}
    "{{ data_class.read_list_name }}",{% if data_class.is_writable or data_class.is_interface %}
    "{{ data_class.write_list_name }}",
    "{{ data_class.read_name }}ApplyList",{% endif %}{% if data_class.has_field_of_type(ft.BasePrimitiveField) %}
    "{{ data_class.field_names }}",{% endif %}{% if data_class.has_primitive_field_of_type((dm.Text, dm.CDFExternalIdReference)) %}
    "{{ data_class.text_field_names }}",{% endif %}
]


{{ data_class.text_field_names }} = Literal["external_id", {{ data_class.text_fields_literals }}]
{{ data_class.field_names }} = Literal["external_id", {{ data_class.fields_literals }}]
{{ data_class.properties_dict_name }} = {% raw %}{{% endraw %}
    "external_id": "externalId",
    {% for field in data_class.fields_of_type(ft.BasePrimitiveField) %}
    "{{ field.name }}": "{{ field.prop_name }}",
    {% endfor %}
{% raw %}}{% endraw %}


class {{ data_class.graphql_name }}(GraphQLCore{% if data_class.has_any_field_model_prefix %}, protected_namespaces=(){% endif %}):
    """This represents the reading version of {{ data_class.doc_name }}, used
    when data is retrieved from CDF using GraphQL.

    It is used when retrieving data from CDF using GraphQL.

    Args:
        space: The space where the node is located.
        external_id: The external id of the {{ data_class.doc_name }}.
        data_record: The data record of the {{ data_class.doc_name }} node.
        {% for field in data_class %}
        {{ field.name }}: {{ field.argument_documentation }}
        {% endfor %}
    """
    view_id: ClassVar[dm.ViewId] = dm.ViewId("{{ data_class.view_id.space }}", "{{ data_class.view_id.external_id }}", "{{ data_class.view_id.version }}")
    {% for field in data_class %}
    {{ field.name }}: {{ field.as_graphql_type_hint() }}
    {% endfor %}

    def as_read(self) -> {{ data_class.read_name }}:
        """Convert this GraphQL format of {{ data_class.doc_name }} to the reading format."""
        return {{ data_class.read_name }}.model_validate(as_read_args(self))

{% if data_class.is_writable or data_class.is_interface %}
    def as_write(self) -> {{ data_class.write_name }}:
        """Convert this GraphQL format of {{ data_class.doc_name }} to the writing format."""
        return {{ data_class.write_name }}.model_validate(as_write_args(self))
{% endif %}


class {{ data_class.read_name }}({{ data_class.read_base_class }}{% if data_class.has_any_field_model_prefix %}, protected_namespaces=(){% endif %}):
    """This represents the reading version of {{ data_class.doc_name }}.

    It is used to when data is retrieved from CDF.

    Args:
        space: The space where the node is located.
        external_id: The external id of the {{ data_class.doc_name }}.
        data_record: The data record of the {{ data_class.doc_name }} edge.
        {% for field in data_class %}
        {{ field.name }}: {{ field.argument_documentation }}
        {% endfor %}
    """
    _view_id: ClassVar[dm.ViewId] = dm.ViewId("{{ data_class.view_id.space }}", "{{ data_class.view_id.external_id }}", "{{ data_class.view_id.version }}")
    space: str{% if has_default_instance_space %} = DEFAULT_INSTANCE_SPACE{% endif +%}
    {% for field in data_class.read_fields %}
    {{ field.name }}: {{ field.as_read_type_hint() }}
    {% endfor %}


{% if data_class.is_writable or data_class.is_interface %}
    def as_write(self) -> {{ data_class.write_name }}:
        """Convert this read version of {{ data_class.doc_name }} to the writing version."""
        return {{ data_class.write_name }}.model_validate(as_write_args(self))

    def as_apply(self) -> {{ data_class.write_name }}:
        """Convert this read version of {{ data_class.doc_name }} to the writing version."""
        warnings.warn(
            "as_apply is deprecated and will be removed in v1.0. Use as_write instead.",
            UserWarning,
            stacklevel=2,
        )
        return self.as_write()


_EXPECTED_START_NODES_BY_END_NODE: dict[type[DomainModelWrite], set[type[DomainModelWrite]]] = {% raw %}{{% endraw %}{% for end_class, classes in grouped_edge_classes.items() %}
    {{ end_class }}: {% raw %}{{% endraw %}{{ ", ".join(classes) }}{% raw %}}{% endraw %},{% endfor %}
{% raw %}}{% endraw %}


def _validate_end_node(
    start_node: DomainModelWrite, end_node: {{ data_class.end_node_field.as_validate_function_type_hint() }}
) -> None:
    if isinstance(end_node, str | dm.NodeId):
        # Nothing to validate
        return
    if type(end_node) not in _EXPECTED_START_NODES_BY_END_NODE:
        raise ValueError(
            f"Invalid end node type: {type(end_node)}. "
            f"Should be one of {[t.__name__ for t in _EXPECTED_START_NODES_BY_END_NODE.keys()]}"
        )
    if type(start_node) not in _EXPECTED_START_NODES_BY_END_NODE[type(end_node)]:
        raise ValueError(
            f"Invalid end node type: {type(end_node)}. "
            f"Expected one of: {_EXPECTED_START_NODES_BY_END_NODE[type(end_node)]}"
        )


class {{ data_class.write_name }}({{ data_class.write_base_class }}{% if data_class.has_any_field_model_prefix %}, protected_namespaces=(){% endif %}):
    """This represents the writing version of {{ data_class.doc_name }}.

    It is used to when data is sent to CDF.

    Args:
        space: The space where the node is located.
        external_id: The external id of the {{ data_class.doc_name }}.
        data_record: The data record of the {{ data_class.doc_name }} edge.
        {% for field in data_class %}
        {% if field.is_write_field %}
        {{ field.name }}: {{ field.argument_documentation }}
        {% endif %}
        {% endfor %}
    """
    {% if data_class.has_container_fields_writable %}
    _container_fields: ClassVar[tuple[str, ...]] = ({{ data_class.container_fields_writable_comma_sep }},)
    {% endif %}
    {% if data_class.has_outwards_edges %}
    _outwards_edges: ClassVar[tuple[tuple[str, dm.DirectRelationReference], ...]] = ({{ data_class.outwards_edges_comma_sep }},)
    {% endif %}
    {% if data_class.has_inwards_edges %}
    _inwards_edges: ClassVar[tuple[tuple[str, dm.DirectRelationReference], ...]] = ({{ data_class.inwards_edges_comma_sep }},)
    {% endif %}
    {% if data_class.has_direct_relations_with_source %}
    _direct_relations: ClassVar[tuple[str, ...]] = ({{ data_class.direct_relations_with_source_comma_sep }},)
    {% endif %}
    _validate_end_node = _validate_end_node

    _view_id: ClassVar[dm.ViewId] = dm.ViewId("{{ data_class.view_id.space }}", "{{ data_class.view_id.external_id }}", "{{ data_class.view_id.version }}")
    {% for field in data_class.write_fields %}
    {{ field.name }}: {{ field.as_write_type_hint() }}
    {% endfor %}


class {{ data_class.read_name }}Apply({{ data_class.write_name }}):
    def __new__(cls, *args, **kwargs) -> {{ data_class.read_name }}Apply:
        warnings.warn(
            "{{ data_class.read_name }}Apply is deprecated and will be removed in v1.0. "
            "Use {{ data_class.write_name }} instead."
            "The motivation for this change is that Write is a more descriptive name for the writing version of the"
            "{{ data_class.read_name }}.",
            UserWarning,
            stacklevel=2,
        )
        return super().__new__(cls)

{% endif %}
class {{ data_class.read_list_name }}(DomainRelationList[{{ data_class.read_name }}]):
    """List of {{ data_class.doc_list_name }} in the reading version."""

    _INSTANCE = {{ data_class.read_name }}
{% if data_class.is_writable or data_class.is_interface %}
    def as_write(self) -> {{ data_class.write_list_name }}:
        """Convert this read version of {{ data_class.doc_name }} list to the writing version."""
        return {{ data_class.write_list_name }}([edge.as_write() for edge in self])

    def as_apply(self) -> {{ data_class.write_list_name }}:
        """Convert these read versions of {{ data_class.doc_name }} list to the writing versions."""
        warnings.warn(
            "as_apply is deprecated and will be removed in v1.0. Use as_write instead.",
            UserWarning,
            stacklevel=2,
        )
        return self.as_write()


class {{ data_class.write_list_name }}(DomainRelationWriteList[{{ data_class.write_name }}]):
    """List of {{ data_class.doc_list_name }} in the writing version."""

    _INSTANCE = {{ data_class.write_name }}


class {{ data_class.read_name }}ApplyList({{ data_class.write_list_name }}): ...

{% endif %}
def {{ data_class.filter_name }}(
    edge_type: dm.DirectRelationReference,
    view_id: dm.ViewId,
    {% if has_default_instance_space %}
    start_node: str | list[str] | dm.NodeId | list[dm.NodeId] | None = None,
    start_node_space: str = DEFAULT_INSTANCE_SPACE,
    end_node: str | list[str] | dm.NodeId | list[dm.NodeId] | None = None,
    space_end_node: str = DEFAULT_INSTANCE_SPACE,
    {% else %}
    start_node: dm.NodeId | list[dm.NodeId] | None = None,
    end_node: dm.NodeId | list[dm.NodeId] | None = None,
    {% endif %}
    {% for parm in list_method.parameters %}
    {{ parm.name }}: {{ parm.annotation }} = {{ parm.default }},
    {% endfor %}
    filter: dm.Filter | None = None,
) -> dm.Filter:
    filters: list[dm.Filter] = [
        dm.filters.Equals(
            ["edge", "type"],
            {"space": edge_type.space, "externalId": edge_type.external_id},
        )
    ]
    {% if has_default_instance_space %}
    if start_node and isinstance(start_node, str):
        filters.append(
            dm.filters.Equals(["edge", "startNode"], value={"space": start_node_space, "externalId": start_node})
        )
    {% endif %}
    if start_node and isinstance(start_node, dm.NodeId):
        filters.append(
            dm.filters.Equals(
                ["edge", "startNode"], value=start_node.dump(camel_case=True, include_instance_type=False)
            )
        )
    if start_node and isinstance(start_node, list):
        filters.append(
            dm.filters.In(
                ["edge", "startNode"],
                values=[{% if has_default_instance_space %}
                    {"space": start_node_space, "externalId": ext_id}
                    if isinstance(ext_id, str)
                    else {% endif %}ext_id.dump(camel_case=True, include_instance_type=False)
                    for ext_id in start_node
                ],
            )
        )
    {% if has_default_instance_space %}
    if end_node and isinstance(end_node, str):
        filters.append(dm.filters.Equals(["edge", "endNode"], value={"space": space_end_node, "externalId": end_node}))
    {% endif %}
    if end_node and isinstance(end_node, dm.NodeId):
        filters.append(
            dm.filters.Equals(["edge", "endNode"], value=end_node.dump(camel_case=True, include_instance_type=False))
        )
    if end_node and isinstance(end_node, list):
        filters.append(
            dm.filters.In(
                ["edge", "endNode"],
                values=[{% if has_default_instance_space %}
                    {"space": space_end_node, "externalId": ext_id}
                    if isinstance(ext_id, str)
                    else {% endif %}ext_id.dump(camel_case=True, include_instance_type=False)
                    for ext_id in end_node
                ],
            )
        )
    {% for filter in list_method.implementations %}
    if {{ filter.condition }}:
        filters.append({{ filter.filter_call }}({{ filter.arguments }}))
    {% endfor %}
    if filter:
        filters.append(filter)
    return dm.filters.And(*filters)


class _{{ data_class.query_cls_name }}(EdgeQueryCore[T_DomainList, {{ data_class.read_list_name }}]):
    _view_id = {{ data_class.read_name }}._view_id
    _result_cls = {{ data_class.read_name }}
    _result_list_cls_end = {{ data_class.read_list_name }}

    def __init__(
        self,
        created_types: set[type],
        creation_path: list[QueryCore],
        client: CogniteClient,
        result_list_cls: type[T_DomainList],
        end_node_cls: type[NodeQueryCore],
        expression: dm.query.ResultSetExpression | None = None,
        connection_name: str | None = None,
        connection_property: ViewPropertyId | None = None,
    ):
        {% for dependency_class in data_class.dependencies %}
        {% if dependency_class.file_name != data_class.file_name %}
        from .{{ dependency_class.file_name }} import _{{ dependency_class.query_cls_name }}
        {% endif %}
        {% endfor %}

        super().__init__(created_types, creation_path, client, result_list_cls, expression, None, connection_name, connection_property)
        if end_node_cls not in created_types:
            self.end_node = end_node_cls(
                created_types=created_types.copy(),
                creation_path=self._creation_path,
                client=client,
                result_list_cls=result_list_cls,  # type: ignore[type-var]
                expression=dm.query.NodeResultSetExpression(),
                connection_property=ViewPropertyId(self._view_id, "end_node")
            )

        {% for field in data_class.fields_of_type(ft.BaseConnectionField) %}
        {% if not field.is_no_source_direct_relation %}
        if _{{ field.destination_class.query_cls_name }} not in created_types:
            self.{{ field.name }} = _{{ field.destination_class.query_cls_name }}(
                created_types.copy(),
                self._creation_path,
                client,
                result_list_cls,  # type: ignore[type-var]
                {% if field.is_direct_relation %}
                dm.query.NodeResultSetExpression(
                    through=self._view_id.as_property_ref("{{ field.prop_name }}"),
                    direction="outwards",
                ),
                {% elif field.is_reverse_direct_relation %}
                dm.query.NodeResultSetExpression(
                    through={{ field.reverse_property.data_class.view_id_str }}.as_property_ref("{{ field.prop_name }}"),
                    direction="inwards",
                ),
                {% else %}
                dm.query.EdgeResultSetExpression(
                    direction="{{ field.edge_direction }}",
                    chain_to="destination",
                ),
                {% endif %}
                connection_name="{{ field.name }}",
                connection_property=ViewPropertyId(self._view_id, "{{ field.prop_name }}"),
            )
        {% endif %}

        {% endfor %}
        self.space = StringFilter(self, ["node", "space"])
        self.external_id = StringFilter(self, ["node", "externalId"])
        {% for field in data_class.filtering_fields %}
        self.{{ field.filter_attribute }} = {{ field.filtering_cls }}(self, self._view_id.as_property_ref("{{ field.prop_name }}"))
        {% endfor %}
        self._filter_classes.extend([
            self.space,
            self.external_id,
            {% for field in data_class.filtering_fields %}
            self.{{ field.filter_attribute }},
            {% endfor %}
        ])
