import sys
import types
from collections.abc import Generator
from contextlib import contextmanager
from contextvars import ContextVar
from dataclasses import dataclass
from typing import (
    TYPE_CHECKING,
    Annotated,
    Any,
    ForwardRef,
    TypeAlias,
    TypeVar,
    Union,
    get_args,
    get_origin,
)

from annotated_types import MaxLen
from pydantic import VERSION as P_VERSION
from pydantic import BaseModel
from pydantic import ConfigDict as ConfigDict
from pydantic._internal._fields import PydanticMetadata
from pydantic._internal._model_construction import ModelMetaclass as ModelMetaclass
from pydantic._internal._repr import Representation as Representation
from pydantic.fields import FieldInfo
from pydantic_core import PydanticUndefined as Undefined
from pydantic_core import PydanticUndefinedType as PydanticUndefinedType

BaseConfig = ConfigDict
UndefinedType = PydanticUndefinedType
PYDANTIC_MINOR_VERSION = tuple(int(i) for i in P_VERSION.split(".")[:2])


if TYPE_CHECKING:
    from .main import RelationshipInfo, SQLModel

UnionType = getattr(types, "UnionType", Union)
NoneType = type(None)
T = TypeVar("T")
InstanceOrType: TypeAlias = T | type[T]
_TSQLModel = TypeVar("_TSQLModel", bound="SQLModel")


class FakeMetadata:
    max_length: int | None = None
    max_digits: int | None = None
    decimal_places: int | None = None


@dataclass
class ObjectWithUpdateWrapper:
    obj: Any
    update: dict[str, Any]

    def __getattribute__(self, __name: str) -> Any:
        update = super().__getattribute__("update")
        obj = super().__getattribute__("obj")
        if __name in update:
            return update[__name]
        return getattr(obj, __name)


def _is_union_type(t: Any) -> bool:
    return t is UnionType or t is Union


finish_init: ContextVar[bool] = ContextVar("finish_init", default=True)


@contextmanager
def partial_init() -> Generator[None, None, None]:
    token = finish_init.set(False)
    yield
    finish_init.reset(token)


class SQLModelConfig(BaseConfig, total=False):
    table: bool | None
    registry: Any | None


def get_model_fields(model: InstanceOrType[BaseModel]) -> dict[str, "FieldInfo"]:
    # TODO: refactor the usage of this function to always pass the class
    # not the instance, and then remove this extra check
    # this is for compatibility with Pydantic v3
    if isinstance(model, type):
        use_model = model
    else:
        use_model = model.__class__
    return use_model.model_fields


def init_pydantic_private_attrs(new_object: InstanceOrType["SQLModel"]) -> None:
    object.__setattr__(new_object, "__pydantic_fields_set__", set())
    object.__setattr__(new_object, "__pydantic_extra__", None)
    object.__setattr__(new_object, "__pydantic_private__", None)


def get_annotations(class_dict: dict[str, Any]) -> dict[str, Any]:
    raw_annotations: dict[str, Any] = class_dict.get("__annotations__", {})
    if sys.version_info >= (3, 14) and "__annotations__" not in class_dict:
        # See https://github.com/pydantic/pydantic/pull/11991
        from annotationlib import (
            Format,
            call_annotate_function,
            get_annotate_from_class_namespace,
        )

        if annotate := get_annotate_from_class_namespace(class_dict):
            raw_annotations = call_annotate_function(annotate, format=Format.FORWARDREF)
    return raw_annotations


def is_table_model_class(cls: type[Any]) -> bool:
    config = getattr(cls, "model_config", {})
    if config:
        return config.get("table", False) or False
    return False


