# -*- coding: utf-8 -*-
#
import datetime
import inspect

import sys

if sys.version_info.major == 3 and sys.version_info.minor >= 10:
    from collections.abc import Iterable
else:
    from collections import Iterable
from functools import reduce, partial
from itertools import groupby
from uuid import UUID

from django.utils.translation import gettext_lazy as _
from django.db.models import QuerySet as DJQuerySet
from django.db.models import Q
from elasticsearch7 import Elasticsearch
from elasticsearch7.helpers import bulk
from elasticsearch7.exceptions import RequestError, SSLError
from elasticsearch7.exceptions import NotFoundError as NotFoundError7
from elasticsearch7.exceptions import UnsupportedProductError as UnsupportedProductError7

from elasticsearch8.exceptions import NotFoundError as NotFoundError8
from elasticsearch8.exceptions import BadRequestError
from elasticsearch8.exceptions import UnsupportedProductError as UnsupportedProductError8

from common.utils.common import lazyproperty
from common.utils import get_logger
from common.utils.timezone import local_now_date_display
from common.exceptions import JMSException

logger = get_logger(__file__)

UNSUPPORTED_PRODUCT_ERRORS = (
    UnsupportedProductError7,
    UnsupportedProductError8,
)


class InvalidElasticsearch(JMSException):
    default_code = 'invalid_elasticsearch'
    default_detail = _('Invalid elasticsearch config')


class InvalidElasticsearchProduct(JMSException):
    default_code = 'invalid_elasticsearch_product'
    default_detail = _('The server is not Elasticsearch or this product is not supported')


class NotSupportElasticsearch8(JMSException):
    default_code = 'not_support_elasticsearch8'
    default_detail = _('Not Support Elasticsearch8')


class InvalidElasticsearchSSL(JMSException):
    default_code = 'invalid_elasticsearch_SSL'
    default_detail = _(
        'Connection failed: Self-signed certificate used. Please check server certificate configuration')


class ESClient(object):

    def __new__(cls, *args, **kwargs):
        version = get_es_client_version(**kwargs)
        if version == 6:
            return ESClientV6(*args, **kwargs)
        if version == 7:
            return ESClientV7(*args, **kwargs)
        elif version == 8:
            return ESClientV8(*args, **kwargs)
        raise ValueError('Unsupported ES_VERSION %r' % version)


class ESClientBase(object):
    @classmethod
    def get_properties(cls, data, index):
        return data[index]['mappings']['properties']

    @classmethod
    def get_mapping(cls, properties):
        return {'mappings': {'properties': properties}}


class ESClientV7(ESClientBase):
    def __init__(self, *args, **kwargs):
        from elasticsearch7 import Elasticsearch
        self.es = Elasticsearch(*args, **kwargs)

    @classmethod
    def get_sort(cls, field, direction):
        return f'{field}:{direction}'

    @classmethod
    def get_sorts(cls, sorts: list):
        return ','.join(sorts)


class ESClientV6(ESClientV7):

    @classmethod
    def get_properties(cls, data, index):
        return data[index]['mappings']['data']['properties']

    @classmethod
    def get_mapping(cls, properties):
        return {'mappings': {'data': {'properties': properties}}}


class ESClientV8(ESClientBase):
    def __init__(self, *args, **kwargs):
        from elasticsearch8 import Elasticsearch
        self.es = Elasticsearch(*args, **kwargs)

    @classmethod
    def get_sort(cls, field, direction):
        return {field: {'order': direction}}

    @classmethod
    def get_sorts(cls, sorts: list):
        return sorts


def raise_invalid_elasticsearch(exc):
    logger.warning('Elasticsearch validation failed: %s', exc)
    if isinstance(exc, UNSUPPORTED_PRODUCT_ERRORS):
        raise InvalidElasticsearchProduct
    raise InvalidElasticsearch


def get_es_client_version(**kwargs):
    try:
        es = Elasticsearch(**kwargs)
        info = es.info()
        version = int(info['version']['number'].split('.')[0])
        return version
    except SSLError:
        raise InvalidElasticsearchSSL
    except Exception as e:
        raise_invalid_elasticsearch(e)


