import contextlib
import itertools
from collections import defaultdict, namedtuple

from django.core.checks import Error, Warning
from django.core.exceptions import FieldDoesNotExist
from django.db import NotSupportedError
from django.db.backends.utils import names_digest, split_identifier
from django.db.models import FloatField, Index, IntegerField
from django.db.models.lookups import BuiltinLookup
from django.db.models.sql.query import Query
from django.db.models.sql.where import AND, XOR, WhereNode
from pymongo import ASCENDING, DESCENDING
from pymongo.operations import IndexModel, SearchIndexModel

from django_mongodb_backend.fields import ArrayField

from .query_utils import process_rhs

MONGO_INDEX_OPERATORS = {
    "exact": "$eq",
    "gt": "$gt",
    "gte": "$gte",
    "lt": "$lt",
    "lte": "$lte",
    "in": "$in",
}

FieldColumn = namedtuple("FieldColumn", ["field", "column"])


def _get_condition_mql(self, model, schema_editor):
    """Analogous to Index._get_condition_sql()."""
    query = Query(model=model, alias_cols=False)
    where = query.build_where(self.condition)
    compiler = query.get_compiler(connection=schema_editor.connection)
    return where.as_mql_idx(compiler, schema_editor.connection)


def builtin_lookup_idx(self, compiler, connection):
    lhs_mql = self.lhs.target.column
    value = process_rhs(self, compiler, connection)
    try:
        operator = MONGO_INDEX_OPERATORS[self.lookup_name]
    except KeyError:
        raise NotSupportedError(
            f"MongoDB does not support the '{self.lookup_name}' lookup in indexes."
        ) from None
    return {lhs_mql: {operator: value}}


def get_field(model, field_name):
    """
    A version of Model_.meta.get_field() that can retrieve embedded model
    fields.
    """
    path = []
    base_model = model
    *parents, leaf = field_name.split(".")
    for i, name in enumerate(parents):
        field = model._meta.get_field(name)
        path.append(field.column)
        # For EmbeddedModelFields, advance to the embedded model and continue
        # to loop, searching for the next field.
        if hasattr(field, "embedded_model"):
            model = field.embedded_model
        # For PolymorphicEmbeddedModelFields, recurse into each embedded model
        # until the field is found.
        elif models := getattr(field, "embedded_models", None):
            for submodel in models:
                with contextlib.suppress(FieldDoesNotExist):
                    subfield = get_field(submodel, ".".join([*parents[i + 1 :], leaf]))
                    path.extend(subfield.column.split("."))
                    return FieldColumn(subfield.field, ".".join(path))
            raise FieldDoesNotExist(
                f"The models of field '{'.'.join(parents)}' have no field named '{leaf}'."
            )
        else:
            raise FieldDoesNotExist(f"{base_model.__name__} has no field named '{field_name}'.")
    # Add the final field.
    field = model._meta.get_field(leaf)
    path.append(field.column)
    return FieldColumn(field, ".".join(path))


def get_pymongo_index_model(self, model, schema_editor, field=None, column_prefix=""):
    """Return a pymongo IndexModel for this Django Index."""
    if self.contains_expressions:
        return None
    kwargs = {}
    filter_expression = defaultdict(dict)
    if self.condition:
        filter_expression.update(self._get_condition_mql(model, schema_editor))
    if filter_expression:
        kwargs["partialFilterExpression"] = filter_expression
    index_orders = (
        [(column_prefix + field.column, ASCENDING)]
        if field
        else [
            # order is "" if ASCENDING or "DESC" if DESCENDING (see
            # django.db.models.indexes.Index.fields_orders).
            (
                column_prefix + get_field(model, field_name).column,
                ASCENDING if order == "" else DESCENDING,
            )
            for field_name, order in self.fields_orders
        ]
    )
    return IndexModel(index_orders, name=self.name, **kwargs)


def where_node_idx(self, compiler, connection):
    if self.connector == AND:
        operator = "$and"
    elif self.connector == XOR:
        raise NotSupportedError("MongoDB does not support the '^' operator lookup in indexes.")
    else:
        operator = "$or"
    if self.negated:
        raise NotSupportedError("MongoDB does not support the '~' operator in indexes.")
    children_mql = []
    for child in self.children:
        mql = child.as_mql_idx(compiler, connection)
        children_mql.append(mql)
    if len(children_mql) == 1:
        mql = children_mql[0]
    elif len(children_mql) > 1:
        mql = {operator: children_mql}
    else:
        mql = {}
    return mql


