# coding: utf-8
import os
from contextlib import closing

from airflow.providers.common.sql.hooks.sql import return_single_query_results
from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator
from airflow.providers.mysql.hooks.mysql import MySqlHook
import json
import logging
from typing import Dict, Optional
from clickhouse_driver import Client
# 配置详细的日志格式
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - [%(filename)s:%(lineno)d] - %(funcName)s() - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)

logger = logging.getLogger(__name__)


class ClickHouseHook(MySqlHook):
    """Interact with ClickHouse via MySQL protocol.
    """

    def __init__(self, *args, **kwargs):
        super(ClickHouseHook, self).__init__(*args, **kwargs)

    def run(self, sql, autocommit=False, parameters=None, handler=None, split_statements=False, return_last=True):
        """
        Copy from DbApiHook.run, but do not execute commit
        """
        self.descriptions = []

        if isinstance(sql, str):
            if split_statements:
                sql_list = self.split_sql_string(sql)
            else:
                sql_list = [sql] if sql.strip() else []
        else:
            sql_list = sql

        if sql_list:
            self.log.debug("Executing following statements against DB: %s", sql_list)
        else:
            raise ValueError("List of SQL statements is empty")
        _last_result = None
        with closing(self.get_conn()) as conn:
            if self.supports_autocommit:
                self.set_autocommit(conn, autocommit)

            with closing(conn.cursor()) as cur:
                results = []
                for sql_statement in sql_list:
                    sql_statement_lower = sql_statement.lower()
                    if 'select' in sql_statement_lower and 'prefer_column_name_to_alias' not in sql_statement_lower:
                        sql_statement = sql_statement.strip().rstrip(';') + '\nSETTINGS prefer_column_name_to_alias = 0'

                    self._run_command(cur, sql_statement, parameters)

                    if handler is not None:
                        result = handler(cur)
                        if return_single_query_results(sql, return_last, split_statements):
                            _last_result = result
                            _last_description = cur.description
                        else:
                            results.append(result)
                            self.descriptions.append(cur.description)

            # if not self.get_autocommit(conn):
            #     conn.commit()

        if handler is None:
            return None
        if return_single_query_results(sql, return_last, split_statements):
            self.descriptions = [_last_description]
            return _last_result
        else:
            return results


class ClickHouseOperator(SQLExecuteQueryOperator):
    """
    Executes sql code in a specific ClickHouse database.
    """

    template_fields = ('sql', 'parameters')
    template_fields_renderers = {
        'sql': 'mysql',
        'parameters': 'json',
    }
    template_ext = ('.sql', '.json')
    ui_color = "#ededed"

    def __init__(self, *, mysql_conn_id: str = 'mysql_default', database: str | None = None, **kwargs) -> None:
        if database is not None:
            hook_params = kwargs.pop('hook_params', {})
            kwargs['hook_params'] = {'schema': database, **hook_params}

        super().__init__(conn_id=mysql_conn_id, **kwargs)

    def execute(self, context):
        self.log.info("Executing: %s", self.sql)
        hook = ClickHouseHook(mysql_conn_id=self.conn_id, schema=self.database)
        output = hook.run(
            sql=self.sql,
            autocommit=self.autocommit,
            parameters=self.parameters,
            handler=self.handler if self.do_xcom_push else None,
            split_statements=self.split_statements,
            return_last=self.return_last,
        )
        if return_single_query_results(self.sql, self.return_last, self.split_statements):
            return self._process_output([output], hook.descriptions)[-1]
        return self._process_output(output, hook.descriptions)


def get_clickhouse_hook(mysql_conn_id, schema=None):
    """
    获取ClickHouse Hook实例
    
    Args:
        mysql_conn_id: MySQL连接ID（ClickHouse通过MySQL协议访问）
        schema: 数据库名（schema）
    
    Returns:
        ClickHouseHook实例
    """
    return ClickHouseHook(mysql_conn_id=mysql_conn_id, schema=schema)

