import dis
import io
import marshal
import os
import sys
import textwrap

from bytecode import fix_spacing, format_bytecode


def disassemble_pyc_file(pyc_file) -> str:
    with open(pyc_file, "rb") as f:
        header = f.read(16)
        magic_number = header[:2]
        if 3394 != int.from_bytes(magic_number, "little"):
            print(f"{pyc_file} not compiled with Python 3.7")
            return None

        bytecode = marshal.load(f)
    original_stdout = sys.stdout
    string_output = io.StringIO()
    sys.stdout = string_output
    dis.dis(bytecode)
    sys.stdout = original_stdout

    disassembled_pyc = string_output.getvalue()
    byte_code = fix_spacing(textwrap.dedent(str(disassembled_pyc)).strip())
    formatted_bytecode = format_bytecode(bytecode=byte_code)

    return formatted_bytecode


def convert_pyc_to_bytecode(file_path: str):
    formatted_bytecode = disassemble_pyc_file(file_path)

    if formatted_bytecode is None:
        return

    with open(file_path + "b", "w") as bytecode_file:
        print(f"Converted {file_path} to bytecode")
        bytecode_file.write(formatted_bytecode)


if __name__ == "__main__":
    path = sys.argv[1]
    if os.path.isfile(path) and path.endswith(".pyc"):
        convert_pyc_to_bytecode(path)
    elif os.path.isdir(path):
        for root, _, files in os.walk(path):
            for file in files:
                file_path = os.path.join(root, file)

                if not file.endswith(".pyc"):
                    continue

                convert_pyc_to_bytecode(file_path)