# convert randgen test file generated by gen_fullstack_test_file.sh to full compact mode
import os
import errno
import sys


def write_force_full_compact_command(file, db_name, table_name):
    file.write("ti_ch> DBGInvoke query_mapped('manage table \$d.\$t flush', {}, {})\n"
               .format(db_name, table_name.lower()))
    file.write("ti_ch> DBGInvoke query_mapped('manage table \$d.\$t merge delta', {}, {})\n\n"
               .format(db_name, table_name.lower()))


def convert(src_path, dst_path):
    db_and_tables = []
    select_line_and_results = []
    drop_lines_after_select = []
    full_compact_mark = "# FULL COMPACT MODE"
    find_full_compact_mark = False
    find_select = False
    # memorize all lines after we find select stmt
    begin_select = False
    with open(src_path, "r") as src_file, open(dst_path, "w") as dst_file:
        lines = src_file.readlines()
        for line in lines:
            if "wait_table" in line:
                db_and_table_name = line.split("wait_table")[1].strip().split(" ")
                db_name = db_and_table_name[0].strip()
                table_name = db_and_table_name[1].strip()
                db_and_tables.append((db_name, table_name))

            # the following contents is generated by full compact mode, ignore the following lines
            if full_compact_mark in line:
                find_full_compact_mark = True

            # move drop statement to the end of the file
            if find_select and ("drop database" in line or "DROP DATABASE" in line):
                drop_lines_after_select.append(line)
                continue

            # rewrite all sql(except drop) after full compact mode mark
            if find_full_compact_mark:
                continue

            # do another select after full compact
            if "select" in line or "SELECT" in line:
                begin_select = True

            if begin_select:
                select_line_and_results.append(line)
                find_select = True

            dst_file.write(line)

        dst_file.write(full_compact_mark)
        dst_file.write("\n")
        for db_and_table_name in db_and_tables:
            write_force_full_compact_command(dst_file, db_and_table_name[0], db_and_table_name[1])
        for line in select_line_and_results:
            dst_file.write(line)
        for line in drop_lines_after_select:
            dst_file.write(line)
        return True


def try_convert_file(full_path):
    tmp_full_path = full_path + ".tmp"
    if convert(full_path, tmp_full_path):
        os.rename(tmp_full_path, full_path)
    else:
        os.remove(tmp_full_path)


def main():
    if len(sys.argv) != 2:
        print("usage: <bin> test_file_path/test_file_dir")
        sys.exit(1)

    test_path = sys.argv[1]
    if os.path.isdir(test_path):
        files = os.listdir(test_path)
        for name in files:
            print("processing file {}".format(name))
            if name.endswith(".tmp"):
                continue
            full_path = test_path + "/" + name
            try_convert_file(full_path)
    else:
        try_convert_file(test_path)


if __name__ == "__main__":
    main()