# coding: utf-8

from datetime import datetime, timedelta

from airflow.models import Variable
from airflow.hooks.base import BaseHook
from airflow.providers.mysql.hooks.mysql import MySqlHook
from airflow.providers.microsoft.mssql.hooks.mssql import MsSqlHook
from airflow.providers.apache.hive.hooks.hive import HiveServer2Hook
from airflow.operators.python import PythonOperator

from utils import dag_start_date, clickhouse, deployment_env
from utils.table_structure_checker import check_and_handle_structure_changes

DEFAULT_DAG_ARGS = {
    'owner': 'airflow',
    'depends_on_past': False,
    'start_date': dag_start_date(start_hour=5),
    'retries': 5,
    "retry_delay": timedelta(minutes=3),
    "mode": "reschedule"  # reschedule模式每次检查后释放worker资源
}

DATETIME_FMT = '%Y-%m-%d %H:%M:%S.%f'


class FetchMode:
    APPEND = 'append'  # check_column > last_value
    APPEND_INCLUDE_PREV = 'append_include_prev'  # check_column >= last_value
    APPEND_STR = 'append_str'  # cast(check_column as int) > last_value
    APPEND_STR_INCLUDE_PREV = 'append_str_include_prev'  # cast(check_column as int) >= last_value
    LAST_MODIFIED = 'last_modified'  # check_column >= last_value
    FULL = 'full'


class WriteMode:
    APPEND = 'append'
    OVERWRITE = 'overwrite'


class TableEngine:
    MERGE_TREE = "ReplicatedMergeTree"
    REPLACING_MERGE_TREE = "ReplicatedReplacingMergeTree"


class IntegrationEngine:
    JDBC_BRIDGE = 'jdbc_bridge'
    MYSQL_ENGINE = 'mysql_engine'


ALL_OPTIONS = {
    'src_conn_id': '数据源连接ID',
    'integration_engine': '数据集成引擎',
    # Jdbc-Bridge中定义的NamedSource，或MySQL Database Engine对应的数据库名
    'src_name_in_ch': '数据源连接在ClickHouse中的名称',
    'src_table': '源表名',
    'src_columns': '源表列，逗号分隔，可使用表达式，可选',
    'where_condition': '查询过滤条件，可选',
    # FetchMode
    'src_fetch_mode': '源表查询方式',
    'src_check_column': '源表增量查询列',
    'src_last_value': '源表上次查询值，last_modified方式时为时间戳，append方式时为整数',
    'dst_conn_id': '数据目标连接ID',
    'dst_table': '目标表名',
    'dst_schema': '目标表schema，可选',
    # WriteMode
    'dst_write_mode': '目标表写入方式',
    'dst_order_columns': '目标表排序列，逗号分隔',
    'dst_partition_expr': '目标表分区表达式',
    # TableEngine
    'dst_engine': '目标表引擎',
}
REQUIRED_OPTIONS = set(ALL_OPTIONS.keys()) - {'src_columns', 'src_check_column', 'src_last_value', 'where_condition',
                                              'dst_schema', 'dst_partition_expr'}


# 验证配置
def validate_init_conf(init_conf):
    if not isinstance(init_conf, dict):
        raise TypeError('init_conf应为dict')

    all_options = set(init_conf.keys())

    lack_options = REQUIRED_OPTIONS - all_options
    if lack_options:
        raise ValueError('缺少配置项: ' + ', '.join(lack_options))

    src_fetch_mode = init_conf['src_fetch_mode']
    if src_fetch_mode != FetchMode.FULL:
        if 'src_check_column' not in init_conf:
            raise ValueError('增量查询缺少配置项: src_check_column')
        if 'src_last_value' not in init_conf:
            raise ValueError('增量查询缺少配置项: src_last_value')

        if src_fetch_mode in (
                FetchMode.APPEND, FetchMode.APPEND_INCLUDE_PREV, FetchMode.APPEND_STR,
                FetchMode.APPEND_STR_INCLUDE_PREV):
            int(init_conf['src_last_value'])
        elif src_fetch_mode == FetchMode.LAST_MODIFIED:
            datetime.strptime(init_conf['src_last_value'], DATETIME_FMT)

    integration_engine = init_conf.get('integration_engine', '')
    if integration_engine not in (IntegrationEngine.JDBC_BRIDGE, IntegrationEngine.MYSQL_ENGINE):
        raise ValueError(f'不支持的integration_engine类型: {integration_engine}')


def quote_db_identifier(identifier, db_type=None):
    if db_type == 'mysql':
        return '`' + identifier + '`'
    elif db_type == 'mssql':
        return '[' + identifier + ']'
    elif db_type == 'hiveserver2':
        return '`' + identifier + '`'  # Hive 使用反引号
    else:
        return identifier


def cast_db_column_to_int(column, for_clickhouse=False):
    return f'toInt64OrNull({column})' if for_clickhouse else f'cast({column} as bigint)'


