__all__ = 'CK',

import pandas as pd
from datetime import datetime
from datetime import timedelta

from clickhouse_driver import Client
from pandas import DatetimeTZDtype


class CK(Client):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def get_one(self, db, tb, **where) -> dict:
        """
        注意 还原时区
        :param db:
        :param tb:
        :param where:
        :return:
        """
        sql = f"select * from {db}.{tb} where 1"
        for k, v in where.items():
            sql += f" and `{k}`='{v}'"
        sql += ' limit 1'
        data, columns = self.execute(sql, with_column_types=True)
        res = dict()
        if data:
            data = {k[0]: v for k, v in zip(columns, data[0])}
            for k, v in data.items():
                if isinstance(v, datetime):
                    res[k] = (v + timedelta(hours=data['#zone_offset'])).strftime('%Y-%m-%d %H:%M:%S')
                else:
                    res[k] = v
            return res

    def get_all(self, db, tb, where: str) -> dict:
        """
        注意 还原时区
        :param db:
        :param tb:
        :param where:
        :return:
        """
        sql = f"select * from {db}.{tb} where "
        sql += where
        data, columns = self.execute(sql, columnar=True, with_column_types=True)
        df = pd.DataFrame({col[0]: d for d, col in zip(data, columns)})
        tz = df['#zone_offset'].apply(lambda x: timedelta(hours=x))
        for t_type in df.select_dtypes(include=[DatetimeTZDtype]):
            df[t_type] = (df[t_type] + tz).apply(lambda x: x.strftime('%Y-%m-%d %H:%M:%S'))

        return df.T.to_dict()