from pot_libs.es_util.es_utils import EsUtil
from unify_api import constants
from unify_api.constants import Importance, SDU_ALARM_LIST
from pot_libs.mysql_util.mysql_util import MysqlUtil


async def alarm_count(company_ids):
    query_body = {
        "query": {
            "bool": {
                "filter": [
                    {"terms": {"cid": company_ids}},
                    {
                        "terms": {
                            "importance": [
                                Importance.First.value,
                                Importance.Second.value,
                                Importance.Third.value,
                            ]
                        }
                    },
                ]
            }
        },
        "size": 0,
        "aggs": {"cid_alarm_aggs": {"terms": {"field": "cid", "size": 10000}}},
    }
    async with EsUtil() as es:
        es_result = await es.search_origin(body=query_body,
                                           index=constants.POINT_1MIN_EVENT)
    cid_alarm_buckets = (
        es_result.get("aggregations", {}).get("cid_alarm_aggs", {}).get(
            "buckets", [])
    )
    total_alarm_cnt = sum([i["doc_count"] for i in cid_alarm_buckets])
    return total_alarm_cnt


async def load_alarm_cnt_sdu(cids):
    sql = f"select count(1) doc_count from point_1min_event " \
          f"where cid in %s and importance in (1, 2, 3) " \
          f"and event_type in {tuple(SDU_ALARM_LIST)}"
    async with MysqlUtil() as conn:
        data = await conn.fetchone(sql, args=(cids,))
    return data["doc_count"] or 0