def get_latest_value(variable_name, default_conf,ods_db_name):
    # 检查表结构
    Variable.setdefault(variable_name, default=default_conf, deserialize_json=True)
    check_and_handle_structure_changes(ods_db_name, variable_name)
    conf = Variable.get(variable_name, deserialize_json=True)
    # 清理已保存Variable的内容
    copied = default_conf.copy()
    copied.pop('src_latest_value', None)
    copied.pop('src_last_value', None)
    # 更新配置信息
    conf.update(copied)
    # 去除无用项，但保留运行时产生的值和表结构缓存
    to_remove = []
    preserved_keys = ('src_latest_value', 'src_last_value', 'src_structure_cache', 'src_columns_order_cache', 'time_cache')
    for key in conf.keys():
        if key not in copied and key not in preserved_keys:
            to_remove.append(key)
    print(f'remove keys: {to_remove}')
    for key in to_remove:
        del conf[key]

    fetch_mode = conf.get('src_fetch_mode')

    if fetch_mode == FetchMode.FULL:
        return

    src_conn_id = conf['src_conn_id']
    src_conn = BaseHook.get_connection(src_conn_id)
    src_type = src_conn.conn_type
    print(f'src_type: {src_type}')

    src_hook = None
    if src_type == 'mysql':
        src_hook = MySqlHook(mysql_conn_id=src_conn_id)
    elif src_type == 'mssql':
        src_hook = MsSqlHook(mssql_conn_id=src_conn_id)
    elif src_type == 'hiveserver2':
        src_hook = HiveServer2Hook(hiveserver2_conn_id=src_conn_id)

    check_column = quote_db_identifier(conf['src_check_column'], src_type)
    src_schema = src_conn.schema
    src_table = conf['src_table']
    check_expression = cast_db_column_to_int(check_column) if fetch_mode in (
        FetchMode.APPEND_STR, FetchMode.APPEND_STR_INCLUDE_PREV) else check_column
    
    # Hive 表名格式：schema.table 或 table（如果 schema 为空）
    if src_schema:
        table_full_name = f'{src_schema}.{src_table}'
    else:
        table_full_name = src_table
    
    latest_value_query = f'select max({check_expression}) as latest_value from {table_full_name} t'
    print(f'latest_value_query: {latest_value_query}')

    res = src_hook.get_records(sql=latest_value_query)
    latest_value = res[0][0]
    if latest_value is not None:
        if fetch_mode == FetchMode.LAST_MODIFIED:
            latest_value = latest_value.strftime(DATETIME_FMT)[:-3]  # SQLServer不支持微秒

        conf['src_latest_value'] = latest_value
        Variable.set(variable_name, conf, serialize_json=True)

        print(f'"{variable_name}".src_latest_value = {latest_value}')


