# -*- coding: utf-8 -*-

from contextlib import closing

from airflow.hooks.base import BaseHook
from airflow.providers.microsoft.mssql.hooks.mssql import MsSqlHook
from airflow.providers.mysql.hooks.mysql import MySqlHook

from utils import mysql, today

CONN_DW_MISC = 'dw_misc'


def last_synced(dst_table, columns, ctl_table='sync_counterdb_ctl'):
    if isinstance(columns, (tuple, list)):
        columns = ', '.join(columns)
    columns += ', last_full_sync'

    dst_hook = MySqlHook(mysql_conn_id=CONN_DW_MISC)
    res = dst_hook.get_first(f'select {columns} from {ctl_table} where dst_table=%s', (dst_table,))

    # 上次增量同步位置, 上次全量同步时间
    if res:
        return res[:-1], res[-1]
    else:
        return None, None


def should_full_sync(last_pos, periodic_full_sync=False, last_full_sync=None):
    # 未同步过，全量
    if not last_pos:
        return True

    # 定期全量同步，且上次全量同步时间不在今天，全量
    if periodic_full_sync and (last_full_sync is None or last_full_sync < today()):
        return True

    # 增量
    return False


def update_synced(dst_table, synced_rows, columns, curr_pos, last_full_sync=None, use_prev_pos_if_none=False,
                  ctl_table='sync_counterdb_ctl'):
    if synced_rows <= 0 or not curr_pos:
        return

    columns = list(columns)
    curr_pos = list(curr_pos)

    sql_prev_pos = f"select {', '.join(columns)} from {ctl_table} where dst_table='{dst_table}'"

    columns.append('dst_table')
    if last_full_sync:
        columns.append('last_full_sync')
    placeholders = ['%s'] * len(columns)
    sql_insert = f"replace into {ctl_table} ({', '.join(columns)}) values ({', '.join(placeholders)})"

    dst_hook = MySqlHook(mysql_conn_id=CONN_DW_MISC)
    with closing(dst_hook.get_conn()) as conn:
        with closing(conn.cursor()) as cur:
            if use_prev_pos_if_none:
                cur.execute(sql_prev_pos)
                prev_pos = cur.fetchone()

                if prev_pos:
                    for idx, val in enumerate(prev_pos):
                        if curr_pos[idx] is None:
                            curr_pos[idx] = val

            curr_pos.append(dst_table)
            if last_full_sync:
                curr_pos.append(last_full_sync)

            curr_pos = tuple(curr_pos)
            cur.execute(sql_insert, curr_pos)
            conn.commit()


