#include "aoc.h"

typedef struct {
    i32 *data;
    isize len, cap;
} Ints;

static bool
is_safe(i32 *nums, isize len) {
    bool expect_increase = nums[1] > nums[0];
    for (isize i = 0; i < len - 1; i++) {
        i32 a = nums[i], b = nums[i + 1];
        i32 diff = b - a;
        if (expect_increase) {
            if (diff < 0) return false;
            if (diff > 3) return false;
            if (diff < 1) return false;
        }
        else {
            if (diff > 0) return false;
            if (diff < -3) return false;
            if (diff > -1) return false;
        }
    }
    return true;
}

static bool
is_safe_with_exception(i32 *nums, i32 len, Arena temp) {
    i32 *selection = ARENA_ALLOC_ARRAY(&temp, i32, len - 1);
    for (isize skip = 0; skip < len; skip++) {
        memcpy(selection, nums, skip * sizeof(i32));
        memcpy(selection + skip, nums + skip + 1, (len - skip - 1) * sizeof(i32));
        if (is_safe(selection, len - 1)) return true;
    }
    return false;
}

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

    i32 part_1 = 0;
    i32 part_2 = 0;

    for (isize i = 0; i < lines.len; i++) {
        str line = lines.tokens[i];
        Ints values = {0};
        Tokens nums = str_tokenize(arena, line, STR(" "));
        for (isize j = 0; j < nums.len; j++) {
            *push(&values, arena) = parse_i64(str_trim(nums.tokens[j]), *arena);
        }
        part_1 += is_safe(values.data, values.len);
        part_2 += is_safe_with_exception(values.data, values.len, *arena);
    }
    printf("%d\n", part_1);
    ASSERT(part_1 == 411);

    printf("%d\n", part_2);
    ASSERT(part_2 == 465);
}