from __future__ import annotations

import itertools
from abc import ABC, abstractmethod
from collections import defaultdict
from collections.abc import Iterable, Iterator
from dataclasses import dataclass, field
from typing import cast

from cognite.client.data_classes import data_modeling as dm
from cognite.client.data_classes.data_modeling.data_types import ListablePropertyType
from typing_extensions import Self

from cognite.pygen import config as pygen_config
from cognite.pygen.utils.text import create_name

_PRIMITIVE_TYPES = (dm.Text, dm.Boolean, dm.Float32, dm.Float64, dm.Int32, dm.Int64, dm.Timestamp, dm.Date, dm.Json)
_EXTERNAL_TYPES = (dm.TimeSeriesReference, dm.FileReference, dm.SequenceReference)


@dataclass(frozen=True)
class ViewSpaceExternalId:
    """
    This represents a view id.

    The motivation for this class is that API class and data classes are independent of the view version (given
    that all properties are equal). This enables reuse of data classes and APIs across view versions.

    Note this could be considered an edge case not worth supporting. However, when you publish data models
    from a .graphql format, the autogenerated view versions for identical types will be different when those types
    are part of two different data models.
    """

    space: str
    external_id: str

    @classmethod
    def from_(cls, view_id: dm.ViewId | dm.View) -> Self:
        return cls(space=view_id.space, external_id=view_id.external_id)


@dataclass(frozen=True)
class Field(ABC):
    """
    A field represents a pydantic field in the generated pydantic class.

    Args:
        name: The name of the field. This is used in the generated Python code.
        prop_name: The name of the property in the data model. This is used when reading and writing to CDF.
        pydantic_field: The name to use for the import 'from pydantic import Field'. This is used in the edge case
                        when the name 'Field' name clashes with the data model class name.

    """

    name: str
    prop_name: str
    pydantic_field: str

    @property
    def need_alias(self) -> bool:
        return self.name != self.prop_name

    @classmethod
    def from_property(
        cls,
        prop_name: str,
        prop: dm.MappedProperty | dm.ConnectionDefinition,
        data_class_by_view_id: dict[ViewSpaceExternalId, DataClass],
        field_naming: pygen_config.FieldNaming,
        view_name: str,
        pydantic_field: str = "Field",
    ) -> Field:
        name = create_name(prop_name, field_naming.name)
        if isinstance(prop, dm.SingleHopConnectionDefinition):
            variable = create_name(prop_name, field_naming.variable)

            edge_api_class_input = f"{view_name}_{prop_name}"
            edge_api_class = f"{create_name(edge_api_class_input, field_naming.edge_api_class)}API"
            edge_api_attribute = create_name(prop_name, field_naming.api_class_attribute)
            return EdgeOneToMany(
                name=name,
                prop_name=prop_name,
                prop=prop,
                data_class=data_class_by_view_id[ViewSpaceExternalId(prop.source.space, prop.source.external_id)],
                variable=variable,
                pydantic_field=pydantic_field,
                edge_api_class=edge_api_class,
                edge_api_attribute=edge_api_attribute,
            )
        elif isinstance(prop, dm.MappedProperty) and (
            isinstance(prop.type, _PRIMITIVE_TYPES) or isinstance(prop.type, _EXTERNAL_TYPES)
        ):
            type_ = _to_python_type(prop.type)
            if isinstance(prop.type, ListablePropertyType) and prop.type.is_list:
                return PrimitiveListField(
                    name=name,
                    prop_name=prop_name,
                    type_=type_,
                    is_nullable=prop.nullable,
                    prop=prop,
                    pydantic_field=pydantic_field,
                )
            elif isinstance(prop.type, dm.CDFExternalIdReference):
                # Note: these are only CDF External Fields that are not listable. Listable CDF External Fields
                # are handled above.
                edge_api_class_input = f"{view_name}_{prop_name}"
                edge_api_class = f"{create_name(edge_api_class_input, field_naming.edge_api_class)}"
                edge_api_attribute = create_name(prop_name, field_naming.api_class_attribute)
                return CDFExternalField(
                    name=name,
                    prop_name=prop_name,
                    type_=type_,
                    is_nullable=prop.nullable,
                    prop=prop,
                    pydantic_field=pydantic_field,
                    edge_api_class=edge_api_class,
                    edge_api_attribute=edge_api_attribute,
                )
            else:
                return PrimitiveField(
                    name=name,
                    prop_name=prop_name,
                    type_=type_,
                    is_nullable=prop.nullable,
                    default=prop.default_value,
                    prop=prop,
                    pydantic_field=pydantic_field,
                )
        elif isinstance(prop, dm.MappedProperty) and isinstance(prop.type, dm.DirectRelation):
            # For direct relation the source is required.
            view_id = cast(dm.ViewId, prop.source)
            target_data_class = data_class_by_view_id[ViewSpaceExternalId(view_id.space, view_id.external_id)]
            return EdgeOneToOne(
                name=name,
                prop_name=prop_name,
                prop=prop,
                data_class=target_data_class,
                pydantic_field=pydantic_field,
            )

        else:
            raise NotImplementedError(f"Property type={type(prop)!r} is not supported")

    @abstractmethod
    def as_read_type_hint(self) -> str:
        raise NotImplementedError()

    @abstractmethod
    def as_write_type_hint(self) -> str:
        raise NotImplementedError()

    @property
    @abstractmethod
    def is_edge(self) -> bool:
        raise NotImplementedError()

    @property
    @abstractmethod
    def is_time_field(self) -> bool:
        raise NotImplementedError()


