#!/usr/bin/env python3

import argparse
import os
import pathlib
import parse
import logging

logger = logging.getLogger("split-ir")


def dump_ir(
    outdir: pathlib.Path, number: int, pass_name: str, target: str, contents: list
):
    logger.info("dump pass: {}, target: {}".format(pass_name, target))
    output_file_name = os.path.join(
        outdir, "{}-{}-{}.ll".format(number, pass_name, target)
    )
    with open(output_file_name, "w") as f:
        f.writelines(contents)


def main():
    parser = argparse.ArgumentParser(
        prog="split-ir",
        description="Split the IR dumped by `opt --print-after` or `opt --print-after-all` into multiple IR files.",
    )

    parser.add_argument(
        "-o",
        "--output-dir",
        type=pathlib.Path,
        required=True,
        help="output directory that contains all the IR files splitted from input",
    )
    parser.add_argument(
        "-v",
        "--verbose",
        action="store_true",
        default=False,
        help="verbose mode",
    )
    parser.add_argument(
        "input",
        type=pathlib.Path,
        help="input IR file that dumped from `opt -S --print-after` or `opt -S --print-after-all`",
    )

    args = parser.parse_args()

    if args.verbose:
        logging.basicConfig(level=logging.INFO)

    assert os.path.isfile(args.input), "file {} doesn't exist".format(args.input)

    if not os.path.exists(args.output_dir):
        os.mkdir(args.output_dir)
    assert os.path.isdir(args.output_dir), "file {} doesn't exist".format(
        args.output_dir
    )

    number = -1
    pass_name = None
    target = None
    contents = []

    with open(args.input, "r") as f:
        for line in f:
            if line.startswith("***"):
                line = line.strip()
                result = parse.parse("*** IR Dump After {pass} on {target} ***", line)
                assert result is not None, "broken splitter {}".format(line)
                if contents:
                    dump_ir(args.output_dir, number, pass_name, target, contents)
                    contents.clear()
                pass_name = result["pass"]
                target = (
                    result["target"] if result["target"] != "[module]" else "module"
                )
                number += 1
            else:
                contents.append(line)

    if contents:
        dump_ir(args.output_dir, number, pass_name, target, contents)


if __name__ == "__main__":
    main()