class EmbeddedFieldIndexMixin:
    def check(self, model, connection):
        # The parent check reports E012 (nonexistent fields) based on top-level
        # fields. These errors are discarded and the checks redone, accounting
        # for embedded model fields.
        errors = [e for e in super().check(model, connection) if e.id != "models.E012"]
        return errors + self._check_local_fields(model)

    def _check_local_fields(self, model):
        errors = []
        forward_fields = self._get_forward_fields(model)
        for field_name in self.fields:
            if field_name not in forward_fields:
                errors.append(
                    Error(
                        f"'{self.meta_option_name}' refers to the nonexistent "
                        f"field '{field_name}'.",
                        obj=model,
                        id="models.E012",
                    )
                )
        return errors

    def _get_forward_fields(self, model):
        """
        Return the set of forward field paths for the given model, including
        embedded fields.
        """
        forward_fields = set()
        for field in model._meta._get_fields(reverse=False):
            # Recurse into embedded models and flatten their forward fields
            # using dotted paths (e.g. "field.subfield").
            if embedded_model := getattr(field, "embedded_model", None):
                embedded_models = [embedded_model]
            elif embedded_models := getattr(field, "embedded_models", None):
                pass
            else:
                embedded_models = []
            for embedded_model in embedded_models:
                sub_forward_fields = self._get_forward_fields(embedded_model)
                for column in sub_forward_fields:
                    forward_fields.add(f"{field.name}.{column}")
                    if hasattr(field, "attname"):
                        forward_fields.add(f"{field.attname}.{column}")
            # For each field, both the public field name and its attribute name
            # (attname), when present, are included in order to mirror Django's
            # field resolution.
            forward_fields.add(field.name)
            if hasattr(field, "attname"):
                forward_fields.add(field.attname)
        return forward_fields


class EmbeddedFieldIndex(EmbeddedFieldIndexMixin, Index):
    meta_option_name = "indexes"

    def set_name_with_model(self, model):
        """
        Generate a unique name for the index.

        The name is divided into three parts: the table name, the field name,
        and a hash plus suffix. Each part is truncated to fit within Django's
        index name length constraints.

        This method overrides Django's base implementation which uses
        model._meta.get_field(field_name).column, which doesn't work for
        embedded fields.
        """
        _, table_name = split_identifier(model._meta.db_table)
        field_names_with_order = [
            (f"-{column_name}" if order else column_name)
            for column_name, order in self.fields_orders
        ]
        hash_data = [table_name, *field_names_with_order, self.suffix]
        self.name = (
            f"{table_name[:11]}_{self.fields_orders[0][0][:7]}_"
            f"{names_digest(*hash_data, length=6)}_{self.suffix}"
        )
        if self.name[0] == "_" or self.name[0].isdigit():
            self.name = f"D{self.name[1:]}"


class SearchIndex(Index):
    suffix = "six"

    def __init__(
        self, *, fields=(), field_mappings=None, name=None, analyzer=None, search_analyzer=None
    ):
        if field_mappings and not isinstance(field_mappings, dict):
            raise ValueError(
                "field_mappings must be a dictionary mapping field names to their "
                "MongoDB Search index options."
            )
        if analyzer and not isinstance(analyzer, str):
            raise ValueError(f"analyzer must be a string; got: {type(analyzer)}.")
        if search_analyzer and not isinstance(search_analyzer, str):
            raise ValueError(f"search_analyzer must be a string; got: {type(search_analyzer)}.")
        self.field_mappings = field_mappings
        self.analyzer = analyzer
        self.search_analyzer = search_analyzer
        if field_mappings:
            if fields:
                raise ValueError("Cannot provide fields and field_mappings.")
            fields = [*self.field_mappings.keys()]
        super().__init__(fields=fields, name=name)

    def deconstruct(self):
        path, args, kwargs = super().deconstruct()
        if self.field_mappings:
            kwargs["field_mappings"] = self.field_mappings
            del kwargs["fields"]
        if self.analyzer:
            kwargs["analyzer"] = self.analyzer
        if self.search_analyzer:
            kwargs["search_analyzer"] = self.search_analyzer
        return path, args, kwargs

    def check(self, model, connection):
        errors = []
        if not connection.features.supports_search:
            errors.append(
                Warning(
                    f"This MongoDB server does not support {self.__class__.__name__}.",
                    hint=(
                        "The index won't be created. Use a Search-enabled version of MongoDB, "
                        "or silence this warning if you don't care about it."
                    ),
                    obj=model,
                    id="mongodb.indexes.search.W001",
                )
            )
        return errors

    def search_index_data_types(self, db_type):
        """
        Map a model field's type to search index type.
        https://www.mongodb.com/docs/atlas/atlas-search/define-field-mappings/#data-types
        """
        if db_type in {"double", "int", "long"}:
            return "number"
        if db_type == "binData":
            return "string"
        if db_type == "bool":
            return "boolean"
        if db_type == "object":
            return "document"
        if db_type == "array":
            return "embeddedDocuments"
        return db_type

    def get_pymongo_index_model(self, model, schema_editor, field=None, column_prefix=""):
        if not schema_editor.connection.features.supports_search:
            return None
        fields = {}
        for field_name, _ in self.fields_orders:
            field_path = column_prefix + model._meta.get_field(field_name).column
            if self.field_mappings:
                fields[field_path] = self.field_mappings[field_name]
            else:
                field = model._meta.get_field(field_name)
                type_ = self.search_index_data_types(field.db_type(schema_editor.connection))
                fields[field_path] = {"type": type_}
        extra = {}
        if self.analyzer:
            extra["analyzer"] = self.analyzer
        if self.search_analyzer:
            extra["searchAnalyzer"] = self.search_analyzer
        return SearchIndexModel(
            definition={"mappings": {"dynamic": False, "fields": fields}, **extra},
            name=self.name,
        )