@dataclass(frozen=True)
class PrimitiveFieldCore(Field, ABC):
    type_: str
    is_nullable: bool
    prop: dm.MappedProperty

    @property
    def is_edge(self) -> bool:
        return False

    @property
    def is_time_field(self) -> bool:
        return self.type_ in ("datetime.datetime", "datetime.date")


@dataclass(frozen=True)
class PrimitiveField(PrimitiveFieldCore):
    """
    This represents a basic type such as str, int, float, bool, datetime.datetime, datetime.date.
    """

    default: str | int | dict | None = None

    def as_read_type_hint(self) -> str:
        rhs = str(self.default)
        if self.need_alias:
            rhs = f'{self.pydantic_field}({rhs}, alias="{self.prop_name}")'

        return f"Optional[{self.type_}] = {rhs}"

    def as_write_type_hint(self) -> str:
        out_type = self.type_
        if self.is_nullable:
            out_type = f"Optional[{out_type}] = {self.default}"
        return out_type


@dataclass(frozen=True)
class PrimitiveListField(PrimitiveFieldCore):
    """
    This represents a list of basic types such as list[str], list[int], list[float], list[bool],
    list[datetime.datetime], list[datetime.date].
    """

    def as_read_type_hint(self) -> str:
        alias = ""
        if self.need_alias:
            alias = f', alias="{self.prop_name}"'

        if alias:
            rhs = f" = {self.pydantic_field}(default_factory=list{alias})"
        else:
            rhs = " = []"
        return f"list[{self.type_}]{rhs}"

    def as_write_type_hint(self) -> str:
        rhs = ""
        if self.is_nullable:
            rhs = " = []"
        return f"list[{self.type_}]{rhs}"


@dataclass(frozen=True)
class CDFExternalField(PrimitiveFieldCore):
    edge_api_class: str
    edge_api_attribute: str

    def as_read_type_hint(self) -> str:
        rhs = "None"
        if self.need_alias:
            rhs = f'{self.pydantic_field}({rhs}, alias="{self.prop_name}")'

        return f"Optional[{self.type_}] = {rhs}"

    def as_write_type_hint(self) -> str:
        out_type = self.type_
        if self.is_nullable:
            out_type = f"Optional[{out_type}] = None"
        return out_type


@dataclass(frozen=True)
class EdgeField(Field, ABC):
    """
    This represents an edge field linking to another data class.
    """

    data_class: DataClass

    @property
    def is_edge(self) -> bool:
        return True

    @property
    def is_time_field(self) -> bool:
        return False


@dataclass(frozen=True)
class EdgeOneToOne(EdgeField):
    """
    This represents an edge field linking to another data class.
    """

    prop: dm.MappedProperty

    def as_read_type_hint(self) -> str:
        return "Optional[str] = None"

    def as_write_type_hint(self) -> str:
        return f"Union[{self.data_class.write_name}, str, None] = {self.pydantic_field}(None, repr=False)"


@dataclass(frozen=True)
class EdgeOneToMany(EdgeField):
    """
    This represents a list of edge fields linking to another data class.
    """

    variable: str
    edge_api_class: str
    edge_api_attribute: str
    prop: dm.SingleHopConnectionDefinition

    def as_read_type_hint(self) -> str:
        return "list[str] = []"

    def as_write_type_hint(self) -> str:
        return (
            f"Union[list[{self.data_class.write_name}], list[str]] "
            f"= {self.pydantic_field}(default_factory=list, repr=False)"
        )