class ES(object):

    def __init__(
            self, config, properties, keyword_fields,
            exact_fields=None, fuzzy_fields=None, match_fields=None, **kwargs
    ):
        self.version = 7
        self.config = config
        hosts = self.config.get('HOSTS')
        kwargs = self.config.get('OTHER', {})

        ignore_verify_certs = kwargs.pop('IGNORE_VERIFY_CERTS', False)
        if ignore_verify_certs:
            kwargs['verify_certs'] = None
        try:
            self.client = ESClient(hosts=hosts, max_retries=0, **kwargs)
        except Exception as e:
            raise_invalid_elasticsearch(e)
        self.es = self.client.es
        self.index_prefix = self.config.get('INDEX') or 'jumpserver'
        self.is_index_by_date = bool(self.config.get('INDEX_BY_DATE', False))

        self.index = None
        self.query_index = None
        self.properties = properties
        self.exact_fields, self.match_fields, self.keyword_fields, self.fuzzy_fields = set(), set(), set(), set()

        if isinstance(keyword_fields, Iterable):
            self.keyword_fields.update(keyword_fields)
        if isinstance(exact_fields, Iterable):
            self.exact_fields.update(exact_fields)
        if isinstance(match_fields, Iterable):
            self.match_fields.update(match_fields)
        if isinstance(fuzzy_fields, Iterable):
            self.fuzzy_fields.update(fuzzy_fields)

        self.init_index()
        self.doc_type = self.config.get("DOC_TYPE") or '_doc'
        if self.is_new_index_type():
            self.doc_type = '_doc'
            self.exact_fields.update(self.keyword_fields)
        else:
            self.match_fields.update(self.keyword_fields)

    def init_index(self):
        if self.is_index_by_date:
            date = local_now_date_display()
            self.index = '%s-%s' % (self.index_prefix, date)
            self.query_index = '%s-alias' % self.index_prefix
        else:
            self.index = self.config.get("INDEX") or 'jumpserver'
            self.query_index = self.config.get("INDEX") or 'jumpserver'

    def is_new_index_type(self):
        if not self.ping(timeout=2):
            return False

        try:
            # 获取索引信息，如果没有定义，直接返回
            data = self.es.indices.get_mapping(index=self.index)
        except (NotFoundError8, NotFoundError7):
            return False

        try:
            properties = self.client.get_properties(data=data, index=self.index)
            for keyword in self.keyword_fields:
                if not properties[keyword]['type'] == 'keyword':
                    break
            else:
                return True
        except KeyError:
            return False

    def pre_use_check(self):
        if not self.ping(timeout=3):
            raise InvalidElasticsearch
        self._ensure_index_exists()

    def _ensure_index_exists(self):
        try:
            mappings = self.client.get_mapping(self.properties)

            if self.is_index_by_date:
                mappings['aliases'] = {
                    self.query_index: {}
                }

            try:
                self.es.indices.create(index=self.index, body=mappings)
            except (RequestError, BadRequestError) as e:
                if e.error == 'resource_already_exists_exception':
                    logger.warning(e)
                else:
                    logger.exception(e)
        except Exception as e:
            logger.error(e, exc_info=True)

    def make_data(self, data):
        return []

    def save(self, **kwargs):
        data = self.make_data(kwargs)
        return self.es.index(index=self.index, doc_type=self.doc_type, body=data)

    def bulk_save(self, command_set, raise_on_error=True):
        actions = []
        for command in command_set:
            data = dict(
                _index=self.index,
                _type=self.doc_type,
                _source=self.make_data(command),
            )
            actions.append(data)
        return bulk(self.es, actions, index=self.index, raise_on_error=raise_on_error)

    def get(self, query: dict):
        item = None
        data = self.filter(query, size=1)
        if len(data) >= 1:
            item = data[0]
        return item

    def filter(self, query: dict, from_=None, size=None, sort=None):
        try:
            data = self._filter(query, from_, size, sort)
        except Exception as e:
            logger.error('ES filter error: {}'.format(e))
            data = []
        return data

    def _filter(self, query: dict, from_=None, size=None, sort=None):
        body = self.get_query_body(**query)
        search_params = {
            'index': self.query_index,
            'body': body,
            'from_': from_,
            'size': size
        }
        if sort is not None:
            search_params['sort'] = sort
        logger.info('search_params: {}'.format(search_params))
        data = self.es.search(**search_params)

        source_data = []
        for item in data['hits']['hits']:
            if item:
                item['_source'].update({'es_id': item['_id']})
                source_data.append(item['_source'])

        return source_data

    def count(self, **query):
        try:
            body = self.get_query_body(**query)
            data = self.es.count(index=self.query_index, body=body)
            count = data["count"]
        except Exception as e:
            logger.error('ES count error: {}'.format(e))
            count = 0
        return count

    def __getattr__(self, item):
        return getattr(self.es, item)

    def all(self):
        """返回所有数据"""
        raise NotImplementedError("Not support")

    def ping(self, timeout=None):
        try:
            return self.es.ping(request_timeout=timeout)
        except Exception:
            return False

    @staticmethod
    def handler_time_field(data):
        datetime__gte = data.get('datetime__gte')
        datetime__lte = data.get('datetime__lte')
        datetime_range = {}

        if datetime__gte:
            if isinstance(datetime__gte, datetime.datetime):
                datetime__gte = datetime__gte.strftime('%Y-%m-%d %H:%M:%S')
            datetime_range['gte'] = datetime__gte
        if datetime__lte:
            if isinstance(datetime__lte, datetime.datetime):
                datetime__lte = datetime__lte.strftime('%Y-%m-%d %H:%M:%S')
            datetime_range['lte'] = datetime__lte
        return 'datetime', datetime_range

    @staticmethod
    def handle_exact_fields(exact):
        _filter = []
        for k, v in exact.items():
            query = 'term'
            if isinstance(v, list):
                query = 'terms'
            _filter.append({
                query: {k: v}
            })
        return _filter

    @staticmethod
    def handle_fuzzy_fields(exact):
        _filter = []
        for k, v in exact.items():
            _filter.append({'wildcard': {k: f'*{v}*'}})
        return _filter

    @staticmethod
    def is_keyword(props: dict, field: str) -> bool:
        return props.get(field, {}).get("type", "keyword") == "keyword"

    @staticmethod
    def is_long(props: dict, field: str) -> bool:
        return props.get(field, {}).get("type") == "long"

    def get_query_body(self, **kwargs):
        new_kwargs = {}
        for k, v in kwargs.items():
            if isinstance(v, UUID):
                v = str(v)
            if k == 'pk':
                k = 'id'
            if k.endswith('__in'):
                k = k.replace('__in', '')
            new_kwargs[k] = v
        kwargs = new_kwargs

        index_in_field = 'id__in'
        keyword_fields = self.keyword_fields
        exact_fields = self.exact_fields
        match_fields = self.match_fields
        fuzzy_fields = self.fuzzy_fields

        match = {}
        search = []
        exact = {}
        fuzzy = {}
        index = {}

        if index_in_field in kwargs:
            index['values'] = kwargs[index_in_field]

        mapping = self.es.indices.get_mapping(index=self.index)
        props = (
            mapping
            .get(self.index, {})
            .get('mappings', {})
            .get('properties', {})
        )

        common_keyword_able = exact_fields | keyword_fields

        for k, v in kwargs.items():
            if k in ("org_id", "session") and self.is_keyword(props, k):
                exact[k] = v

            elif self.is_long(props, k):
                exact[k] = v

            elif k in common_keyword_able:
                exact[f"{k}.keyword"] = v

            elif k in fuzzy_fields:
                fuzzy[f"{k}.keyword"] = v

            elif k in match_fields:
                match[k] = v

        args = kwargs.get('search', [])
        for item in args:
            for k, v in item.items():
                if k in match_fields:
                    search.append(item)

        # 处理时间
        time_field_name, time_range = self.handler_time_field(kwargs)

        # 处理组织
        should = []
        org_id = match.get('org_id')

        real_default_org_id = '00000000-0000-0000-0000-000000000002'
        root_org_id = '00000000-0000-0000-0000-000000000000'

        if org_id == root_org_id:
            match.pop('org_id')
        elif org_id in (real_default_org_id, ''):
            match.pop('org_id')
            should.append({
                'bool': {
                    'must_not': [
                        {
                            'wildcard': {'org_id': '*'}
                        }
                    ]}
            })
            should.append({'match': {'org_id': real_default_org_id}})

        # 构建 body
        body = {
            'query': {
                'bool': {
                    'must': [],
                    'should': should + [
                        {'match': {k: v}} for k, v in match.items()
                    ] + [
                                  {'match': item} for item in search
                              ],
                    'filter': self.handle_exact_fields(exact) +
                              self.handle_fuzzy_fields(fuzzy) +
                              [
                                  {
                                      'range': {
                                          time_field_name: time_range
                                      }
                                  }
                              ] + [
                                  {
                                      'ids': {k: v}
                                  } for k, v in index.items()
                              ]
                }
            },
        }
        return body