def get_relationship_to(
    name: str,
    rel_info: "RelationshipInfo",
    annotation: Any,
) -> Any:
    origin = get_origin(annotation)
    use_annotation = annotation
    # Direct relationships (e.g. 'Team' or Team) have None as an origin
    if origin is None:
        if isinstance(use_annotation, ForwardRef):
            use_annotation = use_annotation.__forward_arg__
        else:
            return use_annotation
    # If Union (e.g. Optional), get the real field
    elif _is_union_type(origin):
        use_annotation = get_args(annotation)
        if len(use_annotation) > 2:
            raise ValueError("Cannot have a (non-optional) union as a SQLAlchemy field")
        arg1, arg2 = use_annotation
        if arg1 is NoneType and arg2 is not NoneType:
            use_annotation = arg2
        elif arg2 is NoneType and arg1 is not NoneType:
            use_annotation = arg1
        else:
            raise ValueError(
                "Cannot have a Union of None and None as a SQLAlchemy field"
            )

    # If a list, then also get the real field
    elif origin is list:
        use_annotation = get_args(annotation)[0]

    return get_relationship_to(name=name, rel_info=rel_info, annotation=use_annotation)


def is_field_noneable(field: "FieldInfo") -> bool:
    if getattr(field, "nullable", Undefined) is not Undefined:
        return field.nullable  # type: ignore
    origin = get_origin(field.annotation)
    if origin is not None and _is_union_type(origin):
        args = get_args(field.annotation)
        if any(arg is NoneType for arg in args):
            return True
    if not field.is_required():
        if field.default is Undefined:
            return False
        if field.annotation is None or field.annotation is NoneType:
            return True
        return False
    return False


def get_sa_type_from_type_annotation(annotation: Any) -> Any:
    # Resolve Optional fields
    if annotation is None:
        raise ValueError("Missing field type")
    origin = get_origin(annotation)
    if origin is None:
        return annotation
    elif origin is Annotated:
        return get_sa_type_from_type_annotation(get_args(annotation)[0])
    if _is_union_type(origin):
        bases = get_args(annotation)
        if len(bases) > 2:
            raise ValueError("Cannot have a (non-optional) union as a SQLAlchemy field")
        # Non-optional unions are not allowed
        if bases[0] is not NoneType and bases[1] is not NoneType:
            raise ValueError("Cannot have a (non-optional) union as a SQLAlchemy field")
        # Optional unions are allowed
        use_type = bases[0] if bases[0] is not NoneType else bases[1]
        return get_sa_type_from_type_annotation(use_type)
    return origin


def get_sa_type_from_field(field: Any) -> Any:
    type_: Any = field.annotation
    return get_sa_type_from_type_annotation(type_)


def get_field_metadata(field: Any) -> Any:
    for meta in field.metadata:
        if isinstance(meta, (PydanticMetadata, MaxLen)):
            return meta
    return FakeMetadata()


