"""This module contains the definition of `TransformationRules` pydantic model and all
its sub-models and validators.
"""

import math
import sys
import types
from abc import ABC, abstractmethod
from collections.abc import Callable, Hashable, Iterator, MutableSequence, Sequence
from datetime import datetime
from typing import (
    Annotated,
    Any,
    ClassVar,
    Literal,
    SupportsIndex,
    TypeVar,
    get_args,
    get_origin,
    overload,
)

import pandas as pd
from cognite.client import data_modeling as dm
from pydantic import (
    BaseModel,
    BeforeValidator,
    ConfigDict,
    Field,
    GetCoreSchemaHandler,
    PlainSerializer,
    field_validator,
    model_serializer,
)
from pydantic.main import IncEx
from pydantic_core import core_schema
from rdflib import Namespace, URIRef

from cognite.neat._constants import DEFAULT_NAMESPACE
from cognite.neat._rules.models._types import (
    ContainerEntityType,
    DataModelExternalIdType,
    DmsPropertyType,
    SpaceType,
    StrListType,
    URIRefType,
    VersionType,
    ViewEntityType,
)
from cognite.neat._rules.models.data_types import DataType
from cognite.neat._rules.models.entities import EdgeEntity, ReverseConnectionEntity, ViewEntity
from cognite.neat._utils.rdf_ import uri_display_name

if sys.version_info >= (3, 11):
    from enum import StrEnum
else:
    from backports.strenum import StrEnum


METADATA_VALUE_MAX_LENGTH = 5120


def _get_required_fields(model: type[BaseModel], use_alias: bool = False) -> set[str]:
    """Get required fields from a pydantic model.

    Parameters
    ----------
    model : type[BaseModel]
        Pydantic data model
    use_alias : bool, optional
        Whether to return field alias name, by default False

    Returns
    -------
    list[str]
        List of required fields
    """
    required_fields = set()
    for name, field in model.model_fields.items():
        if not field.is_required():
            continue

        alias = getattr(field, "alias", None)
        if use_alias and alias:
            required_fields.add(alias)
        else:
            required_fields.add(name)
    return required_fields


class SchemaCompleteness(StrEnum):
    complete = "complete"
    partial = "partial"
    extended = "extended"


class ExtensionCategory(StrEnum):
    addition = "addition"
    reshape = "reshape"
    rebuild = "rebuild"


class DataModelType(StrEnum):
    solution = "solution"
    enterprise = "enterprise"


class DataModelAspect(StrEnum):
    conceptual = "conceptual"
    logical = "logical"
    physical = "physical"


class RoleTypes(StrEnum):
    domain_expert = "domain expert"
    information = "information architect"
    dms = "DMS Architect"


class SchemaModel(BaseModel):
    model_config: ClassVar[ConfigDict] = ConfigDict(
        populate_by_name=True,
        str_strip_whitespace=True,
        arbitrary_types_allowed=True,
        strict=False,
        extra="ignore",
        use_enum_values=True,
    )

    @classmethod
    def mandatory_fields(cls, use_alias=False) -> set[str]:
        """Returns a set of mandatory fields for the model."""
        return _get_required_fields(cls, use_alias)

    @field_validator("*", mode="before")
    def strip_string(cls, value: Any) -> Any:
        if isinstance(value, str):
            return value.strip()
        elif isinstance(value, list):
            return [entry.strip() if isinstance(entry, str) else entry for entry in value]
        return value


