import os
import sys
from collections import defaultdict
from collections.abc import Iterable, Iterator, Mapping, Sequence
from contextlib import contextmanager
from functools import cached_property
from typing import TYPE_CHECKING, Any, Literal, Union

from django.core.exceptions import FieldDoesNotExist, FieldError
from django.db import models
from django.db.models.base import Model
from django.db.models.constants import LOOKUP_SEP
from django.db.models.expressions import Expression
from django.db.models.fields import AutoField, CharField, Field
from django.db.models.fields.related import ForeignKey, RelatedField
from django.db.models.fields.reverse_related import ForeignObjectRel
from django.db.models.lookups import Exact
from django.db.models.sql.query import Query
from mypy.checker import TypeChecker
from mypy.nodes import TypeInfo
from mypy.plugin import MethodContext
from mypy.typeanal import make_optional_type
from mypy.types import AnyType, Instance, ProperType, TypeOfAny, UnionType, get_proper_type
from mypy.types import Type as MypyType

from mypy_django_plugin.exceptions import UnregisteredModelError
from mypy_django_plugin.lib import fullnames, helpers

# This import fails when `psycopg2` is not installed, avoid crashing the plugin.
try:
    from django.contrib.postgres.fields import ArrayField
except ImportError:

    class ArrayField:  # type: ignore[no-redef]
        pass


if TYPE_CHECKING:
    from django.apps.registry import Apps
    from django.conf import LazySettings
    from django.db.models.options import _AnyField


@contextmanager
def temp_environ() -> Iterator[None]:
    """Allow the ability to set os.environ temporarily"""
    environ = dict(os.environ)
    try:
        yield
    finally:
        os.environ.clear()
        os.environ.update(environ)


def initialize_django(settings_module: str) -> tuple["Apps", "LazySettings"]:
    with temp_environ():
        os.environ["DJANGO_SETTINGS_MODULE"] = settings_module

        # add current directory to sys.path
        sys.path.append(os.getcwd())

        from django.apps import apps
        from django.conf import settings

        apps.get_swappable_settings_name.cache_clear()  # type: ignore[attr-defined]
        apps.clear_cache()

        if not settings.configured:
            settings._setup()  # type: ignore[misc]
        apps.populate(settings.INSTALLED_APPS)

    assert apps.apps_ready, "Apps are not ready"
    assert settings.configured, "Settings are not configured"

    return apps, settings


class LookupsAreUnsupported(Exception):
    pass


def get_field_type_from_model_type_info(info: TypeInfo | None, field_name: str) -> Instance | None:
    if info is None:
        return None
    field_node = info.get(field_name)
    if field_node is None:
        return None
    field_type = get_proper_type(field_node.type)
    if not isinstance(field_type, Instance):
        return None
    # Field declares a set and a get type arg. Fallback to `None` when we can't find any args
    if len(field_type.args) != 2:
        return None
    return field_type


def _get_field_set_type_from_model_type_info(info: TypeInfo | None, field_name: str) -> MypyType | None:
    field_type = get_field_type_from_model_type_info(info, field_name)
    if field_type is not None:
        return field_type.args[0]
    return None


def _get_field_get_type_from_model_type_info(info: TypeInfo | None, field_name: str) -> MypyType | None:
    field_type = get_field_type_from_model_type_info(info, field_name)
    if field_type is not None:
        return field_type.args[1]
    return None


