import os
from pathlib import Path
from argparse import ArgumentParser
from .vminit import vminit
from .memory import translate_memory
from .arith_logic import translate_arith_logic
from .compare import translate_compare
from .branching import translate_branching
from .function import translate_function, translate_return
from .utils import *


def vm_translate(input_path, verbose, init):
    # if input_path is a file, translate the file
    # if it is a directory, translate every .vm file inside
    try:
        filenames = os.listdir(input_path)
        files = [Path(input_path / f) for f in filenames]
        vm_files = filter(lambda fn: fn.suffix == ".vm", files)
        output_path = input_path / (input_path.name + ".asm")
    except NotADirectoryError:
        vm_files = [Path(input_path)]
        output_path = input_path.with_suffix(".asm")
    except:
        exit_on_error(f"Cannot open input path: {input_path}", EXIT_CODE_FILE_ERROR)

    # assembly translated from all .vm files will be concatenated here
    asm = vminit(input_path.stem) if init else ""

    for input_fn in vm_files:
        # load all vm commands from file, e.g. program/Main.vm
        try:
            input_file = open(input_fn)
        except:
            exit_on_error(f"Cannot open input file: {input_fn}", EXIT_CODE_FILE_ERROR)

        prog = input_fn.stem  # e.g. Main
        vm_cmds = []
        line = input_file.readline()
        while line:
            cmd = line.rstrip("\n").split("//", maxsplit=1)[0].strip()
            if cmd:
                vm_cmds.append(cmd)
            line = input_file.readline()

        input_file.close()

        for cmd in vm_cmds:
            # tokenize vm_line, hand to sub-translators based on first token
            match cmd.split():
                case []:
                    continue
                case [("push" | "pop") as action, segment, index]:
                    asm += translate_memory(action, segment, index, prog, verbose)
                case [("add" | "sub" | "neg" | "and" | "or" | "not") as command]:
                    asm += translate_arith_logic(command, verbose)
                case [("lt" | "eq" | "gt") as command]:
                    asm += translate_compare(command, verbose)
                case [("label" | "goto" | "if-goto") as action, label]:
                    asm += translate_branching(action, label, prog, verbose)
                case [("call" | "function") as action, function, n]:
                    asm += translate_function(action, function, n, prog, verbose)
                case ["return"]:
                    asm += translate_return(verbose)
                case _:
                    exit_on_error(
                        f"Syntax error: invalid command: {cmd}", EXIT_CODE_SYNTAX_ERROR
                    )

    try:
        output_file = open(output_path, "w")
    except:
        exit_on_error(f"Cannot open output file: {output_path}", EXIT_CODE_FILE_ERROR)

    output_file.write(asm)
    output_file.close()
    print(f"Assembly written to {output_path}")


if __name__ == "__main__":
    parser = ArgumentParser("hack-vm")
    parser.add_argument("-v", "--verbose", action="store_true", help="verbose mode")
    parser.add_argument("-I", "--no-init", action="store_true", help="do not write bootstrap code")
    parser.add_argument("input_path", help="HACK VM file or directory")
    args = parser.parse_args()
    vm_translate(Path(args.input_path), args.verbose, not args.no_init)