#!/usr/bin/python
# ****************************************************************#
# ScriptName: restore_from_downloads
# Author: @alibaba-inc.com
# Function:
# ***************************************************************#

import os
import sys

STRUCTURE_FILE_NAME = "structure.sql"
DATA_PATH_PREFIX = "data"


def create_database(db_host, db_port, db_user, db_pass, create_stmt_file):
    cmd = "mysql -h" + db_host + " -P" + db_port + " -u" + db_user + " -p" + db_pass + " <" + create_stmt_file
    if os.system(cmd) != 0:
        print("[ERROR]: execute SQL failed. command: " + cmd)
        exit(1)


def create_table(db_host, db_port, db_user, db_pass, db_name, create_stmt_file):
    cmd = "mysql -h" + db_host + " -P" + db_port + " -u" + db_user + " -p" + db_pass + " -D" + db_name + " <" + create_stmt_file
    if os.system(cmd) != 0:
        print("[ERROR]: execute SQL failed. command: " + cmd)
        exit(1)


def import_file_csv(db_host, db_port, db_user, db_pass, csv_file, db, table):
    load_cmd = "load data local infile \"" + csv_file + "\" into table `" + db + "`.`" + table + "` character set utf8mb4 FIELDS TERMINATED BY \",\" enclosed by \"\\\"\""

    cmd = "mysql -h" + db_host + " -P" + db_port + " -u" + db_user + " -p" + db_pass + " -e '" + load_cmd + "'"
    print("[INFO]: trying to exec: " + cmd)
    if os.system(cmd) != 0:
        print("[ERROR]: execute SQL failed. command: " + cmd)
        exit(1)


def import_file_sql(db_host, db_port, db_user, db_pass, sql_file):
    cmd = "mysql -h" + db_host + " -P" + db_port + " -u" + db_user + " -p" + db_pass + " <" + sql_file
    print("[INFO]: trying to exec: " + cmd)
    if os.system(cmd) != 0:
        print("[ERROR]: execute SQL failed. command: " + cmd)
        exit(1)


def print_usage():
    print(
        "Usage: python ./restore_from_downloads.py [backupset_directory] [database_host] [database_port] [database_username] [database_password]")


if __name__ == '__main__':
    if len(sys.argv) != 6:
        print_usage()
        exit()

    root_dir = os.path.abspath(sys.argv[1])
    db_host = sys.argv[2]
    db_port = sys.argv[3]
    db_user = sys.argv[4]
    db_pass = sys.argv[5]
    print("[INFO]: restore data from " + root_dir + " to " + db_host + ":" + db_port)

    db_dirs = os.listdir(root_dir)
    for db_dir in db_dirs:
        dir_path = os.path.join(root_dir, db_dir)
        if not os.path.isdir(dir_path):
            continue
        db_structure_file = os.path.join(dir_path, STRUCTURE_FILE_NAME)
        create_database(db_host, db_port, db_user, db_pass, db_structure_file)
        print("[INFO]: restore structure database: " + db_dir + " ends")

        table_dirs = os.listdir(dir_path)
        for table_dir in table_dirs:
            table_dir_path = os.path.join(dir_path, table_dir)
            if not os.path.isdir(table_dir_path):
                continue
            table_structure_file = os.path.join(table_dir_path, STRUCTURE_FILE_NAME)
            create_table(db_host, db_port, db_user, db_pass, db_dir, table_structure_file)
            print("[INFO]: restore structure table: " + table_dir + " ends")

            table_data_dir_path = os.path.join(table_dir_path, DATA_PATH_PREFIX)
            if not os.path.isdir(table_data_dir_path):
                continue
            if len(os.listdir(table_data_dir_path)) == 0:
                print("[INFO]: table " + db_dir + "." + table_dir + " is empty, nothing to do")
                continue

            files_format = os.listdir(table_data_dir_path)[0].split(".")[-1]
            if files_format == "csv":
                csv_files = os.listdir(table_data_dir_path)
                csv_count = 0
                for csv_file in csv_files:
                    csv_file_path = os.path.join(table_data_dir_path, csv_file)
                    file_size = os.path.getsize(csv_file_path)
                    if file_size > 0:
                        import_file_csv(db_host, db_port, db_user, db_pass, csv_file_path, db_dir, table_dir)
                        csv_count = csv_count + 1
                        print("[INFO]: restore data [" + str(csv_count) + "/" + str(
                            len(csv_files)) + "] of table " + db_dir + "." + table_dir)
            elif files_format == "sql":
                sql_files = os.listdir(table_data_dir_path)
                sql_count = 0
                for sql_file in sql_files:
                    sql_file_path = os.path.join(table_data_dir_path, sql_file)
                    file_size = os.path.getsize(sql_file_path)
                    if file_size > 0:
                        import_file_sql(db_host, db_port, db_user, db_pass, sql_file_path)
                        sql_count = sql_count + 1
                        print("[INFO]: restore data [" + str(sql_count) + "/" + str(
                            len(sql_files)) + "] of table " + db_dir + "." + table_dir)