class VectorSearchIndex(SearchIndex):
    suffix = "vsi"
    VALID_FIELD_TYPES = frozenset(("boolean", "date", "number", "objectId", "string", "uuid"))
    VALID_SIMILARITIES = frozenset(("cosine", "dotProduct", "euclidean"))

    def __init__(self, *, fields=(), name=None, similarities):
        super().__init__(fields=fields, name=name)
        self.similarities = similarities
        self._multiple_similarities = isinstance(similarities, (tuple, list))
        for func in similarities if self._multiple_similarities else (similarities,):
            if func not in self.VALID_SIMILARITIES:
                raise ValueError(
                    f"'{func}' isn't a valid similarity function "
                    f"({', '.join(sorted(self.VALID_SIMILARITIES))})."
                )
        seen_fields = set()
        for field_name, _ in self.fields_orders:
            if field_name in seen_fields:
                raise ValueError(f"Field '{field_name}' is duplicated in fields.")
            seen_fields.add(field_name)

    def check(self, model, connection):
        errors = super().check(model, connection)
        num_arrayfields = 0
        for field_name, _ in self.fields_orders:
            field = model._meta.get_field(field_name)
            if isinstance(field, ArrayField):
                num_arrayfields += 1
                try:
                    int(field.size)
                except (ValueError, TypeError):
                    errors.append(
                        Error(
                            f"VectorSearchIndex requires 'size' on field '{field_name}'.",
                            obj=model,
                            id="mongodb.indexes.search.E002",
                        )
                    )
                if not isinstance(field.base_field, (FloatField, IntegerField)):
                    errors.append(
                        Error(
                            "VectorSearchIndex requires the base field of "
                            f"ArrayField '{field.name}' to be FloatField or "
                            "IntegerField but is "
                            f"{field.base_field.get_internal_type()}.",
                            obj=model,
                            id="mongodb.indexes.search.E003",
                        )
                    )
            else:
                search_type = self.search_index_data_types(field.db_type(connection))
                if search_type not in self.VALID_FIELD_TYPES:
                    errors.append(
                        Error(
                            "VectorSearchIndex does not support field "
                            f"'{field_name}' ({field.get_internal_type()}).",
                            obj=model,
                            id="mongodb.indexes.search.E004",
                            hint=f"Allowed types are {', '.join(sorted(self.VALID_FIELD_TYPES))}.",
                        )
                    )
        if self._multiple_similarities and num_arrayfields != len(self.similarities):
            errors.append(
                Error(
                    f"VectorSearchIndex requires the same number of similarities "
                    f"and vector fields; {model._meta.object_name} has "
                    f"{num_arrayfields} ArrayField(s) but similarities "
                    f"has {len(self.similarities)} element(s).",
                    obj=model,
                    id="mongodb.indexes.search.E005",
                )
            )
        if num_arrayfields == 0:
            errors.append(
                Error(
                    "VectorSearchIndex requires at least one ArrayField to store vector data.",
                    obj=model,
                    id="mongodb.indexes.search.E006",
                    hint="If you want to perform search operations without vectors, "
                    "use SearchIndex instead.",
                )
            )
        return errors

    def deconstruct(self):
        path, args, kwargs = super().deconstruct()
        kwargs["similarities"] = self.similarities
        return path, args, kwargs

    def get_pymongo_index_model(self, model, schema_editor, field=None, column_prefix=""):
        if not schema_editor.connection.features.supports_search:
            return None
        similarities = (
            itertools.cycle([self.similarities])
            if not self._multiple_similarities
            else iter(self.similarities)
        )
        fields = []
        for field_name, _ in self.fields_orders:
            field_ = model._meta.get_field(field_name)
            field_path = column_prefix + model._meta.get_field(field_name).column
            mappings = {"path": field_path}
            if isinstance(field_, ArrayField):
                mappings.update(
                    {
                        "type": "vector",
                        "numDimensions": int(field_.size),
                        "similarity": next(similarities),
                    }
                )
            else:
                mappings["type"] = "filter"
            fields.append(mappings)
        return SearchIndexModel(definition={"fields": fields}, name=self.name, type="vectorSearch")


def register_indexes():
    BuiltinLookup.as_mql_idx = builtin_lookup_idx
    Index._get_condition_mql = _get_condition_mql
    Index.get_pymongo_index_model = get_pymongo_index_model
    WhereNode.as_mql_idx = where_node_idx