class QuerySet(DJQuerySet):
    custom = True
    default_days_ago = 7
    max_result_window = 10000

    def __init__(self, es_instance):
        self._method_calls = []
        self._slice = None  # (from_, size)
        self._storage = es_instance

        # 命令列表模糊搜索时报错
        super().__init__()

    @lazyproperty
    def _grouped_method_calls(self):
        _method_calls = {k: list(v) for k, v in groupby(self._method_calls, lambda x: x[0])}
        return _method_calls

    def _grouped_search_args(self, query):
        conditions = {}
        for q in query:
            for c in q.children:
                if isinstance(c, Q):
                    child = self._grouped_search_args(c)
                    [conditions.setdefault(k, []).extend(v) for k, v in child.items()]
                else:
                    conditions.setdefault(c[0], []).append(c[1])
        return conditions

    @lazyproperty
    def _filter_kwargs(self):
        _method_calls = self._grouped_method_calls
        filter_calls = _method_calls.get('filter')
        if not filter_calls:
            return {}
        names, multi_args, multi_kwargs = zip(*filter_calls)

        # input 输入
        multi_args = tuple(reduce(lambda x, y: x + y, (sub for sub in multi_args if sub), ()))
        args = self._grouped_search_args(multi_args)
        striped_args = [{k.replace('__icontains', ''): v} for k, values in args.items() for v in values]

        kwargs = reduce(lambda x, y: {**x, **y}, multi_kwargs, {})
        striped_kwargs = {'search': striped_args}
        for k, v in kwargs.items():
            k = k.replace('__exact', '')
            k = k.replace('__startswith', '')
            k = k.replace('__icontains', '')
            striped_kwargs[k] = v
        return striped_kwargs

    @lazyproperty
    def _sort(self):
        order_by = self._grouped_method_calls.get('order_by')
        _sorts = [self._storage.client.get_sort('_score', 'desc')]
        if order_by:
            for call in reversed(order_by):
                fields = call[1]
                if fields:
                    field = fields[-1]

                    if field.startswith('-'):
                        direction = 'desc'
                    else:
                        direction = 'asc'
                    field = field.lstrip('-+')
                    sort = self._storage.client.get_sort(field, direction)
                    _sorts.append(sort)
                    break
        sorts = self._storage.client.get_sorts(_sorts)
        return sorts

    def __execute(self):
        _filter_kwargs = self._filter_kwargs
        _sort = self._sort
        from_, size = self._slice or (None, None)
        data = self._storage.filter(_filter_kwargs, from_=from_, size=size, sort=_sort)
        return self.model.from_multi_dict(data)

    def __stage_method_call(self, item, *args, **kwargs):
        _clone = self.__clone()
        _clone._method_calls.append((item, args, kwargs))
        return _clone

    def __clone(self):
        uqs = QuerySet(self._storage)
        uqs._method_calls = self._method_calls.copy()
        uqs._slice = self._slice
        uqs.model = self.model
        return uqs

    def get(self, **kwargs):
        kwargs.update(self._filter_kwargs)
        return self._storage.get(kwargs)

    def count(self, limit_to_max_result_window=True):
        filter_kwargs = self._filter_kwargs
        count = self._storage.count(**filter_kwargs)
        if limit_to_max_result_window:
            count = min(count, self.max_result_window)
        return count

    def __getattribute__(self, item):
        if any((
                item.startswith('__'),
                item in QuerySet.__dict__,
        )):
            return object.__getattribute__(self, item)

        origin_attr = object.__getattribute__(self, item)
        if not inspect.ismethod(origin_attr):
            return origin_attr

        attr = partial(self.__stage_method_call, item)
        return attr

    def __getitem__(self, item):
        max_window = self.max_result_window
        if isinstance(item, slice):
            if self._slice is None:
                clone = self.__clone()
                from_ = item.start or 0
                if item.stop is None:
                    size = self.max_result_window - from_
                else:
                    size = item.stop - from_

                if from_ + size > max_window:
                    if from_ >= max_window:
                        from_ = max_window
                        size = 0
                    else:
                        size = max_window - from_
                clone._slice = (from_, size)
                return clone
        return self.__execute()[item]

    def __repr__(self):
        return self.__execute().__repr__()

    def __iter__(self):
        return iter(self.__execute())

    def __len__(self):
        return self.count()
