import datetime

from aioch import Client
import pandas as pd


class CKDrive:
    client: Client = None

    async def execute(self, sql) -> dict:
        data, columns = await self.client.execute(sql, with_column_types=True, columnar=True)
        df = pd.DataFrame({col[0]: d for d, col in zip(data, columns)})
        return df.T.to_dict()

    async def query_dataframe(self, sql):
        data, columns = await self.client.execute(sql, with_column_types=True, columnar=True)
        df = pd.DataFrame({col[0]: d for d, col in zip(data, columns)})
        return df

    async def count(self, db: str, tb: str):
        sql = f'select count() as `count` from {db}.{tb}'
        res = await self.execute(sql)
        return res[0]['count']

    async def distinct_count(self, db: str, tb: str, field: str):
        sql = f'select count(distinct `{field}`)  as `count` from {db}.{tb}'
        res = await self.execute(sql)
        return res[0]['count']

    async def field_count(self, db: str, tb: str):
        sql = f"select count(name) as `count` from system.columns where database='{db}' and table='{tb}'"
        res = await self.execute(sql)
        return res[0]['count']

    async def distinct(self, db: str, tb: str, field: str):
        sql = f'select distinct `{field}` as v from {db}.{tb}'
        res = await self.query_dataframe(sql)
        return res['v'].to_list()

    async def yesterday_event_count(self, db: str):
        today = datetime.date.today()
        yesterday = today - datetime.timedelta(days=1)
        today_str = today.strftime('%Y-%m-%d %H:%M:%S')
        yesterday_str = yesterday.strftime('%Y-%m-%d %H:%M:%S')
        sql = f"select `#event_name` as event_name, count() as v from {db}.event where `#event_time`>='{yesterday_str}' and `#event_time`<'{today_str}' group by `#event_name`"
        df = await self.query_dataframe(sql)
        return df.set_index('event_name').T.to_dict()


ckdb = CKDrive()


def get_ck_db() -> CKDrive:
    return ckdb