import json
import logging
import os
import threading
import typing

from elasticsearch import AsyncElasticsearch
import numpy

logger = logging.getLogger(__name__)


class DataDao:

    def __init__(self, es_client: AsyncElasticsearch):
        self.es_client = es_client

        self._lock = threading.RLock()

    async def search(self, tb_name: str,
                     embedding=typing.Optional[typing.Union[numpy.ndarray, list]],
                     filters=None,
                     from_=typing.Optional[int],
                     size=typing.Optional[int],
                     ):
        search_body = dict()

        from_ = from_ if isinstance(from_, int) and from_ >= 0 else 0
        size = size if isinstance(size, int) and size > 0 else 10
        search_body.update({
            'from': from_,
            'size': size,
        })

        return_embedding = os.getenv('RETURN_EMBEDDING', 'false').lower().strip()
        if return_embedding != 'true':
            search_body.update({
                "_source": {
                    "excludes": ["embedding"]
                }
            })

        condition_query = {
            "match_all": {}
        }
        if isinstance(filters, typing.Collection) and len(filters) > 0:
            tb_mapping = await self._get_mapping(tb_name)
            exist_properties = tb_mapping.get('mappings', {}).get('properties', {})

            clauses = []
            for filter in filters:
                k = filter.name
                v = filter.value
                opt = filter.opt
                if k not in exist_properties:
                    continue
                clause = get_filter_clause(k, v, opt, exist_properties[k])
                if clause:
                    clauses.append(clause)
            condition_query = {
                "bool": {
                    "filter": clauses
                }
            }
        if embedding is not None:
            use_knn = os.getenv('USE_KNN', 'false').lower().strip()
            if use_knn != 'true':
                query_body = {
                    "script_score": {
                        "query": condition_query,
                        "script": {
                            "source": """
                            // 1. 计算原始分数 (范围 [0, 2])
                            double rawScore = cosineSimilarity(params.query_vector, 'embedding') + 1.0;
                            
                            // 2. 归一化到 [0, 1]
                            double normalizedScore = rawScore / 2.0;
                            
                            // 3. 强制截断：确保最小值为 0.0，最大值为 1.0
                            // Math.max 防止出现负数，Math.min 防止出现 > 1.0 的数
                            double clampedScore = Math.max(0.0, Math.min(1.0, normalizedScore));
                            
                            // 4. 四舍五入到 4 位小数
                            // 原理：乘以 10000 -> 四舍五入取整 -> 除以 10000
                            // clampedScore = Math.round(clampedScore * 10000.0) / 10000.0;
                            return clampedScore;
                            """,
                            "params": {
                                "query_vector": embedding.tolist() if isinstance(embedding, numpy.ndarray) else list(embedding)
                            }
                        }
                    }
                }
                search_body.update({
                    "query": query_body
                })
            else:
                search_body.update({
                    "knn": {
                        "field": "embedding",
                        "query_vector": embedding.tolist() if isinstance(embedding, numpy.ndarray) else list(embedding),
                        "k": from_ + size,
                        "num_candidates": (from_ + size) * 20,
                        "filter": condition_query,
                    }
                })
        else:
            query_body = condition_query
            search_body.update({
                "query": query_body
            })

        logger.info(f'index {tb_name} search body: {json.dumps(search_body, ensure_ascii=False, indent=4)}')
        resp = await self.es_client.search(index=tb_name, body=search_body)
        logger.info(f'index {tb_name} search response: {resp}')

        body = resp.body
        hits = body.get('hits', {}).get('hits', [])
        result_data = [
            {
                'id': hit['_id'],
                'score': round(hit['_score'], 4),
                'kwargs': hit['_source'],
            }
            for hit in hits
        ]
        return result_data

    async def delete_by_pks(self, tb_name: str, pks: typing.List[str]):
        query_body = {
            "query": {
                "terms": {
                    "_id": pks
                }
            }
        }
        resp = await self.es_client.delete_by_query(index=tb_name, body=query_body)
        logger.info(f'index {tb_name} delete_by_query response: {resp}')

        body = resp.body
        # return body.get('deleted')

    async def _get_mapping(self, tb_name: str):
        resp = await self.es_client.indices.get_mapping(index=tb_name)
        logger.info(f'index {tb_name} get_mapping response: {resp}')
        return resp.body.get(tb_name, {})

    async def _tb_exist(self, tb_name, mapping=None):
        resp = await self.es_client.indices.exists(index=tb_name)
        result = resp.body
        if result and mapping is not None:
            tb_mapping = await self._get_mapping(tb_name)
            exist_properties = tb_mapping.get('mappings', {}).get('properties', {})

            expect_properties = mapping.get('mappings', {}).get('properties', {})
            if exist_properties and expect_properties:
                shared_keys = set(exist_properties.keys()).intersection(set(expect_properties.keys()))
                for k in shared_keys or []:
                    if exist_properties[k]['type'] != expect_properties[k]['type']:
                        raise Exception(
                            f'index {tb_name} `{k}` type not match, expect: {expect_properties[k]["type"]}, actual: {exist_properties[k]["type"]}')
                    else:
                        if exist_properties[k]['type'] == 'dense_vector':
                            if exist_properties[k].get('dims') != expect_properties[k].get('dims'):
                                raise Exception(
                                    f'index {tb_name} dims not match, expect: {expect_properties[k].get("dims")}, actual: {exist_properties[k].get("dims")}')
        return result

    async def _create_index(self, tb_name, mapping=typing.Optional[dict]) -> bool:
        result = await self._tb_exist(tb_name=tb_name, mapping=mapping)
        if result:
            return True
        with self._lock:
            result = await self._tb_exist(tb_name=tb_name, mapping=mapping)
            if result:
                return True
            resp = await self.es_client.indices.create(index=tb_name, body=mapping)
            logger.info(f'index {tb_name} create response: {resp}')
            return resp.body.get('acknowledged') is True

    async def upsert(self, tb_name, id, embedding, params: dict):
        doc = {
            'embedding': embedding,
            **params
        }

        auto_create_index = os.getenv('AUTO_CREATE_INDEX', 'true').lower().strip()
        if auto_create_index == 'true':
            mapping = generate_mapping(doc)
            # logger.info(f'Possible mapping: {json.dumps(mapping, indent=4)}')
            result = await self._create_index(tb_name=tb_name, mapping=mapping)

        if isinstance(embedding, numpy.ndarray):
            embedding = embedding.tolist()
        # elif isinstance(embedding, typing.Iterable):
        #     if isinstance(embedding, typing.Mapping):
        #         pass
        #     else:
        #         embedding = list(embedding)
        doc['embedding'] = embedding

        resp = await self.es_client.index(index=tb_name, id=id, body=doc)
        logger.info(f'index {tb_name} index response: {resp}')

        body = resp.body
        assert body.get('result') in ['created', 'updated'], f'数据插入失败, id: {id}'
        return body.get('_id')