@dataclass(frozen=True)
class DataClass:
    """
    This represents a data class. It is created from a view.
    """

    view_name: str
    read_name: str
    write_name: str
    read_list_name: str
    write_list_name: str
    variable: str
    variable_list: str
    file_name: str
    view_id: ViewSpaceExternalId
    fields: list[Field] = field(default_factory=list)

    @classmethod
    def from_view(cls, view: dm.View, data_class: pygen_config.DataClassNaming) -> Self:
        view_name = (view.name or view.external_id).replace(" ", "_")
        class_name = create_name(view_name, data_class.name)
        variable_name = create_name(view_name, data_class.variable)
        variable_list = create_name(view_name, data_class.variable_list)
        if variable_name == variable_list:
            variable_list = f"{variable_list}_list"
        file_name = f"_{create_name(view_name, data_class.file)}"
        return cls(
            view_name=view_name,
            read_name=class_name,
            write_name=f"{class_name}Apply",
            read_list_name=f"{class_name}List",
            write_list_name=f"{class_name}ApplyList",
            variable=variable_name,
            variable_list=variable_list,
            file_name=file_name,
            view_id=ViewSpaceExternalId.from_(view),
        )

    def update_fields(
        self,
        properties: dict[str, dm.MappedProperty | dm.ConnectionDefinition],
        data_class_by_view_id: dict[ViewSpaceExternalId, DataClass],
        field_naming: pygen_config.FieldNaming,
    ) -> None:
        pydantic_field = self.pydantic_field
        for prop_name, prop in properties.items():
            field_ = Field.from_property(
                prop_name, prop, data_class_by_view_id, field_naming, self.view_name, pydantic_field=pydantic_field
            )
            self.fields.append(field_)

    @property
    def pydantic_field(self) -> str:
        if any(
            name == "Field" for name in [self.read_name, self.write_name, self.read_list_name, self.write_list_name]
        ):
            return "pydantic.Field"
        else:
            return "Field"

    @property
    def init_import(self) -> str:
        return (
            f"from .{self.file_name} "
            f"import {self.read_name}, {self.write_name}, {self.read_list_name}, {self.write_list_name}"
        )

    def __iter__(self) -> Iterator[Field]:
        return iter(self.fields)

    @property
    def one_to_one_edges(self) -> Iterable[EdgeOneToOne]:
        return (field_ for field_ in self.fields if isinstance(field_, EdgeOneToOne))

    @property
    def one_to_many_edges(self) -> Iterable[EdgeOneToMany]:
        return (field_ for field_ in self.fields if isinstance(field_, EdgeOneToMany))

    @property
    def primitive_fields(self) -> Iterable[PrimitiveField]:
        return (field_ for field_ in self.fields if isinstance(field_, PrimitiveField))

    @property
    def cdf_external_fields(self) -> Iterable[CDFExternalField]:
        return (field_ for field_ in self.fields if isinstance(field_, CDFExternalField))

    @property
    def single_timeseries_fields(self) -> Iterable[CDFExternalField]:
        return (field_ for field_ in self.cdf_external_fields if isinstance(field_.prop.type, dm.TimeSeriesReference))

    @property
    def has_one_to_many_edges(self) -> bool:
        return any(isinstance(field_, EdgeOneToMany) for field_ in self.fields)

    @property
    def has_edges(self) -> bool:
        return any(isinstance(field_, EdgeField) for field_ in self.fields)

    @property
    def has_only_one_to_many_edges(self) -> bool:
        return all(isinstance(field_, EdgeOneToMany) for field_ in self.fields)

    @property
    def fields_by_container(self) -> dict[dm.ContainerId, list[PrimitiveFieldCore | EdgeOneToOne]]:
        result: dict[dm.ContainerId, list[PrimitiveFieldCore | EdgeOneToOne]] = defaultdict(list)
        for field_ in self:
            if isinstance(field_, (PrimitiveFieldCore, EdgeOneToOne)):
                result[field_.prop.container].append(field_)
        return dict(result)

    @property
    def has_time_field(self) -> bool:
        return any(field_.is_time_field for field_ in self.fields)

    @property
    def _field_type_hints(self) -> Iterable[str]:
        return (hint for field_ in self.fields for hint in (field_.as_read_type_hint(), field_.as_write_type_hint()))

    @property
    def use_optional_type(self) -> bool:
        return any("Optional" in hint for hint in self._field_type_hints)

    @property
    def use_pydantic_field(self) -> bool:
        pydantic_field = self.pydantic_field
        return any(pydantic_field in hint for hint in self._field_type_hints)

    @property
    def import_pydantic_field(self) -> str:
        if self.pydantic_field == "Field":
            return "from pydantic import Field"
        else:
            return "import pydantic"

    @property
    def dependencies(self) -> list[DataClass]:
        unique: dict[ViewSpaceExternalId, DataClass] = {}
        for field_ in self.fields:
            if isinstance(field_, EdgeField):
                # This will overwrite any existing data class with the same view id
                # however, this is not a problem as all data classes are uniquely identified by their view id
                unique[field_.data_class.view_id] = field_.data_class
        return sorted(unique.values(), key=lambda x: x.write_name)

    @property
    def has_single_timeseries_fields(self) -> bool:
        return any(
            isinstance(field_.prop.type, dm.TimeSeriesReference) and not isinstance(field_, PrimitiveListField)
            for field_ in self.single_timeseries_fields
        )

    @property
    def primitive_fields_literal(self) -> str:
        return ", ".join(
            f'"{field_.prop_name}"' for field_ in self if isinstance(field_, (PrimitiveField, CDFExternalField))
        )