class BaseMetadata(SchemaModel):
    """
    Metadata model for data model
    """

    role: ClassVar[RoleTypes] = Field(description="Role of the person creating the data model")
    aspect: ClassVar[DataModelAspect] = Field(description="Aspect of the data model")
    space: SpaceType = Field(description="The space where the data model is defined")
    external_id: DataModelExternalIdType = Field(
        alias="externalId", description="External identifier for the data model"
    )
    version: VersionType = Field(description="Version of the data model")

    name: str | None = Field(
        None,
        description="Human readable name of the data model",
        min_length=1,
        max_length=255,
    )

    description: str | None = Field(
        None, min_length=1, max_length=1024, description="Short description of the data model"
    )

    creator: StrListType = Field(
        description=(
            "List of contributors (comma seperated) to the data model creation, "
            "typically information architects are considered as contributors."
        ),
    )

    created: datetime = Field(
        description="Date of the data model creation",
    )

    updated: datetime = Field(
        description="Date of the data model update",
    )

    source_id: URIRefType | None = Field(
        None,
        description="Id of source that produced this rules",
        alias="sourceId",
    )

    @field_validator("*", mode="before")
    def strip_string(cls, value: Any) -> Any:
        if isinstance(value, str):
            return value.strip()
        return value

    @field_validator("description", mode="before")
    def nan_as_none(cls, value):
        if isinstance(value, float) and math.isnan(value):
            return None
        return value

    def to_pandas(self) -> pd.Series:
        """Converts Metadata to pandas Series."""
        return pd.Series(self.model_dump())

    def _repr_html_(self) -> str:
        """Returns HTML representation of Metadata."""
        return self.to_pandas().to_frame("value")._repr_html_()  # type: ignore[operator]

    @classmethod
    def mandatory_fields(cls, use_alias=False) -> set[str]:
        """Returns a set of mandatory fields for the model."""
        return _get_required_fields(cls, use_alias)

    @model_serializer(mode="wrap")
    def include_role(self, serializer: Callable) -> dict:
        return {"role": self.role.value, **serializer(self)}

    @property
    def prefix(self) -> str:
        return self.space

    def get_prefix(self) -> str:
        return self.prefix

    @property
    def identifier(self) -> URIRef:
        """Globally unique identifier for the data model.

        !!! note
            Unlike namespace, the identifier does not end with "/" or "#".

        """
        return DEFAULT_NAMESPACE[f"data-model/verified/{self.aspect}/{self.space}/{self.external_id}/{self.version}"]

    @property
    def namespace(self) -> Namespace:
        """Namespace for the data model used for the entities in the data model."""
        return Namespace(f"{self.identifier}/")

    def as_data_model_id(self) -> dm.DataModelId:
        return dm.DataModelId(space=self.space, external_id=self.external_id, version=self.version)

    def as_identifier(self) -> str:
        return repr(self.as_data_model_id())


class BaseRules(SchemaModel, ABC):
    """
    Rules is a core concept in `neat`. This represents fusion of data model
    definitions and (optionally) the transformation rules used to transform the data/graph
    from the source representation to the target representation defined by the data model.
    The rules are defined in an Excel sheet and then parsed into a `Rules` object. The
    `Rules` object is then used to generate data model and the `RDF` graph made of data
    model instances.

    Args:
        metadata: Data model metadata
    """

    metadata: BaseMetadata

    @classmethod
    def headers_by_sheet(cls, by_alias: bool = False) -> dict[str, list[str]]:
        """Returns a list of headers for the model, typically used by ExcelExporter"""
        headers_by_sheet: dict[str, list[str]] = {}
        for field_name, field in cls.model_fields.items():
            if field_name in ["validators_to_skip", "post_validate"]:
                continue
            sheet_name = (field.alias or field_name) if by_alias else field_name
            annotation = field.annotation

            if isinstance(annotation, types.UnionType):
                annotation = annotation.__args__[0]

            try:
                if isinstance(annotation, types.GenericAlias) and get_origin(annotation) is SheetList:
                    # We know that this is a SheetList, so we can safely access the annotation
                    # which is the concrete type of the SheetEntity.
                    model_fields = get_args(annotation)[0].model_fields  # type: ignore[union-attr]
                elif isinstance(annotation, type) and issubclass(annotation, BaseModel):
                    model_fields = annotation.model_fields
                else:
                    model_fields = {}
            except TypeError:
                # Python 3.10 raises TypeError: issubclass() arg 1 must be a class
                # when calling issubclass(annotation, SheetList) with the dict annotation
                model_fields = {}
            headers_by_sheet[sheet_name] = [
                (field.alias or field_name) if by_alias else field_name
                for field_name, field in model_fields.items()
                if field_name != "validators_to_skip" and not field.exclude
            ]
        return headers_by_sheet

    @property
    def display_name(self):
        return uri_display_name(self.metadata.identifier)

    def dump(
        self,
        entities_exclude_defaults: bool = True,
        sort: bool = False,
        mode: Literal["python", "json"] = "python",
        by_alias: bool = False,
        exclude: IncEx | None = None,
        exclude_none: bool = False,
        exclude_unset: bool = False,
        exclude_defaults: bool = False,
    ) -> dict[str, Any]:
        """Dump the model to a dictionary.

        This is used in the Exporters to dump rules in the required format.

        Args:
            entities_exclude_defaults: Whether to exclude default prefix (and version) for entities.
                For example, given a class that is dumped as 'my_prefix:MyClass', if the prefix for the rules
                set in metadata.prefix = 'my_prefix', then this class will be dumped as 'MyClass' when this flag is set.
                Defaults to True.
            sort: Whether to sort the entities in the output.
            mode: The mode in which `to_python` should run.
                If mode is 'json', the output will only contain JSON serializable types.
                If mode is 'python', the output may contain non-JSON-serializable Python objects.
            by_alias: Whether to use the field's alias in the dictionary key if defined.
            exclude: A set of fields to exclude from the output.
            exclude_none: Whether to exclude fields that have a value of `None`.
            exclude_unset: Whether to exclude fields that have not been explicitly set.
            exclude_defaults: Whether to exclude fields that are set to their default value.
        """
        if sort:
            for field_name in self.model_fields.keys():
                value = getattr(self, field_name)
                # Ensure deterministic order of properties, classes, views, and so on
                if isinstance(value, SheetList):
                    value.sort(key=lambda x: x._identifier())

        context: dict[str, Any] = {}
        if entities_exclude_defaults:
            context["metadata"] = self.metadata

        exclude_input: IncEx | None = exclude

        return self.model_dump(
            mode=mode,
            by_alias=by_alias,
            exclude=exclude_input,
            exclude_none=exclude_none,
            exclude_unset=exclude_unset,
            exclude_defaults=exclude_defaults,
            context=context,
        )


