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 load_cmpy_power(cids):
    sql = "SELECT sum(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):
    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-%%d')='{month_str}' GROUP BY " \
          f"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