main.py

import argparse
import sys
import traceback
import os
from luamc.bytecode_parser import BytecodeParser, LuaFunctionPrototype
from luamc.instruction import Instruction, decode_instruction
from luamc.mc_writer import MCWriter
from luamc.translator import Translator

def format_instruction(inst: Instruction, proto: LuaFunctionPrototype) -> str:
    op_name = inst.op_name
    line_str = f"[{inst.line:03d}]"
    base_str = f"{line_str} 0x{inst.raw:08X}  {op_name:<9} {inst.a:<3}"
    
    def format_rk(val: int) -> str:
        if (val & 0x100): # BITRK
            idx = val & 0xFF
            const = proto.constants[idx] if idx < len(proto.constants) else '?'
            return f"K({idx}) ; {const!r}"
        else:
            return f"R({val})"

    if inst.b is not None and inst.c is not None:
        if op_name in ("ADD", "SUB", "MUL", "DIV", "EQ", "LT", "LE", "GETTABLE", "SETTABLE", "SELF"):
            b_str, c_str = format_rk(inst.b), format_rk(inst.c)
        else:
            b_str = f"R({inst.b})"
            if op_name == 'LOADBOOL':
                c_str = str(inst.c)
            elif op_name == 'MOVE':
                c_str = ""
            else:
                 c_str = f"R({inst.c})"
        return f"{base_str} {b_str:<18} {c_str}"
    
    elif inst.bx is not None: # iABx
        bx_str = str(inst.bx)
        comment = ""
        if op_name in ("LOADK", "GETGLOBAL", "SETGLOBAL", "CLOSURE"):
            const = proto.constants[inst.bx] if inst.bx < len(proto.constants) else '?'
            comment = f" ; {const!r}"
        return f"{base_str} {bx_str}{comment}"
        
    elif inst.sbx is not None: # iAsBx
        target = inst.line + 1 + inst.sbx
        return f"{base_str} {inst.sbx:<4} ; to {target}"
    
    return base_str

def print_prototype(proto: LuaFunctionPrototype, indent: int = 0):
    prefix = "  " * indent
    source = proto.source_name if proto.source_name and '@' in proto.source_name else "<chunk>"
    print(f"{prefix}--- Prototype @ {source} (lines {proto.line_defined}-{proto.last_line_defined}) ---")
    print(f"{prefix}params: {proto.num_params}, stack: {proto.max_stack_size}, upvalues: {proto.num_upvalues}")
    if proto.constants:
        print(f"{prefix}constants ({len(proto.constants)}):")
        for i, const in enumerate(proto.constants):
            print(f"{prefix}  [{i}] = {const!r}")
    if proto.instructions:
        print(f"{prefix}code ({len(proto.instructions)}):")
        for i, inst_raw in enumerate(proto.instructions):
            decoded_inst = decode_instruction(inst_raw, i)
            formatted_str = format_instruction(decoded_inst, proto)
            print(f"{prefix}  {formatted_str}")
    if proto.prototypes:
        print(f"{prefix}prototypes ({len(proto.prototypes)}):")
        for sub_proto in proto.prototypes:
            print_prototype(sub_proto, indent + 1)
    print(f"{prefix}--------------------------------------------------\n")


def main():
    parser = argparse.ArgumentParser(
        description="LuaMC - A tool to parse, decode, and translate Lua 5.1 bytecode."
    )
    parser.add_argument(
        "luac_file",
        help="Path to the compiled .luac file."
    )
    parser.add_argument(
        "-o", "--output",
        default="luamc_datapack",
        help="Path to generate the output datapack."
    )
    parser.add_argument(
        "-t", "--templates",
        default="luamc/templates",
        help="Path to the template datapack directory."
    )
    args = parser.parse_args()

    try:
        print(f"[*] Parsing file: {args.luac_file}\n")
        parser = BytecodeParser(args.luac_file)
        main_prototype = parser.parse()
        
        print("[*] Parsing successful! Decoded structure:\n")
        print_prototype(main_prototype)

        if not os.path.isdir(args.templates):
            print(f"[!] Template directory not found at: {args.templates}", file=sys.stderr)
            sys.exit(1)
        writer = MCWriter(args.output, template_dir=args.templates)
        writer.setup_datapack()
        translator = Translator(main_prototype, writer)
        translator.translate()

    except (FileNotFoundError, ValueError, EOFError, TypeError) as e:
        print(f"[!] An error occurred: {e}", file=sys.stderr)
        traceback.print_exc()
        sys.exit(1)

if __name__ == "__main__":
    main()