def sync_table(src_conn_id, src_table, src_columns, dst_conn_id, dst_table, dst_columns, where_clause=None,
               where_params=None, inc_columns=None, inc_columns_independent=True, truncate_before_insert=False,
               column_modifiers=None, column_generators=None):
    """
    将源表(MySQL/SQLServer)数据导入目标表(MySQL)

    :param src_conn_id 源数据库连接ID
    :param src_table 源表
    :param src_columns 从源表查询的列
    :param dst_conn_id 目标数据库连接ID
    :param dst_table 目标表
    :param dst_columns 目标表中与src_columns相对应的列
    :param where_clause 源表过滤子句
    :param where_params 源表过滤子句参数
    :param inc_columns 源表中作为增量标记的列
    :param inc_columns_independent 源表中作为增量标记的列是否相互独立。如果独立则针对每个列单独计算最大值；
        如果不独立则按照列的顺序依次取最大值，位置在前的列若取到更大的值，其后的列置空
    :param truncate_before_insert 在写入目标表前是否先truncate
    :param column_modifiers 对源数据更新的方法，格式为{'src_column_name'=>function(src_column_value)->new_value}
    :param column_generators 根据源数据生成新列的方法，格式为{'new_column_name'=>function(src_row)->new_value}
        或{('new_column_name1','new_column_name2'...)=>function(src_row)->(new_value1,new_value2...)}
    """
    print(
        f'sync start: src_conn_id={src_conn_id}, src_table={src_table}, src_columns={src_columns}, '
        f'dst_conn_id={dst_conn_id}, dst_table={dst_table}, dst_columns={dst_columns}, where_clause={where_clause}, '
        f'where_params={where_params}, inc_columns={inc_columns}, truncate_before_insert={truncate_before_insert}')

    # 需要查询的列
    if isinstance(src_columns, str):
        src_columns = src_columns.split(',')

    # 增量处理列
    if inc_columns:
        if isinstance(inc_columns, str):
            inc_columns = inc_columns.split(',')
        inc_columns = [src_columns.index(col) for col in inc_columns]
        inc_columns = list(zip(range(len(inc_columns)), inc_columns))  # [(在inc_columns中的顺序, 在src_columns中的顺序)]

    if column_modifiers:
        column_modifiers = {src_columns.index(col): modifier for col, modifier in column_modifiers.items()}

    # 目标列
    if isinstance(dst_columns, str):
        dst_columns = dst_columns.split(',')
    if len(src_columns) != len(dst_columns):
        raise ValueError('源和目标列数不相等')

    dst_columns = list(dst_columns)
    if column_generators:
        for new_columns in column_generators.keys():
            if isinstance(new_columns, str):
                new_columns = (new_columns,)
            for new_column in new_columns:
                if new_column in dst_columns:
                    raise ValueError(f'已存在目标列{new_column}')
                dst_columns.append(new_column)

    src_conn = BaseHook.get_connection(src_conn_id)
    if src_conn.conn_type == 'mysql':
        src_hook = MySqlHook(mysql_conn_id=src_conn_id)
    elif src_conn.conn_type == 'mssql':
        src_hook = MsSqlHook(mssql_conn_id=src_conn_id)
    else:
        raise ValueError(f'{src_conn_id}不能作为源')

    dst_conn = BaseHook.get_connection(dst_conn_id)
    if dst_conn.conn_type == 'mysql':
        dst_hook = MySqlHook(mysql_conn_id=dst_conn_id)
    else:
        raise ValueError(f'{dst_conn_id}不能作为目标')

    select_columns = ', '.join(src_columns)
    if where_clause:
        where_clause = ' where ' + where_clause
    else:
        where_clause = ''

    query = f'select {select_columns} from {src_table} t{where_clause}'
    print(query)

    src_cursor = src_hook.get_cursor()
    if where_params:
        src_cursor.execute(query, where_params)
    else:
        src_cursor.execute(query)

    if truncate_before_insert:
        dst_cursor = dst_hook.get_cursor()
        dst_cursor.execute(f'truncate table {dst_table}')

    row_count = 0
    # 增量处理列上的最大值
    inc_max = None
    if inc_columns:
        inc_max = [None for _ in inc_columns]

    with closing(dst_hook.get_conn()) as conn:
        while True:
            rows = src_cursor.fetchmany(10000)
            if not rows:
                break

            if inc_columns:
                if inc_columns_independent:
                    for row in rows:
                        for inc_idx, row_idx in inc_columns:
                            if inc_max[inc_idx] is None or (
                                    row[row_idx] is not None and row[row_idx] > inc_max[inc_idx]):
                                inc_max[inc_idx] = row[row_idx]
                else:
                    for row in rows:
                        update_max = False
                        for inc_idx, row_idx in inc_columns:
                            if inc_max[inc_idx] is None or (
                                    row[row_idx] is not None and row[row_idx] > inc_max[inc_idx]):
                                update_max = True
                                break
                            elif row[row_idx] is None or row[row_idx] < inc_max[inc_idx]:
                                break
                        if update_max:
                            for inc_idx, row_idx in inc_columns:
                                inc_max[inc_idx] = row[row_idx]

            if column_modifiers:
                modified_rows = []
                for row in rows:
                    new_row = list(row)
                    for mod_idx, modifier in column_modifiers.items():
                        new_row[mod_idx] = modifier(row[mod_idx])
                    modified_rows.append(tuple(new_row))
                rows = modified_rows

            if column_generators:
                generated_rows = []
                for row in rows:
                    new_row = list(row)
                    for generator in column_generators.values():
                        generated_values = generator(row)
                        if isinstance(generated_values, (tuple, list)):
                            new_row.extend(generated_values)
                        else:
                            new_row.append(generated_values)
                    generated_rows.append(new_row)
                rows = generated_rows

            mysql.insert_rows(conn, table=dst_table, rows=rows, target_fields=dst_columns, upsert=True,
                              batch_size=10000)

            row_count += len(rows)
            print(f'{row_count} rows synced')

    src_cursor.close()

    print(f'sync done: src_conn_id={src_conn_id}, src_table={src_table}, dst_conn_id={dst_conn_id}, rows={row_count}')

    return row_count, inc_max