@dataclass(frozen=True)
class APIClass:
    client_attribute: str
    name: str
    file_name: str
    view_id: ViewSpaceExternalId

    @classmethod
    def from_view(cls, view: dm.View, api_class: pygen_config.APIClassNaming) -> APIClass:
        raw_name = view.name or view.external_id

        raw_name = raw_name.replace(" ", "_")

        return cls(
            client_attribute=create_name(raw_name, api_class.client_attribute),
            name=f"{create_name(raw_name, api_class.name)}API",
            file_name=create_name(raw_name, api_class.file_name),
            view_id=ViewSpaceExternalId.from_(view),
        )


@dataclass(frozen=True)
class MultiAPIClass:
    """
    This represents a set of APIs which are generated from a single data model.

    The motivation for having this class is the case when you want to create one SDK for multiple data models.
    """

    sub_apis: list[APIClass]
    client_attribute: str
    name: str
    model: dm.DataModel[dm.View]

    @property
    def model_id(self) -> dm.DataModelId:
        return self.model.as_id()

    @classmethod
    def from_data_model(
        cls,
        data_model: dm.DataModel[dm.View],
        api_class_by_view_id: dict[ViewSpaceExternalId, APIClass],
        multi_api_class: pygen_config.MultiAPIClassNaming,
    ) -> MultiAPIClass:
        sub_apis = sorted(
            [api_class_by_view_id[ViewSpaceExternalId.from_(view)] for view in data_model.views],
            key=lambda api: api.name,
        )

        data_model_name = data_model.name or data_model.external_id

        return cls(
            sub_apis=sub_apis,
            client_attribute=create_name(data_model_name, multi_api_class.client_attribute),
            name=f"{create_name(data_model_name, multi_api_class.name)}APIs",
            model=data_model,
        )


@dataclass
class ListParameter:
    name: str
    type_: str
    default: None = None

    @property
    def annotation(self) -> str:
        return f"{self.type_} | None"

    @property
    def is_time(self) -> bool:
        return self.type_ in ("datetime.datetime", "datetime.date")


@dataclass
class ListFilter:
    filter: type[dm.Filter]
    prop_name: str
    keyword_arguments: dict[str, ListParameter]

    @property
    def condition(self) -> str:
        if self.filter is dm.filters.In:
            parameter = next(iter(self.keyword_arguments.values())).name
            return f"{parameter} and isinstance({parameter}, list)"
        elif self.filter is dm.filters.Equals:
            parameter = next(iter(self.keyword_arguments.values())).name
            return f"{parameter} and isinstance({parameter}, str)"

        return " or ".join(arg.name for arg in self.keyword_arguments.values())

    @property
    def arguments(self) -> str:
        if self.prop_name == "externalId":
            property_ref = '["node", "externalId"], '
        else:
            property_ref = f'view_id.as_property_ref("{self.prop_name}"), '

        filter_args = ", ".join(
            (
                f"{keyword}={arg.name}.isoformat() if {arg.name} else None" if arg.is_time else f"{keyword}={arg.name}"
                for keyword, arg in self.keyword_arguments.items()
            )
        )
        return f"{property_ref}{filter_args}"

    @property
    def filter_call(self) -> str:
        return f"dm.filters.{self.filter.__name__}"