class DjangoContext:
    def __init__(self, django_settings_module: str) -> None:
        self.django_settings_module = django_settings_module

        apps, settings = initialize_django(self.django_settings_module)
        self.apps_registry = apps
        self.settings = settings

    @cached_property
    def model_modules(self) -> dict[str, dict[str, type[Model]]]:
        """All modules that contain Django models."""
        modules: dict[str, dict[str, type[Model]]] = defaultdict(dict)
        for concrete_model_cls in self.apps_registry.get_models(include_auto_created=True, include_swapped=True):
            modules[concrete_model_cls.__module__][concrete_model_cls.__name__] = concrete_model_cls
            # collect abstract=True models
            for model_cls in concrete_model_cls.mro()[1:]:
                if issubclass(model_cls, Model) and hasattr(model_cls, "_meta") and model_cls._meta.abstract:
                    modules[model_cls.__module__][model_cls.__name__] = model_cls
        return modules

    def get_model_class_by_fullname(self, fullname: str) -> type[Model] | None:
        """Returns None if Model is abstract"""
        module, _, model_cls_name = fullname.rpartition(".")
        return self.model_modules.get(module, {}).get(model_cls_name)

    def get_model_fields(self, model_cls: type[Model]) -> Iterator["Field[Any, Any]"]:
        for field in model_cls._meta.get_fields():
            if isinstance(field, Field):
                yield field

    def get_model_foreign_keys(self, model_cls: type[Model]) -> Iterator["ForeignKey[Any, Any]"]:
        for field in model_cls._meta.get_fields():
            if isinstance(field, ForeignKey):
                yield field

    def get_model_related_fields(self, model_cls: type[Model]) -> Iterator["RelatedField[Any, Any]"]:
        """Get model forward relations"""
        for field in model_cls._meta.get_fields():
            if isinstance(field, RelatedField):
                yield field

    def get_model_relations(self, model_cls: type[Model]) -> Iterator[ForeignObjectRel]:
        """Get model reverse relations"""
        for field in model_cls._meta.get_fields():
            if isinstance(field, ForeignObjectRel):
                yield field

    def get_field_lookup_exact_type(
        self, api: TypeChecker, field: Union["Field[Any, Any]", ForeignObjectRel]
    ) -> MypyType:
        if isinstance(field, RelatedField | ForeignObjectRel):
            related_model_cls = self.get_field_related_model_cls(field)
            rel_model_info = helpers.lookup_class_typeinfo(api, related_model_cls)
            if rel_model_info is None:
                return AnyType(TypeOfAny.explicit)

            primary_key_field = self.get_primary_key_field(related_model_cls)
            primary_key_type = self.get_field_get_type(api, rel_model_info, primary_key_field, method="init")

            model_and_primary_key_type = UnionType.make_union([Instance(rel_model_info, []), primary_key_type])
            return make_optional_type(model_and_primary_key_type)

        field_info = helpers.lookup_class_typeinfo(api, field.__class__)
        if field_info is None:
            return AnyType(TypeOfAny.explicit)
        return helpers.get_private_descriptor_type(field_info, "_pyi_lookup_exact_type", is_nullable=field.null)

    def get_related_target_field(
        self, related_model_cls: type[Model], field: "ForeignKey[Any, Any]"
    ) -> "Field[Any, Any] | None":
        # ForeignKey only supports one `to_fields` item (ForeignObject supports many)
        assert len(field.to_fields) == 1
        to_field_name = field.to_fields[0]
        if to_field_name:
            rel_field = related_model_cls._meta.get_field(to_field_name)
            if not isinstance(rel_field, Field):
                return None  # Not supported
            return rel_field
        return self.get_primary_key_field(related_model_cls)

    def get_primary_key_field(self, model_cls: type[Model]) -> "Field[Any, Any]":
        for field in model_cls._meta.get_fields():
            if isinstance(field, Field):
                if field.primary_key:
                    return field
        raise ValueError("No primary key defined")

    def get_expected_types(self, api: TypeChecker, model_cls: type[Model], *, method: str) -> dict[str, MypyType]:
        contenttypes_in_apps = self.apps_registry.is_installed("django.contrib.contenttypes")

        expected_types = {}
        # add pk if not abstract=True
        if not model_cls._meta.abstract:
            primary_key_field = self.get_primary_key_field(model_cls)
            field_set_type = self.get_field_set_type(api, primary_key_field, method=method)
            expected_types["pk"] = field_set_type

        model_info = helpers.lookup_class_typeinfo(api, model_cls)
        for field in model_cls._meta.get_fields():
            if contenttypes_in_apps:
                from django.contrib.contenttypes.fields import GenericForeignKey

                if isinstance(field, GenericForeignKey):
                    # it's generic, so cannot set specific model
                    field_name = field.name
                    gfk_info = helpers.lookup_class_typeinfo(api, field.__class__)
                    if gfk_info is None:
                        gfk_set_type: MypyType = AnyType(TypeOfAny.unannotated)
                    else:
                        gfk_set_type = helpers.get_private_descriptor_type(
                            gfk_info, "_pyi_private_set_type", is_nullable=True
                        )
                    expected_types[field_name] = gfk_set_type
                    continue

            if isinstance(field, Field):
                field_name = field.attname
                # Can not determine target_field for recursive relationship when model is abstract
                if field.related_model == "self" and model_cls._meta.abstract:
                    continue
                # Try to retrieve set type from a model's TypeInfo object and fallback to retrieving it manually
                # from django-stubs own declaration. This is to align with the setter types declared for
                # assignment.
                field_set_type = _get_field_set_type_from_model_type_info(
                    model_info, field_name
                ) or self.get_field_set_type(api, field, method=method)
                expected_types[field_name] = field_set_type

                if isinstance(field, ForeignKey):
                    field_name = field.name
                    foreign_key_info = helpers.lookup_class_typeinfo(api, field.__class__)
                    if foreign_key_info is None:
                        # maybe there's no type annotation for the field
                        expected_types[field_name] = AnyType(TypeOfAny.unannotated)
                        continue

                    try:
                        related_model = self.get_field_related_model_cls(field)
                    except UnregisteredModelError:
                        # Recognise the field but don't say anything about its type..
                        expected_types[field_name] = AnyType(TypeOfAny.from_error)
                        continue

                    if related_model._meta.proxy_for_model is not None:
                        related_model = related_model._meta.proxy_for_model

                    related_model_info = helpers.lookup_class_typeinfo(api, related_model)
                    if related_model_info is None:
                        expected_types[field_name] = AnyType(TypeOfAny.unannotated)
                        continue

                    is_nullable = self.get_field_nullability(field, method)
                    foreign_key_set_type = helpers.get_private_descriptor_type(
                        foreign_key_info, "_pyi_private_set_type", is_nullable=is_nullable
                    )
                    model_set_type = helpers.convert_any_to_type(foreign_key_set_type, Instance(related_model_info, []))

                    expected_types[field_name] = model_set_type

        return expected_types

    @cached_property
    def all_registered_model_classes(self) -> set[type[models.Model]]:
        model_classes = self.apps_registry.get_models()

        all_model_bases = set()
        for model_cls in model_classes:
            for base_cls in model_cls.mro():
                if issubclass(base_cls, models.Model):
                    all_model_bases.add(base_cls)

        return all_model_bases

    @cached_property
    def model_class_fullnames_by_label(self) -> Mapping[str, str]:
        return {
            klass._meta.label: helpers.get_class_fullname(klass)
            for klass in self.all_registered_model_classes
            if klass is not models.Model
        }

    def get_field_nullability(self, field: Union["Field[Any, Any]", ForeignObjectRel], method: str | None) -> bool:
        if method in ("values", "values_list"):
            return field.null

        nullable = field.null
        if not nullable and isinstance(field, CharField) and field.blank:
            return True
        if method == "__init__":
            if (isinstance(field, Field) and field.primary_key) or isinstance(field, ForeignKey):
                return True
        if method == "create":
            if isinstance(field, AutoField):
                return True
        if isinstance(field, Field) and field.has_default():
            return True
        return nullable

    def get_field_set_type(
        self, api: TypeChecker, field: Union["Field[Any, Any]", ForeignObjectRel], *, method: str
    ) -> MypyType:
        """Get a type of __set__ for this specific Django field."""
        target_field = field
        if isinstance(field, ForeignKey):
            try:
                # We gotta be careful for exceptions when we're triggering '__get__'.
                # Related model could very well be unresolvable
                target_field = field.target_field
            except ValueError:
                return AnyType(TypeOfAny.from_error)

        field_info = helpers.lookup_class_typeinfo(api, target_field.__class__)
        if field_info is None:
            return AnyType(TypeOfAny.from_error)

        field_set_type = helpers.get_private_descriptor_type(
            field_info, "_pyi_private_set_type", is_nullable=self.get_field_nullability(field, method)
        )
        if isinstance(target_field, ArrayField):
            argument_field_type = self.get_field_set_type(api, target_field.base_field, method=method)
            field_set_type = helpers.convert_any_to_type(field_set_type, argument_field_type)
        return field_set_type

    def get_field_get_type(
        self,
        api: TypeChecker,
        model_info: TypeInfo | None,
        field: Union["Field[Any, Any]", ForeignObjectRel],
        *,
        method: str,
    ) -> MypyType:
        """Get a type of __get__ for this specific Django field."""
        if isinstance(field, Field):
            get_type = _get_field_get_type_from_model_type_info(model_info, field.attname)
            if get_type is not None:
                return get_type

        field_info = helpers.lookup_class_typeinfo(api, field.__class__)
        if field_info is None:
            return AnyType(TypeOfAny.unannotated)

        is_nullable = self.get_field_nullability(field, method)
        if isinstance(field, RelatedField):
            related_model_cls = self.get_field_related_model_cls(field)
            rel_model_info = helpers.lookup_class_typeinfo(api, related_model_cls)

            if method in ("values", "values_list"):
                primary_key_field = self.get_primary_key_field(related_model_cls)
                return self.get_field_get_type(api, rel_model_info, primary_key_field, method=method)

            model_info = helpers.lookup_class_typeinfo(api, related_model_cls)
            if model_info is None:
                return AnyType(TypeOfAny.unannotated)

            return Instance(model_info, [])
        return helpers.get_private_descriptor_type(field_info, "_pyi_private_get_type", is_nullable=is_nullable)

    def get_field_related_model_cls(self, field: Union["RelatedField[Any, Any]", ForeignObjectRel]) -> type[Model]:
        if isinstance(field, RelatedField):
            related_model_cls = field.remote_field.model
        else:
            related_model_cls = field.field.model

        if related_model_cls is None:
            raise UnregisteredModelError

        if isinstance(related_model_cls, str):
            if related_model_cls == "self":  # type: ignore[unreachable]
                # same model
                related_model_cls = field.model
            elif "." not in related_model_cls:
                # same file model
                related_model_fullname = f"{field.model.__module__}.{related_model_cls}"
                related_model_cls = self.get_model_class_by_fullname(related_model_fullname)
                if related_model_cls is None:
                    raise UnregisteredModelError
            else:
                try:
                    related_model_cls = self.apps_registry.get_model(related_model_cls)
                except LookupError as e:
                    raise UnregisteredModelError from e

        return related_model_cls

    def _resolve_field_from_parts(
        self, field_parts: Iterable[str], model_cls: type[Model]
    ) -> tuple[Union["Field[Any, Any]", ForeignObjectRel], type[Model]]:
        currently_observed_model = model_cls
        field: _AnyField | None = None
        for field_part in field_parts:
            if field_part == "pk":
                field = self.get_primary_key_field(currently_observed_model)
                continue

            field = currently_observed_model._meta.get_field(field_part)
            if isinstance(field, RelatedField):
                currently_observed_model = self.get_field_related_model_cls(field)
                model_name = currently_observed_model._meta.model_name
                if model_name is not None and field_part == (model_name + "_id"):
                    field = self.get_primary_key_field(currently_observed_model)

            if isinstance(field, ForeignObjectRel):
                currently_observed_model = self.get_field_related_model_cls(field)

        # Guaranteed by `query.solve_lookup_type` before.
        assert isinstance(field, Field | ForeignObjectRel)
        return field, currently_observed_model

    def solve_lookup_type(
        self, model_cls: type[Model], lookup: str
    ) -> tuple[Sequence[str], Sequence[str], Expression | Literal[False]] | None:
        query = Query(model_cls)
        if (lookup == "pk" or lookup.startswith("pk__")) and query.get_meta().pk is None:
            # Primary key lookup when no primary key field is found, model is presumably
            # abstract and we can't say anything about 'pk'.
            return None
        try:
            return query.solve_lookup_type(lookup)
        # This occurs when the following conditions are met:
        # - model_cls._meta.abstract = True
        # - part of the lookup is a foreign key defined on model_cls where the 'to' argument is a string
        # On abstract models the 'to' parameter is only coalesced into the actual class whenever a subclass of it is
        # instantiated, therefore it is never swapped out for abstract base classes.
        except AttributeError:
            pass
        query_parts = lookup.split(LOOKUP_SEP)
        try:
            field = query.get_meta().get_field(query_parts[0])
        except FieldDoesNotExist:
            return None

        if len(query_parts) == 1:
            return [], [query_parts[0]], False

        if not isinstance(field, RelatedField | ForeignObjectRel):
            return None

        related_model = self.get_field_related_model_cls(field)
        sub_query = Query(related_model).solve_lookup_type(LOOKUP_SEP.join(query_parts[1:]))
        entire_query_parts = [query_parts[0], *sub_query[1]]
        return sub_query[0], entire_query_parts, sub_query[2]

    def resolve_lookup_into_field(
        self, model_cls: type[Model], lookup: str
    ) -> tuple[Union["Field[Any, Any]", ForeignObjectRel, None], type[Model]]:
        solved_lookup = self.solve_lookup_type(model_cls, lookup)
        if solved_lookup is None:
            return None, model_cls
        lookup_parts, field_parts, _ = solved_lookup
        if lookup_parts:
            raise LookupsAreUnsupported()
        return self._resolve_field_from_parts(field_parts, model_cls)

    def _resolve_lookup_type_from_lookup_class(
        self, ctx: MethodContext, lookup_cls: type, field: Union["Field[Any, Any]", ForeignObjectRel] | None = None
    ) -> MypyType | None:
        """Resolve the expected type for a lookup class (used both for regular fields and annotated fields)

        Args:
            ctx
            lookup_cls: The Django lookup class (e.g., IsNull, Contains)
            field: Optional field for resolving Field-dependent types (None for annotated fields)

        Returns:
            The resolved type, or None if it couldn't be determined
        """
        lookup_info = helpers.lookup_class_typeinfo(helpers.get_typechecker_api(ctx), lookup_cls)
        if lookup_info is None:
            return None

        for lookup_base in helpers.iter_bases(lookup_info):
            if lookup_base.args and isinstance((lookup_type := get_proper_type(lookup_base.args[0])), Instance):
                # if it's Field, consider lookup_type a __get__ of current field
                if lookup_type.type.fullname == fullnames.FIELD_FULLNAME:
                    if field is None:
                        # No field available (e.g., annotation), can't resolve further
                        return None
                    field_info = helpers.lookup_class_typeinfo(helpers.get_typechecker_api(ctx), field.__class__)
                    if field_info is None:
                        return None
                    return get_proper_type(
                        helpers.get_private_descriptor_type(field_info, "_pyi_private_get_type", is_nullable=field.null)
                    )
                return lookup_type

        return None

    def _resolve_annotated_field_lookup(
        self, ctx: MethodContext, lookup: str, model_instance: Instance
    ) -> MypyType | None:
        """Resolve the expected type for a lookup on an annotated field.

        Args:
            ctx
            lookup: The full lookup string (e.g., 'total__gte' or 'total').
            model_instance: The model instance containing extra_attrs from annotations.

        Returns:
            The resolved mypy type for the lookup, or None if the base field
            is not found in the model's annotations.
        """
        if not helpers.is_annotated_model(model_instance.type) or not model_instance.extra_attrs:
            return None

        lookup_base_field, *annotation_lookup_parts = lookup.split(LOOKUP_SEP)

        if lookup_base_field not in model_instance.extra_attrs.attrs:
            return None

        if annotation_lookup_parts:
            lookup_cls = Field().get_lookup(annotation_lookup_parts[-1])
            if lookup_cls is not None:
                lookup_type = self._resolve_lookup_type_from_lookup_class(ctx, lookup_cls)
                if lookup_type is not None:
                    return lookup_type

        # No lookup suffix or Field-dependent lookup: fall back to annotation type
        return model_instance.extra_attrs.attrs[lookup_base_field]

    def resolve_lookup_expected_type(
        self, ctx: MethodContext, model_cls: type[Model], lookup: str, model_instance: Instance
    ) -> MypyType:
        try:
            # solve_lookup_type uses Django's Query.solve_lookup_type(), which raises
            # FieldError for annotated fields since they don't exist on the actual model...
            solved_lookup = self.solve_lookup_type(model_cls, lookup)
        except FieldError as exc:
            # ...so we handle annotation lookups here
            annotation_lookup = self._resolve_annotated_field_lookup(ctx, lookup, model_instance)
            if annotation_lookup is not None:
                return annotation_lookup

            msg = exc.args[0]
            if model_instance.extra_attrs:
                msg = ", ".join((msg, *model_instance.extra_attrs.attrs.keys()))
            ctx.api.fail(msg, ctx.context)
            return AnyType(TypeOfAny.from_error)

        if solved_lookup is None:
            return AnyType(TypeOfAny.implementation_artifact)
        lookup_parts, field_parts, is_expression = solved_lookup
        if is_expression:
            return AnyType(TypeOfAny.explicit)

        field, _ = self._resolve_field_from_parts(field_parts, model_cls)

        lookup_cls = None
        if lookup_parts:
            lookup = lookup_parts[-1]
            lookup_cls = field.get_lookup(lookup)
            if lookup_cls is None:
                # unknown lookup
                return AnyType(TypeOfAny.explicit)

        if lookup_cls is None or isinstance(lookup_cls, Exact):
            return self.get_field_lookup_exact_type(helpers.get_typechecker_api(ctx), field)

        resolved_type = self._resolve_lookup_type_from_lookup_class(ctx, lookup_cls, field)
        if resolved_type is not None:
            return resolved_type

        return AnyType(TypeOfAny.explicit)

    def resolve_f_expression_type(self, f_expression_type: Instance) -> ProperType:
        return AnyType(TypeOfAny.explicit)

    @cached_property
    def is_contrib_auth_installed(self) -> bool:
        return "django.contrib.auth" in self.settings.INSTALLED_APPS
