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.execute(sql)
        return res[0]['v']


ckdb = CKDrive()


def get_ck_db() -> CKDrive:
    return ckdb