from pot_libs.es_util.es_utils import EsUtil
from pot_libs.mysql_util.mysql_util import MysqlUtil
from unify_api import constants


async def power_use_count(company_ids):
    query_body = {
        "query": {"bool": {"filter": [{"terms": {"cid": company_ids}},]}},
        "size": 0,
        "aggs": {"kwh": {"sum": {"field": "kwh"}}},
    }

    async with EsUtil() as es:
        es_result = await es.search_origin(body=query_body, index=constants.COMPANY_15MIN_POWER)

    total_power = round(es_result.get("aggregations", {}).get("kwh", {}).get("value") or 0, 2)
    return total_power


async def power_use_count_new15(cids):
    sql = "SELECT COUNT(kwh) kwh FROM `company_1day_power` where cid in %s"
    async with MysqlUtil() as conn:
        data = await conn.fetchone(sql, args=(cids,))
    total_power = round(data.get("kwh") or 0, 2)
    return total_power


async def inline_power_use_info(inline_ids, month_str):
    """
    获取进线用电信息
    :param inline_id:
    :param es_start_time:
    :param es_end_time:
    :return:
    """
    async with EsUtil() as es:
        filters = [
            {"terms": {"inlid": inline_ids}},
            {"term": {"month": f"{month_str}-01T00:00:00+08:00"}},
        ]
        query_body = {
            "query": {"bool": {"filter": filters}},
            "size": 10000,
        }
        es_result = await es.search_origin(body=query_body, index=constants.INLINE_1MONTH_POWER)

        power_infos = es_result["hits"]["hits"]

        inline_power_info_map = {
            i["_source"]["inlid"]: {
                "kwh": i["_source"]["kwh"],
                "charge": i["_source"]["charge"],
                "p": i["_source"]["p"],
            }
            for i in power_infos
        }
        return inline_power_info_map


async def inline_power_use_info_new15(inline_ids, month_str):
    sql = "SELECT inlid, sum(kwh) kwh, sum(charge) charge, sum(p) p FROM " \
          "`inline_1day_power` where inlid in %s  and " \
          f"DATE_FORMAT(create_time, '%%Y-%%m')='{month_str}' GROUP BY inlid"
    async with MysqlUtil() as conn:
        datas = await conn.fetchall(sql, args=(inline_ids, ))
    inline_power_info_map = {
        i["inlid"]: {
            "kwh": i["kwh"],
            "charge": i["charge"],
            "p": i["p"],
        }
        for i in datas
    }
    return inline_power_info_map