def valid_multi_src_config(src_tables):
    if not isinstance(src_tables, dict):
        raise ValueError('src_tables应是dict')

    fetch_col_num = set()
    for table_name, ctl_info in src_tables.items():
        if ((not isinstance(table_name, str)) or (not isinstance(ctl_info[0], tuple)) or (
                not isinstance(ctl_info[1], tuple)) or len(ctl_info[0]) != len(ctl_info[1])):
            raise ValueError('src_tables的格式应为{源表名:((源增量字段1,[源增量字段2]),(控制表字段1,[控制表字段2]),'
                             '(从源表获取字段1,[从源表获取字段2]))}')

        fetch_col_num.add(len(ctl_info[2]))

    if len(fetch_col_num) > 1:
        raise ValueError('src_tables中从源表获取字段数应一致')


def last_synced_multi_src(dst_table, src_tables, ctl_table='sync_counterdb_multi_src_ctl'):
    valid_multi_src_config(src_tables)

    hook = MySqlHook(mysql_conn_id=CONN_DW_MISC)

    res = {}
    for table_name, ctl_info in src_tables.items():
        cols = []
        for col in ctl_info[1]:
            if col == 'integer1':
                col = "ifnull(max(integer1),0) as integer1"
            elif col == 'integer2':
                col = "ifnull(max(integer2),0) as integer2"
            elif col == 'ts1':
                col = "ifnull(max(ts1),cast('2000-01-01' as datetime)) as ts1"
            elif col == 'ts2':
                col = "ifnull(max(ts2),cast('2000-01-01' as datetime)) as ts2"
            else:
                raise ValueError(f'控制表字段应为integer1/integer2/ts1/ts2，不能为{col}')

            cols.append(col)

        query = f"""select {', '.join(cols)}
        from {ctl_table}
        where
            dst_table='{dst_table}'
            and src_table='{table_name}'"""

        pos = hook.get_first(query)
        res[table_name] = pos

    return res


def current_pos_multi_src(src_conn_id, src_tables):
    valid_multi_src_config(src_tables)

    base_hook = BaseHook.get_connection(src_conn_id)
    if base_hook.conn_type == 'mysql':
        hook = MySqlHook(mysql_conn_id=src_conn_id)
    elif base_hook.conn_type == 'mssql':
        hook = MsSqlHook(mssql_conn_id=src_conn_id)
    else:
        raise ValueError(f'暂不支持连接类型{base_hook.conn_type}')

    res = {}
    with closing(hook.get_conn()) as conn:
        with closing(conn.cursor()) as cursor:
            for table_name, ctl_info in src_tables.items():
                cols = []
                for col in ctl_info[0]:
                    col = f"max({col}) as {col}"
                    cols.append(col)

                query = f"""select {', '.join(cols)} from {table_name}"""
                print(query)

                cursor.execute(query)
                pos = cursor.fetchone()

                res[table_name] = pos

    return res


def changed_from_last_pos(src_conn_id, table_name, fetch_cols, where_cols, last_pos, table_query=None):
    base_hook = BaseHook.get_connection(src_conn_id)
    if base_hook.conn_type == 'mysql':
        hook = MySqlHook(mysql_conn_id=src_conn_id)
    elif base_hook.conn_type == 'mssql':
        hook = MsSqlHook(mssql_conn_id=src_conn_id)
    else:
        raise ValueError(f'暂不支持连接类型{base_hook.conn_type}')

    if (not isinstance(where_cols, tuple)) or (not isinstance(last_pos, tuple)) or len(where_cols) != len(
            last_pos) or len(where_cols) <= 0:
        raise ValueError('where_cols和last_pos都应是tuple，且两者长度应一致')

    if isinstance(fetch_cols, (tuple, list)):
        fetch_cols = ', '.join(fetch_cols)

    where_clauses = []
    for where_col in where_cols:
        where_clause = f'{where_col}>=%s'
        where_clauses.append(where_clause)

    if table_query:
        query = f"select {fetch_cols} from ({table_query}) t where {' or '.join(where_clauses)}"
    else:
        query = f"select {fetch_cols} from {table_name} where {' or '.join(where_clauses)}"

    rows = hook.get_records(query, last_pos)
    print(f'{table_name}变动数: {len(rows)}')

    return rows


def update_synced_multi_src(dst_table, src_tables, current_pos, ctl_table='sync_counterdb_multi_src_ctl'):
    valid_multi_src_config(src_tables)

    hook = MySqlHook(mysql_conn_id=CONN_DW_MISC)

    for table_name, ctl_info in src_tables.items():
        row = [dst_table, table_name]
        for val in current_pos[table_name]:
            row.append(val)

        target_fields = ['dst_table', 'src_table']
        for col in ctl_info[1]:
            target_fields.append(col)

        hook.insert_rows(ctl_table, rows=[tuple(row)], target_fields=target_fields, replace=True)
