#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>

#include "ast.h"
#include "common.h"
#include "parser.h"

#define N 1024

void astNodeListEnsureCapacity(AstNodeList* list, uint32_t additional)
{
    const uint32_t new_len = list->len + additional;
    if (new_len <= list->cap) {
        return;
    }

    const uint32_t new_cap = new_len > list->cap * 2 ? new_len : list->cap * 2;
    list->tags = realloc(list->tags, new_cap * sizeof(AstNodeTag));
    list->main_tokens = realloc(list->main_tokens, new_cap * sizeof(AstTokenIndex));
    list->datas = realloc(list->datas, new_cap * sizeof(AstData));
    if (!list->tags || !list->main_tokens || !list->datas)
        exit(1);
    list->cap = new_cap;
}

AstNodeIndex astNodeListAppend(
    AstNodeList* list,
    AstNodeTag tag,
    AstTokenIndex main_token,
    AstData data)
{
    astNodeListEnsureCapacity(list, 1);
    list->tags[list->len] = tag;
    list->main_tokens[list->len] = main_token;
    list->datas[list->len] = data;
    return list->len++;
}

Ast astParse(const char* source, const uint32_t len)
{
    uint32_t estimated_token_count = len / 8;

    // Initialize token list
    AstTokenList tokens = {
        .len = 0,
        .cap = estimated_token_count,
        .tags = SLICE_INIT(TokenizerTag, estimated_token_count),
        .starts = SLICE_INIT(AstIndex, estimated_token_count)
    };

    // Tokenize
    Tokenizer tok = tokenizerInit(source, len);
    while (true) {
        if (tokens.len >= tokens.cap) {
            fprintf(stderr, "too many tokens, bump estimated_token_count\n");
            exit(1);
        }
        TokenizerToken token = tokenizerNext(&tok);
        tokens.tags[tokens.len] = token.tag;
        tokens.starts[tokens.len] = token.loc.start;
        tokens.len++;
        if (token.tag == TOKENIZER_TAG_EOF)
            break;
    }

    // Initialize node list
    uint32_t estimated_node_count = (tokens.len + 2) / 2;
    AstNodeList nodes = {
        .len = 0,
        .cap = estimated_node_count,
        .tags = SLICE_INIT(AstNodeTag, estimated_node_count),
        .main_tokens = SLICE_INIT(AstTokenIndex, estimated_node_count),
        .datas = SLICE_INIT(AstData, estimated_node_count)
    };

    // Initialize parser
    Parser p = {
        .source = source,
        .source_len = len,
        .token_tags = tokens.tags,
        .token_starts = tokens.starts,
        .tokens_len = tokens.len,
        .tok_i = 0,
        .nodes = nodes,
        .extra_data = {
            .len = 0,
            .cap = N,
            .arr = SLICE_INIT(AstNodeIndex, N) },
        .scratch = { .len = 0, .cap = N, .arr = SLICE_INIT(AstNodeIndex, N) }
    };

    free(p.scratch.arr); // Parser takes ownership

    parseRoot(&p);

    return (Ast) {
        .source = source,
        .source_len = len,
        .tokens = tokens,
        .nodes = p.nodes,
        .extra_data = {
            .len = p.extra_data.len,
            .cap = p.extra_data.cap,
            .arr = p.extra_data.arr,
        },
    };
}