class SheetRow(SchemaModel):
    neatId: URIRefType | None = Field(
        alias="Neat ID",
        description="Globally unique identifier for the property",
        default=None,
    )

    @abstractmethod
    def _identifier(self) -> tuple[Hashable, ...]:
        raise NotImplementedError()

    def __repr__(self) -> str:
        # Simplified representation of the object for debugging
        return f"{self.__class__.__name__}({self._identifier()})"


T_SheetRow = TypeVar("T_SheetRow", bound=SheetRow)


class SheetList(list, MutableSequence[T_SheetRow]):
    @classmethod
    def __get_pydantic_core_schema__(cls, source: Any, handler: GetCoreSchemaHandler) -> core_schema.CoreSchema:
        if args := get_args(source):
            item_type = args[0]
        else:
            # Someone use SheetList without specifying the type
            raise TypeError("SheetList must be used with a type argument, e.g., SheetList[InformationProperty]")

        instance_schema = core_schema.is_instance_schema(cls)
        sequence_row_schema = handler.generate_schema(Sequence[item_type])  # type: ignore[valid-type]

        non_instance_schema = core_schema.no_info_after_validator_function(SheetList, sequence_row_schema)
        return core_schema.union_schema([instance_schema, non_instance_schema])

    def to_pandas(self, drop_na_columns: bool = True, include: list[str] | None = None) -> pd.DataFrame:
        """Converts ResourceDict to pandas DataFrame."""
        df = pd.DataFrame([entity.model_dump() for entity in self])
        if drop_na_columns:
            df = df.dropna(axis=1, how="all")
        if include is not None:
            df = df[include]
        return df

    def _repr_html_(self) -> str:
        """Returns HTML representation of ResourceDict."""
        df = self.to_pandas(drop_na_columns=True)
        if "neatId" in df.columns:
            df = df.drop(columns=["neatId"])
        return df._repr_html_()  # type: ignore[operator]

    # Implemented to get correct type hints
    def __iter__(self) -> Iterator[T_SheetRow]:
        return super().__iter__()

    @overload
    def __getitem__(self, index: SupportsIndex) -> T_SheetRow: ...

    @overload
    def __getitem__(self, index: slice) -> "SheetList[T_SheetRow]": ...

    def __getitem__(self, index: SupportsIndex | slice, /) -> "T_SheetRow | SheetList[T_SheetRow]":
        if isinstance(index, slice):
            return SheetList[T_SheetRow](super().__getitem__(index))
        return super().__getitem__(index)


ExtensionCategoryType = Annotated[
    ExtensionCategory,
    PlainSerializer(
        lambda v: v.value if isinstance(v, ExtensionCategory) else v,
        return_type=str,
        when_used="unless-none",
    ),
    BeforeValidator(lambda v: ExtensionCategory(v) if isinstance(v, str) else v),
]


# Immutable such that this can be used as a key in a dictionary
class ContainerProperty(BaseModel, frozen=True):
    container: ContainerEntityType
    property_: DmsPropertyType


class ContainerDestinationProperty(ContainerProperty, frozen=True):
    value_type: DataType | ViewEntity
    connection: Literal["direct"] | ReverseConnectionEntity | EdgeEntity | None = None


class ViewRef(BaseModel, frozen=True):
    view: ViewEntityType


class ViewProperty(ViewRef, frozen=True):
    property_: DmsPropertyType