# This field is used when creating the list method.
_EXTERNAL_ID_FIELD = PrimitiveField(
    name="external_id",
    prop_name="externalId",
    type_="str",
    is_nullable=False,
    default=None,
    prop=dm.MappedProperty(
        container_property_identifier="externalId",
        type=dm.Text(),
        nullable=False,
        auto_increment=False,
        container=dm.ContainerId("dummy", "dummy"),
    ),
    pydantic_field="Field",
)


@dataclass
class ListMethod:
    parameters: list[ListParameter]
    filters: list[ListFilter]

    @classmethod
    def from_fields(cls, fields: Iterable[Field], config: pygen_config.ListMethodFilters) -> Self:
        parameters_by_name: dict[str, ListParameter] = {}
        list_filters: list[ListFilter] = []

        for field_ in itertools.chain(fields, (_EXTERNAL_ID_FIELD,)):
            if not isinstance(field_, PrimitiveField):
                # Only primitive fields supported for now
                continue

            for selected_filter in config.get(field_.prop.type, field_.prop_name):
                if selected_filter is dm.filters.Equals:
                    if field_.name not in parameters_by_name:
                        parameter = ListParameter(name=field_.name, type_=field_.type_)
                        parameters_by_name[parameter.name] = parameter
                    else:
                        # Equals and In filter share parameter, you have to extend the type hint.
                        parameter = parameters_by_name[field_.name]
                        parameter.type_ = f"{field_.type_} | {parameter.type_}"
                    list_filters.append(
                        ListFilter(
                            filter=selected_filter,
                            prop_name=field_.prop_name,
                            keyword_arguments=dict(value=parameter),
                        )
                    )
                elif selected_filter is dm.filters.In:
                    if field_.name not in parameters_by_name:
                        parameter = ListParameter(field_.name, type_=f"list[{field_.type_}]")
                        parameters_by_name[parameter.name] = parameter
                    else:
                        # Equals and In filter share parameter, you have to extend the typ hint.
                        parameter = parameters_by_name[field_.name]
                        parameter.type_ = f"{parameter.type_} | list[{field_.type_}]"
                    list_filters.append(
                        ListFilter(
                            filter=selected_filter,
                            prop_name=field_.prop_name,
                            keyword_arguments=dict(values=parameter),
                        )
                    )
                elif selected_filter is dm.filters.Prefix:
                    parameter = ListParameter(name=f"{field_.name}_prefix", type_=field_.type_)
                    parameters_by_name[parameter.name] = parameter
                    list_filters.append(
                        ListFilter(
                            filter=selected_filter,
                            prop_name=field_.prop_name,
                            keyword_arguments=dict(value=parameter),
                        )
                    )
                elif selected_filter is dm.filters.Range:
                    min_parameter = ListParameter(name=f"min_{field_.name}", type_=field_.type_)
                    max_parameter = ListParameter(name=f"max_{field_.name}", type_=field_.type_)
                    parameters_by_name[min_parameter.name] = min_parameter
                    parameters_by_name[max_parameter.name] = max_parameter
                    list_filters.append(
                        ListFilter(
                            filter=selected_filter,
                            prop_name=field_.prop_name,
                            keyword_arguments=dict(gte=min_parameter, lte=max_parameter),
                        )
                    )
                else:
                    # This is a filter that is not supported by the list method.
                    continue
        return cls(
            parameters=list(parameters_by_name.values()),
            filters=list_filters,
        )


def _to_python_type(type_: dm.DirectRelationReference | dm.PropertyType) -> str:
    if isinstance(type_, (dm.Int32, dm.Int64)):
        out_type = "int"
    elif isinstance(type_, dm.Boolean):
        out_type = "bool"
    elif isinstance(type_, (dm.Float32, dm.Float64)):
        out_type = "float"
    elif isinstance(type_, dm.Date):
        out_type = "datetime.date"
    elif isinstance(type_, dm.Timestamp):
        out_type = "datetime.datetime"
    elif isinstance(type_, dm.Json):
        out_type = "dict"
    elif isinstance(type_, (dm.Text, dm.DirectRelation, dm.CDFExternalIdReference, dm.DirectRelationReference)):
        out_type = "str"
    else:
        raise ValueError(f"Unknown type {type_}")

    return out_type
