#include <scc_tree_dump.h>

#define DEFAULT_VERTICAL "| "
#define DEFAULT_BRANCH "|-"
#define DEFAULT_LAST_BRANCH "`-"
#define DEFAULT_SPACE "  "
#define NODE_COLOR "\033[34m"
#define VALUE_COLOR "\033[32m"
#define BRANCH_COLOR "\033[33m"
#define RESET_COLOR "\033[0m"

void scc_tree_dump_init(scc_tree_dump_t *td, cbool use_color) {
    scc_str_init(&td->buf);
    scc_vec_init(td->stack);
    td->use_color = use_color;
    td->line_start = true;

    td->vertical = DEFAULT_VERTICAL;
    td->branch = DEFAULT_BRANCH;
    td->last_branch = DEFAULT_LAST_BRANCH;
    td->space = DEFAULT_SPACE;

    td->node_color = use_color ? NODE_COLOR : "";
    td->value_color = use_color ? VALUE_COLOR : "";
    td->branch_color = use_color ? BRANCH_COLOR : "";
    td->reset_color = use_color ? RESET_COLOR : "";
}

void scc_tree_dump_drop(scc_tree_dump_t *td) {
    scc_str_drop(&td->buf);
    scc_vec_free(td->stack);
}

void scc_tree_dump_clear(scc_tree_dump_t *td) {
    scc_str_clear(&td->buf);
    td->line_start = true;
}

const char *scc_tree_dump_cstr(scc_tree_dump_t *td) {
    return scc_str_as_cstr(&td->buf);
}

void scc_tree_dump_flush(scc_tree_dump_t *td,
                         void (*output)(const char *str, usize len, void *user),
                         void *user) {
    if (td->buf.size > 1 && output) {
        output(td->buf.data, td->buf.size - 1, user);
    }
    scc_tree_dump_clear(td);
}

void scc_tree_dump_begin_line(scc_tree_dump_t *td) {
    if (!td->line_start) {
        // 如果不在行首，先换行（表示上一行结束）
        if (td->buf.size > 1) {
            scc_str_append_ch(&td->buf, '\n');
        }
    }
    // 输出缩进
    usize depth = scc_vec_size(td->stack);
    for (usize i = 0; i < depth; i++) {
        cbool last = scc_vec_at(td->stack, i);
        const char *prefix;
        if (i + 1 == depth) {
            prefix = last ? td->last_branch : td->branch;
        } else {
            prefix = last ? td->space : td->vertical;
        }
        if (td->use_color) {
            scc_str_append_fmt(&td->buf, "%s%s%s", td->branch_color, prefix,
                               td->reset_color);
        } else {
            scc_str_append_cstr(&td->buf, prefix, scc_strlen(prefix));
        }
    }
    td->line_start = false;
}

void scc_tree_dump_append(scc_tree_dump_t *td, const char *s) {
    if (td->line_start) {
        scc_tree_dump_begin_line(td);
    }
    scc_str_append_cstr(&td->buf, s, scc_strlen(s));
}

void scc_tree_dump_append_fmt(scc_tree_dump_t *td, const char *fmt, ...) {
    if (td->line_start) {
        scc_tree_dump_begin_line(td);
    }
    va_list args;
    va_start(args, fmt);
    // 计算所需长度
    va_list args_copy;
    va_copy(args_copy, args);
    usize len = scc_vsnprintf(null, 0, fmt, args_copy);
    va_end(args_copy);
    if (len == 0) {
        va_end(args);
        return;
    }
    // 确保缓冲区容量足够
    if (td->buf.size + len + 1 > td->buf.cap) {
        usize new_cap = td->buf.cap;
        while (new_cap < td->buf.size + len + 1) {
            new_cap = new_cap ? new_cap * 2 : 16;
        }
        char *new_data = (char *)scc_realloc(td->buf.data, new_cap);
        Assert(new_data != null);
        td->buf.data = new_data;
        td->buf.cap = new_cap;
    }
    scc_vsnprintf(td->buf.data + td->buf.size - (td->buf.size == 0 ? 0 : 1),
                  len + 1, fmt, args);
    va_end(args);
    if (td->buf.size == 0) {
        td->buf.size = 1;
    }
    td->buf.size += len;
    td->buf.data[td->buf.size - 1] = '\0';
}

void scc_tree_dump_push(scc_tree_dump_t *td, cbool is_last) {
    scc_vec_push(td->stack, is_last);
}

void scc_tree_dump_pop(scc_tree_dump_t *td) { (void)scc_vec_pop(td->stack); }