#!/usr/bin/python3

import argparse
import re
import sys

GLOBAL_TABLE = [
    "index_queue",
    "user_flag",
    "user_info",
    "group_info",
    "administrative",
    "config",
    "team_library",
    "many_team_library_has_many_user_info",
    "mobile_config",
]

GLOBAL_TABLE_SEQ = [i + "_id_seq" for i in GLOBAL_TABLE]


def HandlePgDumpSql(lines):
    output = ""

    multiline = False
    skip = False
    body = []
    sql_type = ""
    for line in lines:
        body.append(line)
        if re.match(r"SELECT pg_catalog.set_config", line):
            skip = True
        elif re.match(
            r"CREATE TRIGGER (create_user_schema_trigger|delete_user_schema_trigger|update_user_schema_trigger|create_delete_team_user_trigger)",
            line,
        ):
            skip = True
        elif re.match(r"(CREATE|ALTER) FUNCTION .*", line):
            skip = True
            sql_type = "FUNCTION"
        elif re.match(r"(CREATE|ALTER) SEQUENCE .*", line):
            res = re.findall(r"(?:CREATE|ALTER) SEQUENCE (?:public\.)?(\w+)", line)
            if len(res) and res[0] in GLOBAL_TABLE_SEQ:
                skip = True
        elif re.match(r"CREATE INDEX .*", line):
            res = re.findall(r"CREATE INDEX (\w+)", line)
            if len(res) and any([s for s in GLOBAL_TABLE if s in res[0]]):
                skip = True
        elif re.match(r"(CREATE|ALTER) TABLE .*", line):
            res = re.findall(
                r"(?:CREATE|ALTER) TABLE (?:ONLY\s+)?(?:public\.)?(\w+)(?:$|.*$)?", line
            )
            if len(res) and res[0] in GLOBAL_TABLE:
                skip = True

        if (not sql_type and line.endswith(";")) or (
            sql_type == "FUNCTION" and line.endswith("$$;")
        ):
            sql_type = ""
            if skip:
                body = []
                skip = False
                continue
            output += "\n".join(body)
            output += "\n"
            body = []

    return output


def HandlePgModelSql(lines):
    output = "\n".join(lines)
    return output


def main():
    parser = argparse.ArgumentParser(description="Process some integers.")
    parser.add_argument(
        "input", type=argparse.FileType("r"), help="sql file dumpped by pgmodelel"
    )
    parser.add_argument(
        "--outfile", "-o", nargs="?", type=argparse.FileType("w"), default=sys.stdout
    )
    parser.add_argument("--dump_source", required=True, choices=["postgres", "pgmodel"])

    args = parser.parse_args()
    dump_source = args.dump_source
    lines = []
    for line in args.input:
        if (
            line.isspace()
            or line.startswith("--")
            or line.startswith("COMMENT")
            or re.search("OWNER", line)
        ):
            continue
        if dump_source == "postgres" and line.startswith("SET"):
            continue
        if re.match(r"CREATE SCHEMA public", line):
            continue

        lines.append(line.rstrip())

    args.input.close()

    if dump_source == "pgmodel":
        args.outfile.write(HandlePgModelSql(lines))
    else:
        args.outfile.write(HandlePgDumpSql(lines))

    args.outfile.close()


if __name__ == "__main__":
    main()