from django.db import NotSupportedError
from django.db.models.expressions import Func
from django.db.models.functions.comparison import Cast, Coalesce, Greatest, Least, NullIf
from django.db.models.functions.datetime import (
    Extract,
    ExtractDay,
    ExtractHour,
    ExtractIsoWeekDay,
    ExtractIsoYear,
    ExtractMinute,
    ExtractMonth,
    ExtractSecond,
    ExtractWeek,
    ExtractWeekDay,
    ExtractYear,
    Now,
    TruncBase,
    TruncDate,
    TruncTime,
)
from django.db.models.functions.math import Ceil, Cot, Degrees, Log, Power, Radians, Random, Round
from django.db.models.functions.text import (
    Concat,
    ConcatPair,
    Left,
    Length,
    Lower,
    LTrim,
    Replace,
    RTrim,
    StrIndex,
    Substr,
    Trim,
    Upper,
)

from .query_utils import process_lhs

MONGO_OPERATORS = {
    Ceil: "ceil",
    Coalesce: "ifNull",
    Degrees: "radiansToDegrees",
    Greatest: "max",
    Least: "min",
    Power: "pow",
    Radians: "degreesToRadians",
    Random: "rand",
}
EXTRACT_OPERATORS = {
    ExtractDay.lookup_name: "dayOfMonth",
    ExtractHour.lookup_name: "hour",
    ExtractIsoWeekDay.lookup_name: "isoDayOfWeek",
    ExtractIsoYear.lookup_name: "isoWeekYear",
    ExtractMinute.lookup_name: "minute",
    ExtractMonth.lookup_name: "month",
    ExtractSecond.lookup_name: "second",
    ExtractWeek.lookup_name: "isoWeek",
    ExtractWeekDay.lookup_name: "dayOfWeek",
    ExtractYear.lookup_name: "year",
}


def cast(self, compiler, connection):
    output_type = connection.data_types[self.output_field.get_internal_type()]
    lhs_mql = process_lhs(self, compiler, connection)[0]
    if max_length := self.output_field.max_length:
        lhs_mql = {"$substrCP": [lhs_mql, 0, max_length]}
    # Skip the conversion for "object" as it doesn't need to be transformed for
    # interpretation by JSONField, which can handle types including int,
    # object, or array.
    if output_type != "object":
        lhs_mql = {"$convert": {"input": lhs_mql, "to": output_type}}
    if decimal_places := getattr(self.output_field, "decimal_places", None):
        lhs_mql = {"$trunc": [lhs_mql, decimal_places]}
    return lhs_mql


def concat(self, compiler, connection):
    return self.get_source_expressions()[0].as_mql(compiler, connection)


def concat_pair(self, compiler, connection):
    # null on either side results in null for expression, wrap with coalesce.
    coalesced = self.coalesce()
    return super(ConcatPair, coalesced).as_mql(compiler, connection)


def cot(self, compiler, connection):
    lhs_mql = process_lhs(self, compiler, connection)
    return {"$divide": [1, {"$tan": lhs_mql}]}


def extract(self, compiler, connection):
    lhs_mql = process_lhs(self, compiler, connection)
    operator = EXTRACT_OPERATORS.get(self.lookup_name)
    if operator is None:
        raise NotSupportedError("%s is not supported." % self.__class__.__name__)
    if timezone := self.get_tzname():
        lhs_mql = {"date": lhs_mql, "timezone": timezone}
    return {f"${operator}": lhs_mql}


def func(self, compiler, connection):
    lhs_mql = process_lhs(self, compiler, connection)
    operator = MONGO_OPERATORS.get(self.__class__, self.function.lower())
    return {f"${operator}": lhs_mql}


def left(self, compiler, connection):
    return self.get_substr().as_mql(compiler, connection)


def length(self, compiler, connection):
    # Check for null first since $strLenCP only accepts strings.
    lhs_mql = process_lhs(self, compiler, connection)
    return {"$cond": {"if": {"$eq": [lhs_mql, None]}, "then": None, "else": {"$strLenCP": lhs_mql}}}


def log(self, compiler, connection):
    # This function is usually log(base, num) but on MongoDB it's log(num, base).
    clone = self.copy()
    clone.set_source_expressions(self.get_source_expressions()[::-1])
    return func(clone, compiler, connection)


def now(self, compiler, connection):  # noqa: ARG001
    return "$$NOW"


