# Author: Junjun
# Date: 2025/9/9

import json

import requests
from elasticsearch import Elasticsearch
from fastapi import HTTPException

from apps.datasource.models.datasource import DatasourceConf


def get_es_connect(conf: DatasourceConf):
    es_client = Elasticsearch(
        [conf.host],  # ES address
        basic_auth=(conf.username, conf.password),
        verify_certs=False,
        compatibility_mode=True
    )
    return es_client


# get tables
def get_es_index(conf: DatasourceConf):
    es_client = get_es_connect(conf)
    indices = es_client.cat.indices(format="json")
    res = []
    if indices is not None:
        for idx in indices:
            index_name = idx.get('index')
            desc = ''
            # get mapping
            mapping = es_client.indices.get_mapping(index=index_name)
            mappings = mapping.get(index_name).get("mappings")
            if mappings.get('_meta'):
                desc = mappings.get('_meta').get('description')
            res.append((index_name, desc))
    return res


# get fields
def get_es_fields(conf: DatasourceConf, table_name: str):
    es_client = get_es_connect(conf)
    index_name = table_name
    mapping = es_client.indices.get_mapping(index=index_name)
    properties = mapping.get(index_name).get("mappings").get("properties")
    res = []
    if properties is not None:
        for field, config in properties.items():
            field_type = config.get("type")
            desc = ''
            if config.get("_meta"):
                desc = config.get("_meta").get('description')

            if field_type:
                res.append((field, field_type, desc))
            else:
                # object、nested...
                res.append((field, ','.join(list(config.keys())), desc))
    return res


def get_es_data(conf: DatasourceConf, sql: str, table_name: str):
    r = requests.post(f"{conf.host}/_sql/translate", json={"query": sql})
    if r.json().get('error'):
        print(json.dumps(r.json()))

    es_client = get_es_connect(conf)
    response = es_client.search(
        index=table_name,
        body=json.dumps(r.json())
    )

    # print(response)
    fields = get_es_fields(conf, table_name)
    res = []
    for hit in response.get('hits').get('hits'):
        item = []
        if 'fields' in hit:
            result = hit.get('fields')  # {'title': ['Python'], 'age': [30]}
            for field in fields:
                v = result.get(field[0])
                item.append(v[0]) if v else item.append(None)
            res.append(tuple(item))
            # print(hit['fields']['title'][0])
        # elif '_source' in hit:
        #     print(hit.get('_source'))
    return res, fields