#include "aoc.h"
#include <ctype.h>

typedef struct Token {
    enum {
        T_EOF,

        T_INT,
        T_STR,
        T_LEFT_PAREN,
        T_RIGHT_PAREN,
        T_COMMA,
    } type;
    union {
        i64 integer;
        str string;
    };
} Token;

Token lex(str *input) {
    if (input->len == 0) {
        return (Token) { .type = T_EOF };
    }
    u8 ch = input->data[0];
    if (ch == '(') {
        *input = str_sub(*input, 1, input->len);
        return (Token) { .type = T_LEFT_PAREN };
    }
    else if (ch == ')') {
        *input = str_sub(*input, 1, input->len);
        return (Token) { .type = T_RIGHT_PAREN };
    }
    else if (ch == ',') {
        *input = str_sub(*input, 1, input->len);
        return (Token) { .type = T_COMMA };
    }
    else if (isdigit(ch)) {
        i64 integer = 0;
        while (isdigit(ch)) {
            integer = integer * 10 + (ch - '0');
            *input = str_sub(*input, 1, input->len);
            ch = input->data[0];
        }
        return (Token) { .type = T_INT, .integer = integer };
    }
    else {
        str string = { .data = input->data, .len = 0 };
        while (input->len > 0 && ch != '(' && ch != ')' && ch != ',' && !isdigit(ch)) {
            string.len++;
            *input = str_sub(*input, 1, input->len);
            ch = input->data[0];
        }
        return (Token) { .type = T_STR, .string = string };
    }
}

int main(int argc, char **argv) {
    Arena *arena = make_arena(Megabytes(1));
    str input = read_file(arena, argv[1]);

    struct {
        Token *data;
        isize len, cap;
    } tokens = {0};

    Token token;

    do {
        token = lex(&input);
        *push(&tokens, arena) = token;
    } while (token.type != T_EOF);

    i64 part_1 = 0;
    i64 part_2 = 0;
    bool enabled = true;
    for (isize i = 0; i < tokens.len; i++) {
        if ((i + 5 < tokens.len)
            && (tokens.data[i].type == T_STR)
            && (str_ends_with(tokens.data[i].string, STR("mul")))
            && (tokens.data[i + 1].type == T_LEFT_PAREN)
            && (tokens.data[i + 2].type == T_INT)
            && (tokens.data[i + 3].type == T_COMMA)
            && (tokens.data[i + 4].type == T_INT)
            && (tokens.data[i + 5].type == T_RIGHT_PAREN))
        {
            i64 factor = tokens.data[i + 2].integer * tokens.data[i + 4].integer;
            part_1 += factor;
            part_2 += enabled ? factor : 0;
        }
        else if (((i + 2) < tokens.len)
            && (tokens.data[i].type == T_STR)
            && (str_ends_with(tokens.data[i].string, STR("do")))
            && (tokens.data[i + 1].type == T_LEFT_PAREN)
            && (tokens.data[i + 2].type == T_RIGHT_PAREN))
        {
            enabled = true;
        }
        else if (((i + 2) < tokens.len)
            && (tokens.data[i].type == T_STR)
            && (str_ends_with(tokens.data[i].string, STR("don't")))
            && (tokens.data[i + 1].type == T_LEFT_PAREN)
            && (tokens.data[i + 2].type == T_RIGHT_PAREN))
        {
            enabled = false;
        }
    }

    printf("%ld\n", part_1);
    printf("%ld\n", part_2);
}