from typing import List


def findLabel(lines: List[List[str]], label: str) -> int:
    """Find the line number of a label"""
    for i, line in enumerate(lines):
        if line[0] == label + ":":
            return i
    raise ValueError(f"Label '{label}' not found")


def parse(args: List[int], values: List[int], arg: str) -> int:
    """Parse an argument string into an integer value"""
    arg_val = 0

    if arg.isdigit():
        arg_val = int(arg)
    elif arg.startswith("L"):
        arg_val = values[int(arg[1:])]
    elif arg.startswith("A"):
        arg_val = args[int(arg[1:])]

    return arg_val


def runL(args: List[int], values: List[int], line: List[str]):
    """Run a Set command"""
    out_var = int(line[0][1:])
    op = line[1]

    if out_var >= len(values):
        values.extend([0] * (out_var - len(values) + 1))

    if op.isdigit():
        values[out_var] = int(op)
    else:
        var1, var2 = line[2], line[3]
        var1_val, var2_val = parse(args, values, var1), parse(args, values, var2)
        if op == "+":
            values[out_var] = var1_val + var2_val
        elif op == "-":
            values[out_var] = var1_val - var2_val
        else:
            raise ValueError(f"Invalid Operator: {op}")


def handleTargetError(target: int, lines: List[List[str]]) -> int | None:
    if target < 0 or target >= len(lines):
        raise ValueError(f"Invalid Jump Target: {target}")
    return target


def runJMP(args: List[int], values: List[int], line: List[str], lines) -> int | None:
    """Return the target line number or None for a JMP command"""
    command = line[0]
    if command == "JMP":
        jump_targ = line[1]
        if jump_targ.isdigit():
            target = int(line[1])
            return handleTargetError(target, lines)
        else:
            target = findLabel(lines, line[1])
            return handleTargetError(target, lines)
    else:
        expr = line[1]
        if command[3] == "+":
            var = int(expr[1])
            if values[var] > 0:
                target = findLabel(lines, line[2])
                return handleTargetError(target, lines)
            return None
        elif command[3] == "0":
            var = int(expr[1])
            if values[var] == 0:
                target = findLabel(lines, line[2])
                return handleTargetError(target, lines)
            return None
        else:
            raise ValueError(f"Invalid Jump Target: {line[0]}")


def runSimpleProgram(program: str, args: List[int]):
    """Run a simple program with given arguments."""
    lines = program.splitlines()
    lines = [line.strip() for line in lines if line.strip()]
    lines = [line.split() for line in lines if not line.startswith("!")]

    values: List[int] = []
    i = 0
    return_value = None
    while not return_value:
        line = lines[i]
        i += 1
        # print(f"executing {i}: {line}, L: {values}, A: {args}")
        command = line[0]
        if ":" in line:
            continue
        if command.startswith("L"):
            runL(args, values, line)
        elif command.startswith("JMP"):
            jump_to = runJMP(args, values, line, lines)
            if jump_to != None:
                i = jump_to
        elif command == "RTN":
            val = line[1]
            # print(f"RTN {val} from A: {args}, L: {values}")
            return_value = parse(args, values, val)
            break
    return return_value


def testRunSimpleProgram():
    print("Testing runSimpleProgram()...", end="")
    largest = """! largest: Returns max(A0, A1)
                   L0 - A0 A1
                   JMP+ L0 a0
                   RTN A1
                   a0:
                   RTN A0"""
    assert runSimpleProgram(largest, [5, 6]) == 6
    assert runSimpleProgram(largest, [6, 5]) == 6

    sumToN = """! SumToN: Returns 1 + ... + A0
                ! L0 is a counter, L1 is the result
                L0 0
                L1 0
                loop:
                L2 - L0 A0
                JMP0 L2 done
                L0 + L0 1
                L1 + L1 L0
                JMP loop
                done:
                RTN L1"""
    assert runSimpleProgram(sumToN, [5]) == 1 + 2 + 3 + 4 + 5
    assert runSimpleProgram(sumToN, [10]) == 10 * 11 // 2
    print("Passed!")


testRunSimpleProgram()