from lib import get_data, str_to_ints


def part_1(data):
    xs = str_to_ints(data)
    i = 0
    while i < len(xs):
        inst = str(xs[i])
        inst = "0" * (5 - len(inst)) + inst
        assert len(inst) == 5
        op = int(inst[3:5])
        mode_p1 = int(inst[2])
        mode_p2 = int(inst[1])
        mode_p3 = int(inst[0])
        match op:
            case 1:
                p1 = xs[xs[i + 1]] if mode_p1 == 0 else xs[i + 1]
                p2 = xs[xs[i + 2]] if mode_p2 == 0 else xs[i + 2]
                assert mode_p3 == 0
                xs[xs[i + 3]] = p1 + p2
                i += 4
            case 2:
                p1 = xs[xs[i + 1]] if mode_p1 == 0 else xs[i + 1]
                p2 = xs[xs[i + 2]] if mode_p2 == 0 else xs[i + 2]
                assert mode_p3 == 0
                xs[xs[i + 3]] = p1 * p2
                i += 4
            case 3:
                print("input", i, 1)
                assert mode_p1 == 0
                xs[xs[i + 1]] = 1
                i += 2
            case 4:
                if mode_p1 == 0:
                    v = xs[xs[i + 1]]
                else:
                    v = xs[i + 1]
                print("output", v)
                i += 2
            case 99:
                break


def part_2(data):
    xs = str_to_ints(data)
    i = 0
    while i < len(xs):
        inst = str(xs[i])
        inst = "0" * (5 - len(inst)) + inst
        assert len(inst) == 5
        op = int(inst[3:5])
        mode_p1 = int(inst[2])
        mode_p2 = int(inst[1])
        mode_p3 = int(inst[0])
        match op:
            case 1:
                p1 = xs[xs[i + 1]] if mode_p1 == 0 else xs[i + 1]
                p2 = xs[xs[i + 2]] if mode_p2 == 0 else xs[i + 2]
                assert mode_p3 == 0
                xs[xs[i + 3]] = p1 + p2
                i += 4
            case 2:
                p1 = xs[xs[i + 1]] if mode_p1 == 0 else xs[i + 1]
                p2 = xs[xs[i + 2]] if mode_p2 == 0 else xs[i + 2]
                assert mode_p3 == 0
                xs[xs[i + 3]] = p1 * p2
                i += 4
            case 3:
                print("input", i, 5)
                assert mode_p1 == 0
                xs[xs[i + 1]] = 5
                i += 2
            case 4:
                if mode_p1 == 0:
                    v = xs[xs[i + 1]]
                else:
                    v = xs[i + 1]
                print("output", v)
                i += 2
            case 99:
                break
            case 5:
                p1 = xs[xs[i + 1]] if mode_p1 == 0 else xs[i + 1]
                p2 = xs[xs[i + 2]] if mode_p2 == 0 else xs[i + 2]
                if p1 != 0:
                    i = p2
                else:
                    i += 3
            case 6:
                p1 = xs[xs[i + 1]] if mode_p1 == 0 else xs[i + 1]
                p2 = xs[xs[i + 2]] if mode_p2 == 0 else xs[i + 2]
                if p1 == 0:
                    i = p2
                else:
                    i += 3
            case 7:
                p1 = xs[xs[i + 1]] if mode_p1 == 0 else xs[i + 1]
                p2 = xs[xs[i + 2]] if mode_p2 == 0 else xs[i + 2]
                assert mode_p3 == 0
                if p1 < p2:
                    xs[xs[i + 3]] = 1
                else:
                    xs[xs[i + 3]] = 0
                i += 4
            case 8:
                p1 = xs[xs[i + 1]] if mode_p1 == 0 else xs[i + 1]
                p2 = xs[xs[i + 2]] if mode_p2 == 0 else xs[i + 2]
                assert mode_p3 == 0
                if p1 == p2:
                    xs[xs[i + 3]] = 1
                else:
                    xs[xs[i + 3]] = 0
                i += 4


def main():
    data = get_data(__file__)
    part_1(data)
    part_2(data)


if __name__ == "__main__":
    main()