def get_filter_clause(k, v, opt, type_property):
    property_type = type_property.get('type')
    if property_type == 'text':
        property_fields = type_property.get('fields') or {}
        if 'keyword' in property_fields:
            k = f'{k}.keyword'

    opt = opt.lower()
    # "eq", "neq", "lt", "gt", "lte", "gte", "like", "in"
    if opt in ['eq', 'in']:
        v = list(v) if isinstance(v, typing.Collection) else [v]
        should_list = []
        if None in v or len(v) == 0:
            should_list.append({
                "bool": {
                    "must_not": [
                        {
                            "exists": {
                                "field": k
                            }
                        }
                    ]
                }
            })
        no_null_v = [item for item in v if item is not None]
        if len(no_null_v) > 0:
            should_list.append({
                "terms": {
                    k: no_null_v
                }
            })
        return should_list[0] if len(should_list) == 1 else {
            "bool": {
                "should": should_list,
                "minimum_should_match": 1
            }
        }
    elif opt == 'neq':
        must_list = []
        v = list(v) if isinstance(v, typing.Collection) else [v]

        if None in v or len(v) == 0:
            must_list.append({
                "bool": {
                    "must": [
                        {
                            "exists": {
                                "field": k
                            }
                        }
                    ]
                }
            })
        no_null_v = [item for item in v if item is not None]
        if len(no_null_v) > 0:
            must_list.append({
                "bool": {
                    "must_not": {
                        "terms": {
                            k: no_null_v
                        }
                    }
                }
            })
        return must_list[0] if len(must_list) == 1 else {
            "bool": {
                "must": must_list
            }
        }
    elif opt in ['lt', 'lte', 'gt', 'gte']:
        v = list(v) if isinstance(v, typing.Collection) else [v]
        if len(v) > 0 and v[0] is not None:
            return {
                "range": {
                    k: {opt: list(v)[0]}
                }
            }
        else:
            return {
                "range": {
                    k: {}
                }
            }
    elif opt == 'like':
        v = list(v) if isinstance(v, typing.Collection) else [v]
        if len(v) > 0 and v[0] is not None:
            return {
                "wildcard": {
                    k: f"*{list(v)[0]}*"
                }
            }
        else:
            return {
                "bool": {
                    "must_not": [
                        {
                            "exists": {
                                "field": k
                            }
                        }
                    ]
                }
            }
    else:
        raise Exception(f'opt {opt} not support')


def generate_mapping(doc):
    def _inner(v):
        if isinstance(v, numpy.ndarray):
            return 'dense_vector'
        elif isinstance(v, str):
            return 'text'
        elif isinstance(v, int):
            return 'long'
        elif isinstance(v, float):
            return 'float'
        elif isinstance(v, typing.Mapping):
            return None
        elif isinstance(v, typing.Collection):
            if len(v) == 0:
                return None
            else:
                return _inner(list(v)[0])
        else:
            return None

    properties = dict()
    for key, value in (doc or {}).items():
        key_type = _inner(value)
        if key_type:
            properties[key] = {
                'type': key_type,
            }
            if key_type == 'dense_vector':
                properties[key]['dims'] = len(value)
                properties[key].update({
                    'dims': len(value),
                    'index': True,
                    'similarity': 'cosine'
                })
            elif key_type == 'text':
                properties[key]['fields'] = {
                    "keyword": {
                        "type": "keyword",
                        # "ignore_above": 256
                    }
                }
    es_mapping = {
        "settings": {
            "number_of_replicas": int(os.getenv('ES_NUMBER_OF_REPLICAS', 0)),
        },
        "mappings": {
            "properties": properties
        }
    }
    return es_mapping