from django.conf import settings
from django.contrib.postgres.fields.jsonb import KeyTransform
from django.contrib.postgres.search import SearchQuery as SearchQueryBase
from django.db.models import Func, Subquery, Case, When, Value, IntegerField, CharField
from django.db.models.functions import Lower, Substr
from django.contrib.postgres.fields import JSONField
from model_utils import Choices


class PythonConverterExpression(Func):
    """
    Класс для пост-обработки данных полученных из БД.
    Работает, как все подклассы Func, но еще
    принимает convert_function - callback, который выполнится
    после того, как данные получены.
    """
    template = '%(expressions)s'

    def __init__(self, *args, **kwargs):
        self.convert_function = kwargs.pop('convert_function', lambda x: x)
        super().__init__(*args, **kwargs)

    def convert_value(self, value, expression, connection):
        return self.convert_function(value)


class MapperExpression(Case):

    def __init__(self, field_name, mapping, **kwargs):
        cases = [When(**{field_name: k, 'then': Value(v)}) for k, v in mapping.items()]
        super().__init__(*cases, **kwargs)


class OrderExpression(Case):
    """
    Вычисляет порядковое значение для нечисловых полей.

    Высчитывает порядок для значения в поле `field_name`,
    отталкиваясь от сортированного списка `choices`.
    Если значение в поле `field_name` отсутствует в `choices`,
    по умолчанию берется -1.
    Можно переопределить через параметр `default`.
    """
    def __init__(self, field_name, choices, **kwargs):
        assert 'output_field' not in kwargs
        if isinstance(choices, Choices):
            choices = (i for i, _ in choices)
        cases = [When(**{field_name: c, 'then': Value(i)}) for i, c in enumerate(choices)]
        kwargs['output_field'] = IntegerField()
        kwargs.setdefault('default', Value(-1))
        super().__init__(*cases, **kwargs)


class ArraySubquery(Subquery):
    """
    Позволяет вернуть результат подзапроса в виде массива.
    """
    template = 'ARRAY(%(subquery)s)'


class RowToDict(Subquery):
    """
    Позволяет вернуть результат подзапроса в виде JSON-объекта.
    Важно! Подзапрос должен возвращать не более одной строки (лучше ограничить по [:1])
    """
    output_field = JSONField()
    template = '(SELECT to_jsonb(t) FROM (%(subquery)s) t)'


class RowsToList(Subquery):
    """
    Позволяет вернуть результат подзапроса в виде списка JSON-объектов.
    """
    output_field = JSONField()
    template = '(SELECT jsonb_agg(row_to_json(t)) FROM (%(subquery)s) t)'


class SearchQuery(SearchQueryBase):
    """
    Query для полнотекстового поиска с использованием нашего конфига `russian_unaccent`
    """
    search_config = settings.PG_TEXT_SEARCH_CONFIG

    def __init__(self, *args, **kwargs):
        kwargs.setdefault('config', self.search_config)
        super().__init__(*args, **kwargs)


class SearchQueryEn(SearchQuery):

    search_config = settings.PG_TEXT_SEARCH_CONFIG_EN


class JsonF(KeyTransform):
    """
    То же что и F, но умеет работать c json-полями
    """
    def __init__(self, field_name: str, key_name: str):
        keys = key_name.split('__')
        previous = field_name
        for i in range(len(keys) - 1):
            previous = KeyTransform(keys[i], previous)
        super().__init__(keys[-1], previous)


class Cardinality(Func):
    """
    Количество элементов ArrayField
    """
    function = 'CARDINALITY'
    arity = 1


class SubstringLevenshtein(Func):
    """
    Считаем расстояние Левенштейна по подстроке без дополнительных аргументов (стоимостей действий)
    """
    function = 'levenshtein'
    arity = 2

    def __init__(self, source_field, query, case_insensitive=False, **kwargs):
        # Берем столько символов, сколько в запросе
        source = Substr(source_field, 1, len(query))
        query = Value(query, output_field=CharField())
        if case_insensitive:
            source, query = Lower(source), Lower(query)

        super(SubstringLevenshtein, self).__init__(source, query, **kwargs)
