#include "ast.h"
#include "collection.h"
#include <assert.h>
#include <stdio.h>
#include <stdlib.h>

static void push(AstStringifier* s, const char* str)
{
    for (; *str; ++str) {
        arena_array_push(s->arena,
            (void**)&s->string,
            &s->string_cap,
            &s->string_size,
            (void*)str,
            sizeof(*str));
    }

    char ch = '\0';
    arena_array_push(s->arena,
        (void**)&s->string,
        &s->string_cap,
        &s->string_size,
        (void*)&ch,
        sizeof(ch));
    s->string_size -= 1;
}

#define pushf(S, ...)                                                          \
    do {                                                                       \
        size_t len = snprintf(nullptr, 0, __VA_ARGS__);                        \
        char* buf = calloc(len + 1, 1);                                        \
        snprintf(buf, len + 1, __VA_ARGS__);                                   \
        push(S, buf);                                                          \
        free(buf);                                                             \
    } while (false)

static void indent(AstStringifier* s)
{
    for (int i = 0; i < s->depth; ++i) {
        push(s, "    ");
    }
}

void ast_stringifier_init(AstStringifier* s, Arena* arena)
{
    *s = (AstStringifier) {
        ARRAY_FIELD_INIT(string),
        .arena = arena,
        .depth = 0,
    };
}

void ast_stmt_block_to_string(AstStringifier* s, const AstStmtBlock* node)
{
    pushf(s, "StmtBlock %d:%d {\n", node->loc.line, node->loc.col);
    s->depth += 1;
    indent(s);
    push(s, "stmts: [");

    if (node->stmts_size) {
        s->depth += 1;
        push(s, "\n");
        for (size_t i = 0; i < node->stmts_size; ++i) {
            indent(s);
            ast_stmt_to_string(s, node->stmts[i]);
            push(s, ",\n");
        }
        s->depth -= 1;
        indent(s);
    }
    push(s, "],");

    s->depth -= 1;
    indent(s);
    push(s, "}");
}

void ast_stmt_to_string(AstStringifier* s, const AstStmt* node)
{
    switch (node->ty) {
        case AstStmt_Decl: {
            pushf(s, "Stmt::Decl %d:%d {\n", node->loc.line, node->loc.col);
            s->depth += 1;

            indent(s);
            push(s, "decl: ");
            ast_decl_to_string(s, node->decl);
            push(s, ",\n");

            s->depth -= 1;
            indent(s);
            push(s, "}");
            break;
        }
        case AstStmt_Expr: {
            pushf(s, "Stmt::Expr %d:%d {\n", node->loc.line, node->loc.col);
            s->depth += 1;
            indent(s);

            push(s, "expr: ");

            ast_expr_to_string(s, node->expr);
            push(s, ",\n");

            s->depth -= 1;
            indent(s);
            push(s, "}");
            break;
        }
    }
}

static void spec_to_string(AstStringifier* s, AstSpec spec)
{
    struct {
        AstSpec specifier;
        const char* ident;
    } spec_idents[] = {
        { AstSpec_Int, "int" },
        { AstSpec_Char, "char" },
        { AstSpec_Unsigned, "unsigned" },
        { AstSpec_Static, "static" },
        { AstSpec_Const, "const" },
        { AstSpec_Inline, "inline" },
        { AstSpec_Typedef, "typedef" },
        { AstSpec_Enum, "enum" },
        { AstSpec_Struct, "struct" },
        { AstSpec_Union, "union" },
    };

    bool first = true;

    for (size_t i = 0; i < sizeof(spec_idents) / sizeof(spec_idents[0]); ++i) {
        if (spec & spec_idents[i].specifier) {
            if (!first)
                push(s, " ");
            first = false;
            push(s, spec_idents[i].ident);
        }
    }
}

void ast_decl_to_string(AstStringifier* s, const AstDecl* node)
{
    pushf(s, "Decl %d:%d {\n", node->loc.line, node->loc.col);
    s->depth += 1;

    indent(s);
    push(s, "spec: [");
    spec_to_string(s, node->spec);
    push(s, "],\n");

    if (node->spec & AstSpec_DefTy) {
        assert(node->type_ident);
        indent(s);
        pushf(s, "type_ident: \"%s\",\n", node->type_ident);
    }

    indent(s);
    push(s, "declars: [");
    if (node->declars_size) {
        push(s, "\n");
        s->depth += 1;

        for (size_t i = 0; i < node->declars_size; ++i) {
            indent(s);
            ast_declar_to_string(s, node->declars[i]);
            push(s, ",\n");
        }

        s->depth -= 1;
        indent(s);
        push(s, "],\n");
    } else {
        push(s, "],\n");
    }

    if (node->function_body) {
        indent(s);
        push(s, "function_body: ");
        ast_stmt_block_to_string(s, node->function_body);
        push(s, ",\n");
    }

    s->depth -= 1;
    indent(s);
    push(s, "}");
}

void ast_declar_to_string(AstStringifier* s, const AstDeclar* node)
{
    switch (node->ty) {
        case AstDeclar_Ident: {
            pushf(s, "Declar::Ident %d:%d {\n", node->loc.line, node->loc.col);
            s->depth += 1;

            indent(s);
            pushf(s, "ident: \"%s\",\n", node->ident);

            break;
        }
        case AstDeclar_Pointer: {
            pushf(
                s, "Declar::Pointer %d:%d {\n", node->loc.line, node->loc.col);
            s->depth += 1;

            indent(s);
            push(s, "declar: ");
            ast_declar_to_string(s, node->pointer.declar);
            push(s, ",\n");

            indent(s);
            push(s, "spec: [");
            spec_to_string(s, node->pointer.spec);
            push(s, "],\n");

            break;
        }
        case AstDeclar_Array: {
            pushf(s, "Declar::Array %d:%d {\n", node->loc.line, node->loc.col);
            s->depth += 1;

            indent(s);
            push(s, "declar: ");
            ast_declar_to_string(s, node->array.declar);
            push(s, ",\n");

            indent(s);
            push(s, "expr: ");
            ast_expr_to_string(s, node->array.expr);
            push(s, ",\n");

            break;
        }
        case AstDeclar_Function: {
            pushf(
                s, "Declar::Function %d:%d {\n", node->loc.line, node->loc.col);
            s->depth += 1;

            indent(s);
            push(s, "declar: ");
            ast_declar_to_string(s, node->array.declar);
            push(s, ",\n");

            indent(s);
            push(s, "params: [");
            if (node->function.params_size) {
                push(s, "\n");
                s->depth += 1;

                for (size_t i = 0; i < node->function.params_size; ++i) {
                    indent(s);
                    ast_decl_to_string(s, node->function.params[i]);
                    push(s, ",\n");
                }

                s->depth -= 1;
                indent(s);
                push(s, "],\n");
            } else {
                push(s, "],\n");
            }

            break;
        }
    }

    if (node->init_expr) {
        indent(s);
        push(s, "decl: ");
        ast_expr_to_string(s, node->init_expr);
        push(s, ",\n");
    }

    s->depth -= 1;
    indent(s);
    push(s, "}");
}

void ast_expr_to_string(AstStringifier* s, const AstExpr* node)
{
    push(s, "<expr>");
}