#include "aoc.h"

typedef struct {
    u64 a, b, c;
    i64 ip;
} CPU;

static u64 
parse_register(str line, Arena arena) {
    str_find_result colon = str_find_left(line, STR(":"));
    ASSERT(colon.found);
    u64 result = (u64) parse_i64(str_trim(str_sub(line, colon.offset + 1, line.len)), arena);
    return result;
}

static u64
get_combo(CPU *cpu, u64 operand) {
    if (operand <= 3) return operand;
    else if (operand == 4) return cpu->a;
    else if (operand == 5) return cpu->b;
    else if (operand == 6) return cpu->c;
    else NOT_REACHABLE();
}

typedef DYNAMIC_ARRAY(u8) Program;

static void
run(Arena *arena, CPU * restrict cpu, Program program, Program *output) {
    output->len = 0;
    while (cpu->ip >= 0 && cpu->ip < program.len) {
        ASSERT(cpu->ip % 2 == 0);
        u8 op = program.data[cpu->ip];
        ASSERT(op >= 0);
        ASSERT(op < 8);
        u8 arg = program.data[cpu->ip + 1];

        bool took_branch = false;
        if (op == 0) {
            // adv
            cpu->a = cpu->a / (1ull << get_combo(cpu, (u64) arg));
        }
        else if (op == 1) {
            // bxl
            cpu->b = cpu->b ^ (u64) arg;
        }
        else if (op == 2) {
            // bst
            cpu->b = get_combo(cpu, (u64) arg) % 8;
        }
        else if (op == 3) {
            // jnz
            if (cpu->a != 0) {
                cpu->ip = (u64) arg;
                took_branch = true;
            }
            cpu->b = 0;
            cpu->c = 0;
        }
        else if (op == 4) {
            // bxc
            cpu->b ^= cpu->c;
        }
        else if (op == 5) {
            // out
            u8 value = get_combo(cpu, (u64) arg) % 8;
            *push(output, arena) = value;
        }
        else if (op == 6) {
            // bdv
            cpu->b = cpu->a / (1ull << get_combo(cpu, (u64) arg));
        }
        else if (op == 7) {
            // cdv
            cpu->c = cpu->a / (1ull << get_combo(cpu, (u64) arg));
        }
        else NOT_REACHABLE();

        if (!took_branch) {
            cpu->ip += 2;
        }
    }
}

static void
dump_program(Program program) {
    for (int i = 0; i < program.len; i++) {
        printf("%d", program.data[i]);
        if (i < program.len - 1) {
            putchar(',');
        }
    }
    puts("");
}

static u8
process(u64 *a) {
    u64 b, c;
    b = *a % 8;
    b ^= 2;
    c = *a >> b;
    b ^= c;
    *a >>= 3;
    b ^= 7;
    return b % 8;
}

static bool
check(Program *target, u64 a) {
    isize index = 0;
    for (;index < target->len; index++) {
        if (process(&a) != target->data[index]) {
            return false;
        }
    }
    return index == target->len;
}

static bool
check_suffix(Program *target, u64 a, i32 suffix_len) {
    isize index = 0;
    for (; index < target->len - suffix_len; index++) {
        process(&a);
    }
    for (; index < target->len; index++) {
        if (process(&a) != target->data[index]) {
            return false;
        }
    }
    return index == target->len;
}

static i64
solve(Program *target, i32 index, u64 a) {
    if (index == -1) {
        return a;
    }
    else {
        for (u64 i = 0; i < 8; i++) {
            u64 new_a = a;
            new_a |= (i << (3 * index));
            if (check_suffix(target, new_a, target->len - index)) {
                i64 result;
                if ((result = solve(target, index - 1, new_a)) != -1) {
                    return result;
                }
            }
        }
        return -1;
    }
}

int main(int argc, char **argv) {
    Arena *arena = make_arena(Megabytes(1));
    Tokens lines = read_lines(arena, argv[1]);

    CPU cpu = {0};
    cpu.a = parse_register(lines.tokens[0], *arena);
    cpu.b = parse_register(lines.tokens[1], *arena);
    cpu.c = parse_register(lines.tokens[2], *arena);

    CPU initial = cpu;

    Program program = {0};
    str token;
    str inp = lines.tokens[lines.len - 1];
    str_find_result colon = str_find_left(inp, STR(":"));
    inp = str_trim(str_sub(inp, colon.offset + 1, inp.len));
    while ((token = str_next_token(&inp, STR(","))).len > 0) {
        token = str_trim(token);
        *push(&program, arena) = (u8) parse_i64(token, *arena);
    }
    /* dump_program(program); */

    for (i32 i = 0; i < program.len; i += 2) {
        u8 op = program.data[i];
        u8 arg = program.data[i + 1];
        char combo[16] = {0};
        if (arg <= 3) {
            snprintf(combo, sizeof(combo), "%d", arg);
        } else if (arg == 4) {
            snprintf(combo, sizeof(combo), "A");
        } else if (arg == 5) {
            snprintf(combo, sizeof(combo), "B");
        } else if (arg == 6) {
            snprintf(combo, sizeof(combo), "C");
        } 
        switch (op) {
        case 0: printf("A = A / (1 << %s)\n", combo); break;
        case 1: printf("B = B ^ %d\n", arg); break;
        case 2: printf("B = %s %% 8\n", combo); break;
        case 3: printf("JNZ A, %d\n", arg); break;
        case 4: printf("B = B ^ C\n"); break;
        case 5: printf("OUT %s %% 8\n", combo); break;
        case 6: printf("B = A / (1 << %s)\n", combo); break;
        case 7: printf("C = A / (1 << %s)\n", combo); break;
        default: NOT_REACHABLE();
        }
    }

    Program output = {0};
    run(arena, &cpu, program, &output);
    dump_program(output);

    // Part 2
    /**
     * Our input program in disassembled form:
     * B = A % 8        ; equivalent to B = A & 0b111
     * B = B ^ 2        ; flips the 2nd bit of B
     * C = A / (1 << B) ; equivalent to C = A >> B
     * B = B ^ C        ; flips the bits of B that are set in C
     * A = A / (1 << 3) ; equivalent to A = A >> 3
     * B = B ^ 7        ; flips lower 3 bits of B
     * OUT B % 8        ; outputs lower 3 bits of B
     * JNZ A, 0
     *
     * Therefore, our input program processes A in chunks of 3 bits at a time,
     * until A becomes 0.
     * Registers B and C do not retain state between iterations of the program.
     * Hence, the length of the program's output is the number of octals in A.
     *
     * Combined:
     * B1 = A % 8
     * B2 = B1 ^ 2
     * C1 = A >> B2
     * B3 = B2 ^ C1
     * A2 = A >> 3
     * B4 = B3 ^ 7
     * OUT (B4 % 8)
     * JNZ A2, 0
     *
     * OUT (B4 % 8) = OUT ((B3 ^ 7) % 8)
     *              = OUT (((B2 ^ C1) ^ 7) % 8)
     *              = OUT (((B2 ^ (A >> B2)) ^ 7) % 8)
     *              = OUT (((((A % 8) ^ 2) ^ (A >> ((A % 8) ^ 2))) ^ 7) % 8)
     * 
     */

    i64 part_2 = solve(&program, program.len - 1, 0);
    ASSERT(part_2 != -1);
    printf("%ld\n", part_2);
    return 0;
}