# coding: utf-8
# OceanBase Deploy.
# Copyright (C) 2021 OceanBase
#
# This file is part of OceanBase Deploy.
#
# OceanBase Deploy is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OceanBase Deploy is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with OceanBase Deploy.  If not, see <https://www.gnu.org/licenses/>.


from __future__ import absolute_import, division, print_function


import re
import os
import time
try:
    import subprocess32 as subprocess
except:
    import subprocess
from ssh import LocalClient


stdio = None


def parse_size(size):
    _bytes = 0
    if not isinstance(size, str) or size.isdigit():
        _bytes = int(size)
    else:
        units = {"B": 1, "K": 1<<10, "M": 1<<20, "G": 1<<30, "T": 1<<40}
        match = re.match(r'([1-9][0-9]*)\s*([B,K,M,G,T])', size.upper())
        _bytes = int(match.group(1)) * units[match.group(2)]
    return _bytes


def exec_cmd(cmd):
    stdio.verbose('execute: %s' % cmd)
    process = subprocess.Popen(cmd, shell=True, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    while process.poll() is None:
        line = process.stdout.readline()
        line = line.strip()
        if line:
            stdio.print(line.decode("utf8", 'ignore'))
    return process.returncode == 0


def run_test(plugin_context, db, cursor, *args, **kwargs):
    def get_option(key, default=''):
        value = getattr(options, key, default)
        if value is None:
            value = default
        return value
    def execute(cursor, query, args=None):
        msg = query % tuple(args) if args is not None else query
        stdio.verbose('execute sql: %s' % msg)
        stdio.verbose("query: %s. args: %s" % (query, args))
        try:
            cursor.execute(query, args)
            return cursor.fetchone()
        except:
            msg = 'execute sql exception: %s' % msg
            stdio.exception(msg)
            raise Exception(msg)

    def local_execute_command(command, env=None, timeout=None):
        return LocalClient.execute_command(command, env, timeout, stdio)

    global stdio
    cluster_config = plugin_context.cluster_config
    stdio = plugin_context.stdio
    clients = plugin_context.clients
    options = plugin_context.options

    not_test_only = not get_option('test_only')

    host = get_option('host', '127.0.0.1')
    port = get_option('port', 2881)
    mysql_db = get_option('database', 'test')
    user = get_option('user', 'root')
    tenant_name = get_option('tenant', 'test')
    password = get_option('password', '')
    ddl_path = get_option('ddl_path')
    tbl_path = get_option('tbl_path')
    sql_path = get_option('sql_path')
    tmp_dir = get_option('tmp_dir')
    obclient_bin = get_option('obclient_bin', 'obclient')

    sql_path = sorted(sql_path, key=lambda x: (len(x), x))

    max_cpu = kwargs.get('max_cpu', 2)
    tenant_id = kwargs.get('tenant_id')
    unit_count = kwargs.get('unit_count', 0)
    cpu_total = 0

    if not_test_only:
        sql_cmd_prefix = '%s -h%s -P%s -u%s@%s %s -A' % (obclient_bin, host, port, user, tenant_name, ("-p'%s'" % password) if password else '')
        ret = local_execute_command('%s -e "%s"' % (sql_cmd_prefix, 'create database if not exists %s' % mysql_db))
        sql_cmd_prefix += ' -D %s' % mysql_db
        if not ret:
            stdio.error(ret.stderr)
            return
    else:
        sql_cmd_prefix = '%s -h%s -P%s -u%s@%s %s -D %s -A' % (obclient_bin, host, port, user, tenant_name, ("-p'%s'" % password) if password else '', mysql_db)

    ret = LocalClient.execute_command('%s -e "%s"' % (sql_cmd_prefix, 'select version();'), stdio=stdio)
    if not ret:
        stdio.error(ret.stderr)
        return

    for server in cluster_config.servers:
        client = clients[server]
        ret = client.execute_command("grep -e 'processor\s*:' /proc/cpuinfo | wc -l")
        if ret and ret.stdout.strip().isdigit():
            cpu_total += int(ret.stdout)
        else:
            server_config = cluster_config.get_server_conf(server)
            cpu_total += int(server_config.get('cpu_count', 0))
    try:
        sql = "select value from oceanbase.__all_virtual_sys_variable where tenant_id = %d and name = 'secure_file_priv'" % tenant_id
        ret = execute(cursor, sql)['value']
        if ret is None:
            stdio.error('Access denied. Please set `secure_file_priv` to "".')
            return
        if ret:
            for path in tbl_path:
                if not path.startswith(ret):
                    stdio.error('Access denied. Please set `secure_file_priv` to "".')
                    return

        parallel_num = int(max_cpu * unit_count)
        
        if not_test_only:
            # 替换并发数
            stdio.start_loading('Format DDL')
            n_ddl_path = []
            for fp in ddl_path:
                _, fn = os.path.split(fp)
                nfp = os.path.join(tmp_dir, fn)
                ret = local_execute_command("sed %s -e 's/partitions cpu_num/partitions %d/' > %s" % (fp, cpu_total, nfp))
                if not ret:
                    raise Exception(ret.stderr)
                n_ddl_path.append(nfp)
            ddl_path = n_ddl_path
            stdio.stop_loading('succeed')

            stdio.start_loading('Create table')
            for path in ddl_path:
                path = os.path.abspath(path)
                stdio.verbose('load %s' % path)
                ret = local_execute_command('%s < %s' % (sql_cmd_prefix, path))
                if not ret:
                    raise Exception(ret.stderr)
            stdio.stop_loading('succeed')

            stdio.start_loading('Load data')
            for path in tbl_path:
                _, fn = os.path.split(path)
                stdio.verbose('load %s' % path)
                ret = local_execute_command("""%s -c -e "load data /*+ parallel(%s) */ infile '%s' into table %s fields terminated by '|';" """ % (sql_cmd_prefix, parallel_num, path, fn[:-4]))
                if not ret:
                    raise Exception(ret.stderr)
            stdio.stop_loading('succeed')

            merge_version = execute(cursor, "select value from oceanbase.__all_zone where name='frozen_version'")['value']
            stdio.start_loading('Merge')
            execute(cursor, 'alter system major freeze')
            sql = "select value from oceanbase.__all_zone where name='frozen_version' and value != %s" % merge_version
            while True:
                if execute(cursor, sql):
                    break
                time.sleep(1)
                
            while True:
                if not execute(cursor, """select * from  oceanbase.__all_zone 
                where name='last_merged_version'
                and value != (select value from oceanbase.__all_zone where name='frozen_version' limit 1)
                and zone in (select zone from  oceanbase.__all_zone where name='status' and info = 'ACTIVE')
                """):
                    break
                time.sleep(5)
            stdio.stop_loading('succeed')


        # 替换并发数
        stdio.start_loading('Format SQL')
        n_sql_path = []
        for fp in sql_path:
            _, fn = os.path.split(fp)
            nfp = os.path.join(tmp_dir, fn)
            ret = local_execute_command("sed %s -e 's/parallel(cpu_num)/parallel(%d)/' > %s" % (fp, cpu_total, nfp))
            if not ret:
                raise Exception(ret.stderr)
            n_sql_path.append(nfp)
        sql_path = n_sql_path
        stdio.stop_loading('succeed')

        #warmup预热
        stdio.start_loading('Warmup')
        times = 2
        for path in sql_path:
            _, fn = os.path.split(path)
            log_path = os.path.join(tmp_dir, '%s.log' % fn)
            ret = local_execute_command('echo source %s | %s -c > %s' % (path, sql_cmd_prefix, log_path))
            if not ret:
                raise Exception(ret.stderr)
        stdio.stop_loading('succeed')

        total_cost = 0
        for path in sql_path:
            start_time = time.time()
            _, fn = os.path.split(path)
            log_path = os.path.join(tmp_dir, '%s.log' % fn)
            stdio.print('[%s]: start %s' % (time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(start_time)), path))
            ret = local_execute_command('echo source %s | %s -c > %s' % (path, sql_cmd_prefix, log_path))
            end_time = time.time()
            cost = end_time - start_time
            total_cost += cost
            stdio.print('[%s]: end %s, cost %.1fs' % (time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(end_time)), path, cost))
            if not ret:
                raise Exception(ret.stderr)
        stdio.print('Total Cost: %.1fs' % total_cost)
        return plugin_context.return_true()
    except KeyboardInterrupt:
        stdio.stop_loading('fail')
    except Exception as e:
        stdio.stop_loading('fail')
        stdio.exception(str(e))