def null_if(self, compiler, connection):
    """Return None if expr1==expr2 else expr1."""
    expr1, expr2 = (expr.as_mql(compiler, connection) for expr in self.get_source_expressions())
    return {"$cond": {"if": {"$eq": [expr1, expr2]}, "then": None, "else": expr1}}


def preserve_null(operator):
    # If the argument is null, the function should return null, not
    # $toLower/Upper's behavior of returning an empty string.
    def wrapped(self, compiler, connection):
        lhs_mql = process_lhs(self, compiler, connection)
        return {
            "$cond": {
                "if": connection.mongo_operators["isnull"](lhs_mql, True),
                "then": None,
                "else": {f"${operator}": lhs_mql},
            }
        }

    return wrapped


def replace(self, compiler, connection):
    expression, text, replacement = process_lhs(self, compiler, connection)
    return {"$replaceAll": {"input": expression, "find": text, "replacement": replacement}}


def round_(self, compiler, connection):
    # Round needs its own function because it's a special case that inherits
    # from Transform but has two arguments.
    return {"$round": [expr.as_mql(compiler, connection) for expr in self.get_source_expressions()]}


def str_index(self, compiler, connection):
    lhs = process_lhs(self, compiler, connection)
    # StrIndex should be 0-indexed (not found) but it's -1-indexed on MongoDB.
    return {"$add": [{"$indexOfCP": lhs}, 1]}


def substr(self, compiler, connection):
    lhs = process_lhs(self, compiler, connection)
    # The starting index is zero-indexed on MongoDB rather than one-indexed.
    lhs[1] = {"$add": [lhs[1], -1]}
    # If no limit is specified, use the length of the string since $substrCP
    # requires one.
    if len(lhs) == 2:
        lhs.append({"$strLenCP": lhs[0]})
    return {"$substrCP": lhs}


def trim(operator):
    def wrapped(self, compiler, connection):
        lhs = process_lhs(self, compiler, connection)
        return {f"${operator}": {"input": lhs}}

    return wrapped


def trunc(self, compiler, connection):
    lhs_mql = process_lhs(self, compiler, connection)
    lhs_mql = {"date": lhs_mql, "unit": self.kind, "startOfWeek": "mon"}
    if timezone := self.get_tzname():
        lhs_mql["timezone"] = timezone
    return {"$dateTrunc": lhs_mql}


def trunc_date(self, compiler, connection):
    # Cast to date rather than truncate to date.
    lhs_mql = process_lhs(self, compiler, connection)
    tzname = self.get_tzname()
    if tzname and tzname != "UTC":
        raise NotSupportedError(f"TruncDate with tzinfo ({tzname}) isn't supported on MongoDB.")
    return {
        "$dateFromString": {
            "dateString": {
                "$concat": [
                    {"$dateToString": {"format": "%Y-%m-%d", "date": lhs_mql}},
                    # Dates are stored with time(0, 0), so by replacing any
                    # existing time component with that, the result of
                    # TruncDate can be compared to DateField.
                    "T00:00:00.000",
                ]
            },
        }
    }


def trunc_time(self, compiler, connection):
    lhs_mql = process_lhs(self, compiler, connection)
    return {
        "$dateFromString": {
            "dateString": {
                "$concat": [
                    # Times are stored with date(1, 1, 1)), so by
                    # replacing any existing date component with that, the
                    # result of TruncTime can be compared to TimeField.
                    "0001-01-01T",
                    {"$dateToString": {"format": "%H:%M:%S.%L", "date": lhs_mql}},
                ]
            }
        }
    }


def register_functions():
    Cast.as_mql = cast
    Concat.as_mql = concat
    ConcatPair.as_mql = concat_pair
    Cot.as_mql = cot
    Extract.as_mql = extract
    Func.as_mql = func
    Left.as_mql = left
    Length.as_mql = length
    Log.as_mql = log
    Lower.as_mql = preserve_null("toLower")
    LTrim.as_mql = trim("ltrim")
    Now.as_mql = now
    NullIf.as_mql = null_if
    Replace.as_mql = replace
    Round.as_mql = round_
    RTrim.as_mql = trim("rtrim")
    StrIndex.as_mql = str_index
    Substr.as_mql = substr
    Trim.as_mql = trim("trim")
    TruncBase.as_mql = trunc
    TruncDate.as_mql = trunc_date
    TruncTime.as_mql = trunc_time
    Upper.as_mql = preserve_null("toUpper")