def sqlmodel_table_construct(
    *,
    self_instance: _TSQLModel,
    values: dict[str, Any],
    _fields_set: set[str] | None = None,
) -> _TSQLModel:
    # Copy from Pydantic's BaseModel.construct()
    # Ref: https://github.com/pydantic/pydantic/blob/v2.5.2/pydantic/main.py#L198
    # Modified to not include everything, only the model fields, and to
    # set relationships
    # SQLModel override to get class SQLAlchemy __dict__ attributes and
    # set them back in after creating the object
    # new_obj = cls.__new__(cls)
    cls = type(self_instance)
    old_dict = self_instance.__dict__.copy()
    # End SQLModel override

    fields_values: dict[str, Any] = {}
    defaults: dict[
        str, Any
    ] = {}  # keeping this separate from `fields_values` helps us compute `_fields_set`
    for name, field in cls.model_fields.items():
        if field.alias and field.alias in values:
            fields_values[name] = values.pop(field.alias)
        elif name in values:
            fields_values[name] = values.pop(name)
        elif not field.is_required():
            defaults[name] = field.get_default(call_default_factory=True)
    if _fields_set is None:
        _fields_set = set(fields_values.keys())
    fields_values.update(defaults)

    _extra: dict[str, Any] | None = None
    if cls.model_config.get("extra") == "allow":
        _extra = {}
        for k, v in values.items():
            _extra[k] = v
    # SQLModel override, do not include everything, only the model fields
    # else:
    #     fields_values.update(values)
    # End SQLModel override
    # SQLModel override
    # Do not set __dict__, instead use setattr to trigger SQLAlchemy
    # object.__setattr__(new_obj, "__dict__", fields_values)
    # instrumentation
    for key, value in {**old_dict, **fields_values}.items():
        setattr(self_instance, key, value)
    # End SQLModel override
    object.__setattr__(self_instance, "__pydantic_fields_set__", _fields_set)
    if not cls.__pydantic_root_model__:
        object.__setattr__(self_instance, "__pydantic_extra__", _extra)

    if cls.__pydantic_post_init__:
        self_instance.model_post_init(None)
    elif not cls.__pydantic_root_model__:
        # Note: if there are any private attributes, cls.__pydantic_post_init__ would exist
        # Since it doesn't, that means that `__pydantic_private__` should be set to None
        object.__setattr__(self_instance, "__pydantic_private__", None)
    # SQLModel override, set relationships
    # Get and set any relationship objects
    for key in self_instance.__sqlmodel_relationships__:
        value = values.get(key, Undefined)
        if value is not Undefined:
            setattr(self_instance, key, value)
    # End SQLModel override
    return self_instance


def sqlmodel_validate(
    cls: type[_TSQLModel],
    obj: Any,
    *,
    strict: bool | None = None,
    from_attributes: bool | None = None,
    context: dict[str, Any] | None = None,
    update: dict[str, Any] | None = None,
) -> _TSQLModel:
    if not is_table_model_class(cls):
        new_obj: _TSQLModel = cls.__new__(cls)
    else:
        # If table, create the new instance normally to make SQLAlchemy create
        # the _sa_instance_state attribute
        # The wrapper of this function should use with _partial_init()
        with partial_init():
            new_obj = cls()
    # SQLModel Override to get class SQLAlchemy __dict__ attributes and
    # set them back in after creating the object
    old_dict = new_obj.__dict__.copy()
    use_obj = obj
    if isinstance(obj, dict) and update:
        use_obj = {**obj, **update}
    elif update:
        use_obj = ObjectWithUpdateWrapper(obj=obj, update=update)
    cls.__pydantic_validator__.validate_python(
        use_obj,
        strict=strict,
        from_attributes=from_attributes,
        context=context,
        self_instance=new_obj,
    )
    # Capture fields set to restore it later
    fields_set = new_obj.__pydantic_fields_set__.copy()
    if not is_table_model_class(cls):
        # If not table, normal Pydantic code, set __dict__
        new_obj.__dict__ = {**old_dict, **new_obj.__dict__}
    else:
        # Do not set __dict__, instead use setattr to trigger SQLAlchemy
        # instrumentation
        for key, value in {**old_dict, **new_obj.__dict__}.items():
            setattr(new_obj, key, value)
    # Restore fields set
    object.__setattr__(new_obj, "__pydantic_fields_set__", fields_set)
    # Get and set any relationship objects
    if is_table_model_class(cls):
        for key in new_obj.__sqlmodel_relationships__:
            value = getattr(use_obj, key, Undefined)
            if value is not Undefined:
                setattr(new_obj, key, value)
    return new_obj


def sqlmodel_init(*, self: "SQLModel", data: dict[str, Any]) -> None:
    old_dict = self.__dict__.copy()
    if not is_table_model_class(self.__class__):
        self.__pydantic_validator__.validate_python(
            data,
            self_instance=self,
        )
    else:
        sqlmodel_table_construct(
            self_instance=self,
            values=data,
        )
    object.__setattr__(
        self,
        "__dict__",
        {**old_dict, **self.__dict__},
    )