def generate_import_sql(variable_name, default_conf):
    conf = Variable.get(variable_name, default_var=default_conf, deserialize_json=True)

    using_mysql_engine = conf.get('integration_engine', IntegrationEngine.JDBC_BRIDGE) == IntegrationEngine.MYSQL_ENGINE

    src_conn = conf['src_conn_id']
    src_type = "Custom"
    try:
        src_conn = BaseHook.get_connection(conf['src_conn_id'])
        src_type = src_conn.conn_type
    except Exception as e:
        print(e)

    src_table = conf['src_table']
    src_columns = conf.get('src_columns', '*')

    filter_conditions = []
    where_condition = conf.get('where_condition', None)
    if where_condition:
        filter_conditions.append(f'({where_condition})')

    fetch_mode = conf['src_fetch_mode']
    if fetch_mode != FetchMode.FULL:
        check_column = quote_db_identifier(conf['src_check_column'], src_type)
        last_value = conf['src_last_value']
        latest_value = conf.get('src_latest_value', None)

        if fetch_mode in (
                FetchMode.APPEND, FetchMode.APPEND_INCLUDE_PREV, FetchMode.APPEND_STR,
                FetchMode.APPEND_STR_INCLUDE_PREV):
            op = '>' if fetch_mode in (FetchMode.APPEND, FetchMode.APPEND_STR) else '>='
            check_expression = cast_db_column_to_int(check_column, for_clickhouse=using_mysql_engine) if fetch_mode in (
                FetchMode.APPEND_STR, FetchMode.APPEND_STR_INCLUDE_PREV) else check_column
            filter_conditions.append(f'{check_expression} {op} {last_value}')
            if latest_value is not None:
                filter_conditions.append(f'{check_expression} <= {latest_value}')
        elif fetch_mode == FetchMode.LAST_MODIFIED:
            sql_time_format = "'{}'" if using_mysql_engine else "''{}''"

            filter_conditions.append(f'{check_column} >= {sql_time_format.format(last_value)}')
            if latest_value is not None:
                filter_conditions.append(f'{check_column} <= {sql_time_format.format(latest_value)}')

    where = ''
    if filter_conditions:
        where = ' and '.join(filter_conditions)
        where = ' where ' + where

    src_name_in_ch = conf['src_name_in_ch']
    if isinstance(src_conn, str):
        src_schema = None
    else:
        src_schema = src_conn.schema
    dst_table = conf['dst_table']
    dst_schema = conf.get('dst_schema', '')
    dst_order_columns = conf['dst_order_columns']
    dst_engine = conf['dst_engine']
    dst_write_mode = conf['dst_write_mode']

    dst_partition_expr = conf.get('dst_partition_expr', '')
    if dst_partition_expr:
        dst_partition_expr = f' partition by {dst_partition_expr}'

    if src_type == 'hiveserver2' and ('sensors' in conf['src_conn_id'].lower()):
        sensors_project = 'production' if deployment_env().upper() == 'PROD' else 'default'
        src_pragma = f' /*SA({sensors_project})*/ /*+remapping_on*/'
    else:
        src_pragma = ''

    res_sql = []

    # 1. create dst table
    create_table_sql = ''
    if using_mysql_engine:
        create_table_sql = f'''create table if not exists {dst_table} on cluster default
        engine={dst_engine}{dst_partition_expr} order by ({dst_order_columns})
        as
        select {src_columns}
        from {src_name_in_ch}.{src_table}
        limit 0'''
    elif src_type == 'mssql':
        create_table_sql = f'''create table if not exists {dst_table} on cluster default
        engine={dst_engine}{dst_partition_expr} order by ({dst_order_columns})
        as
        select *
        from jdbc('{src_name_in_ch}','{dst_schema}','select top 0 {src_columns} from {src_schema}.{src_table}')'''
    elif src_type == 'mysql':
        create_table_sql = f'''create table if not exists {dst_table} on cluster default
        engine={dst_engine}{dst_partition_expr} order by ({dst_order_columns})
        as
        select *
        from jdbc('{src_name_in_ch}','{dst_schema}','select {src_columns} from {src_schema}.{src_table} limit 0')'''
    elif src_type == 'hiveserver2':
        create_table_sql = f'''create table if not exists {dst_table} on cluster default
        engine={dst_engine}{dst_partition_expr} order by ({dst_order_columns})
        as
        select *
        from jdbc('{src_name_in_ch}','{dst_schema}',
            'select {src_columns} from {src_schema}.{src_table} limit 0{src_pragma}')'''
    else:
        create_table_sql = f'''create table if not exists {dst_table} on cluster default
        engine={dst_engine}{dst_partition_expr} order by ({dst_order_columns})
        as
        select *
        from jdbc('{src_name_in_ch}','{dst_schema}',
            'select {src_columns} from {src_schema}.{src_table}{where}{src_pragma}')'''
    res_sql.append(create_table_sql)

    # 2. truncate dst table if overwrite
    if dst_write_mode == WriteMode.OVERWRITE:
        truncate_sql = f"truncate table {dst_table} on cluster default"
        res_sql.append(truncate_sql)

    # 3. insert new data
    if using_mysql_engine:
        insert_sql = f'''insert into {dst_table}
        select {src_columns}
        from {src_name_in_ch}.{src_table}{where}'''
    else:
        insert_sql = f'''insert into {dst_table}
        select *
        from jdbc('{src_name_in_ch}','{dst_schema}',
            'select {src_columns} from {src_schema}.{src_table}{where}{src_pragma}')'''

    res_sql.append(insert_sql)
    res_sql.append(f'optimize table {dst_table} on cluster default final')
    return res_sql


def update_last_value(variable_name,ods_db_name, modify_func=None):
    conf = Variable.get(variable_name, deserialize_json=True)

    fetch_mode = conf['src_fetch_mode']
    if fetch_mode != FetchMode.FULL:
        last_value = conf['src_last_value']
        latest_value = conf.get('src_latest_value', None)

        if latest_value is not None:
            if modify_func:
                latest_value = modify_func(latest_value)

            if latest_value > last_value:
                conf['src_last_value'] = latest_value
            del conf['src_latest_value']
            Variable.set(variable_name, conf, serialize_json=True)
    check_and_handle_structure_changes(ods_db_name, variable_name)

def create_tasks(dag, init_conf, import_pool=None, latest_value_modify_func=None):
    if 'integration_engine' not in init_conf:
        init_conf['integration_engine'] = IntegrationEngine.JDBC_BRIDGE

    validate_init_conf(init_conf)

    src_conn_id = init_conf['src_conn_id']
    src_table = init_conf['src_table']
    src_id = src_conn_id + '.' + src_table
    dst_conn_id = init_conf['dst_conn_id']

    variable_name = 'ods.' + src_id

    latest_task_id = 'latest.' + src_id
    ods_db_name=dst_conn_id
    latest_task = PythonOperator(
        dag=dag,
        task_id=latest_task_id,
        python_callable=get_latest_value,
        op_args=[variable_name, init_conf,ods_db_name],
    )

    import_task_id = 'import.' + src_id
    pool = import_pool if import_pool is not None else 'default_pool'
    import_sql = generate_import_sql(variable_name, init_conf)
    import_task = clickhouse.ClickHouseOperator(
        dag=dag,
        task_id=import_task_id,
        mysql_conn_id=dst_conn_id,
        sql=import_sql,
        pool=pool,
    )

    update_task_id = 'update.' + src_id
    update_task = PythonOperator(
        dag=dag,
        task_id=update_task_id,
        python_callable=update_last_value,
        op_args=[variable_name,ods_db_name, latest_value_modify_func],
    )

    return latest_task, import_task, update_task
