#include "spl_diag.h"
#include "../stage0/include/color.h"

#include <stdarg.h>

static const char *const g_kind_names[] = {
    [SPL_DIAG_KIND_NONE] = "none",
    [SPL_DIAG_KIND_LEX_UNEXPECTED] = "lex_unexpected",
    [SPL_DIAG_KIND_PARSE_EXPECTED] = "parse_expected",
    [SPL_DIAG_KIND_PARSE_UNRECOGNIZED] = "parse_unrecognized",
    [SPL_DIAG_KIND_UNDEFINED_SYMBOL] = "undefined_symbol",
    [SPL_DIAG_KIND_DUPLICATE_SYMBOL] = "duplicate_symbol",
    [SPL_DIAG_KIND_TYPE_MISMATCH] = "type_mismatch",
    [SPL_DIAG_KIND_ARG_COUNT_MISMATCH] = "arg_count_mismatch",
    [SPL_DIAG_KIND_NOT_A_FUNCTION] = "not_a_function",
    [SPL_DIAG_KIND_NO_FIELD] = "no_field",
    [SPL_DIAG_KIND_NO_METHOD] = "no_method",
    [SPL_DIAG_KIND_INDEX_BASE_INVALID] = "index_base_invalid",
    [SPL_DIAG_KIND_TUPLE_INDEX_OOB] = "tuple_index_oob",
    [SPL_DIAG_KIND_FITS_OVERFLOW] = "fits_overflow",
    [SPL_DIAG_KIND_CIRCULAR_DEP] = "circular_dep",
    [SPL_DIAG_KIND_BREAK_OUTSIDE_LOOP] = "break_outside_loop",
    [SPL_DIAG_KIND_CONTINUE_OUTSIDE_LOOP] = "continue_outside_loop",
    [SPL_DIAG_KIND_UNDEFINED_LABEL] = "undefined_label",
    [SPL_DIAG_KIND_CAST_INVALID] = "cast_invalid",
    [SPL_DIAG_KIND_UNKNOWN_ATTR] = "unknown_attr",
    [SPL_DIAG_KIND_UNKNOWN_INTRIN] = "unknown_intrin",
    [SPL_DIAG_KIND_INTRIN_ARG_COUNT] = "intrin_arg_count",
    [SPL_DIAG_KIND_VAR_UNINITIALIZED] = "var_uninitialized",
    [SPL_DIAG_KIND_CANNOT_INFER] = "cannot_infer",
    [SPL_DIAG_KIND_MUST_RETURN_VALUE] = "must_return_value",
    [SPL_DIAG_KIND_UNHANDLED_NODE] = "unhandled_node",
    [SPL_DIAG_KIND_IMPLICIT_CAST_FAIL] = "implicit_cast_fail",
    [SPL_DIAG_KIND_NO_LVALUE] = "no_lvalue",
    [SPL_DIAG_KIND_BAD_FIELD_INDEX] = "bad_field_index",
    [SPL_DIAG_KIND_BAD_BLOCK] = "bad_block",
    [SPL_DIAG_KIND_IR_UNKNOWN_KIND] = "ir_unknown_kind",
    [SPL_DIAG_KIND_FRAME_OVERFLOW] = "frame_overflow",
    [SPL_DIAG_KIND_INTERNAL] = "internal",
};

const char *spl_diag_kind_name(spl_diag_kind_t k) {
    if ((usize)k < sizeof(g_kind_names) / sizeof(g_kind_names[0]) && g_kind_names[k])
        return g_kind_names[k];
    return "unknown";
}

static const char *const g_stage_names[] = {
    [SPL_STAGE_NONE] = "none",     [SPL_STAGE_LEX] = "lex",       [SPL_STAGE_PARSE] = "parse",
    [SPL_STAGE_SEMA] = "sema",     [SPL_STAGE_AST2IR] = "ast2ir", [SPL_STAGE_IR2VM] = "ir2vm",
    [SPL_STAGE_DRIVER] = "driver",
};

const char *spl_diag_stage_name(spl_diag_stage_t s) {
    if ((usize)s < sizeof(g_stage_names) / sizeof(g_stage_names[0]) && g_stage_names[s])
        return g_stage_names[s];
    return "?";
}

const char *spl_diag_level_name(spl_diag_level_t l) {
    switch (l) {
    case SPL_DIAG_LEVEL_NOTE:
        return "note";
    case SPL_DIAG_LEVEL_WARN:
        return "warn";
    case SPL_DIAG_LEVEL_ERROR:
        return "error";
    case SPL_DIAG_LEVEL_FATAL:
        return "fatal";
    }
    return "?";
}

static const char *level_color(spl_diag_level_t l) {
    switch (l) {
    case SPL_DIAG_LEVEL_NOTE:
        return ANSI_FG_CYAN;
    case SPL_DIAG_LEVEL_WARN:
        return ANSI_FG_YELLOW;
    case SPL_DIAG_LEVEL_ERROR:
        return ANSI_FG_RED;
    case SPL_DIAG_LEVEL_FATAL:
        return ANSI_FG_RED ANSI_UNDERLINED;
    }
    return ANSI_NONE;
}

static char *diag_strdup(const char *s) {
    if (!s)
        return NULL;
    usize n = strlen(s);
    char *r = malloc(n + 1);
    if (!r)
        return NULL;
    memcpy(r, s, n);
    r[n] = '\0';
    return r;
}

void spl_diag_init(spl_diag_t *d, void *fout) {
    memset(d, 0, sizeof *d);
    vec_init(d->ctx);
    vec_init(d->records);
    d->collect = true;
    d->fout = fout;
    d->color = true;
    d->stage = SPL_STAGE_NONE;
}

static void frame_free(spl_diag_frame_t *f) {
    free(f->fn_name);
    free(f->label);
    f->fn_name = NULL;
    f->label = NULL;
}

static void record_free(spl_diag_record_t *r) {
    free(r->message);
    r->message = NULL;
    for (usize i = 0; i < vec_size(r->ctx); ++i)
        frame_free(&vec_at(r->ctx, i));
    vec_free(r->ctx);
    for (usize i = 0; i < vec_size(r->notes); ++i)
        free(vec_at(r->notes, i).message);
    vec_free(r->notes);
}

void spl_diag_reset(spl_diag_t *d) {
    for (usize i = 0; i < vec_size(d->ctx); ++i)
        frame_free(&vec_at(d->ctx, i));
    d->ctx.size = 0;

    for (usize i = 0; i < vec_size(d->records); ++i)
        record_free(&vec_at(d->records, i));
    d->records.size = 0;

    d->num_printed = 0;
    d->n_errors = 0;
    d->n_warns = 0;
    d->n_fatals = 0;
}

void spl_diag_drop(spl_diag_t *d) {
    spl_diag_reset(d);
    vec_free(d->ctx);
    vec_free(d->records);
}

void spl_diag_set_stage(spl_diag_t *d, spl_diag_stage_t stage) { d->stage = stage; }

void spl_diag_set_output(spl_diag_t *d, FILE *out) { d->fout = out; }

void spl_diag_set_color(spl_diag_t *d, bool enabled) { d->color = enabled; }

void spl_diag_set_collect(spl_diag_t *d, bool enabled) { d->collect = enabled; }

void spl_diag_push(spl_diag_t *d, const spl_dbg_node_t *loc, const char *fn_name,
                   const char *label) {
    spl_diag_frame_t f = {0};
    if (loc)
        f.loc = *loc;
    f.fn_name = fn_name ? diag_strdup(fn_name) : NULL;
    f.label = label ? diag_strdup(label) : NULL;
    vec_push(d->ctx, f);
}

void spl_diag_pop(spl_diag_t *d) {
    if (vec_size(d->ctx) == 0)
        return;
    frame_free(&vec_at(d->ctx, vec_size(d->ctx) - 1));
    d->ctx.size--;
}

usize spl_diag_depth(const spl_diag_t *d) { return vec_size(d->ctx); }

static spl_diag_frame_vec_t ctx_snapshot(const spl_diag_t *d) {
    spl_diag_frame_vec_t snap;
    vec_init(snap);
    for (usize i = 0; i < vec_size(d->ctx); ++i) {
        spl_diag_frame_t *src = &vec_at(d->ctx, i);
        spl_diag_frame_t f = {0};
        f.loc = src->loc;
        f.fn_name = src->fn_name ? diag_strdup(src->fn_name) : NULL;
        f.label = src->label ? diag_strdup(src->label) : NULL;
        vec_push(snap, f);
    }
    return snap;
}

static void frame_desc(const spl_diag_frame_t *f, char *buf, usize cap) {
    if (f->fn_name && f->label)
        snprintf(buf, cap, "%s %s", f->label, f->fn_name);
    else if (f->fn_name)
        snprintf(buf, cap, "%s", f->fn_name);
    else if (f->label)
        snprintf(buf, cap, "%s", f->label);
    else
        snprintf(buf, cap, "?");
}

static void print_loc(FILE *out, const spl_dbg_node_t *loc) {
    if (loc && loc->fname)
        fprintf(out, "%s:%d:%d", loc->fname, loc->line, loc->col);
    else
        fprintf(out, "<no-loc>");
}

static void print_record(const spl_diag_t *d, const spl_diag_record_t *r) {
    FILE *out = (FILE *)d->fout;
    if (!out)
        return;
    const char *col = d->color ? level_color(r->level) : "";
    const char *rst = d->color ? ANSI_NONE : "";

    fprintf(out, "%s%s[%s:%s]%s: ", col, spl_diag_level_name(r->level),
            spl_diag_stage_name(r->stage), spl_diag_kind_name(r->kind), rst);
    print_loc(out, &r->loc);
    fputc('\n', out);

    fprintf(out, "    %s\n", r->message ? r->message : "");

    for (usize i = vec_size(r->ctx); i > 0; --i) {
        char buf[256];
        frame_desc(&vec_at(r->ctx, i - 1), buf, sizeof buf);
        fprintf(out, "    in %s at ", buf);
        print_loc(out, &vec_at(r->ctx, i - 1).loc);
        fputc('\n', out);
    }

    for (usize i = 0; i < vec_size(r->notes); ++i) {
        spl_diag_note_t *n = &vec_at(r->notes, i);
        const char *ncol = d->color ? ANSI_FG_CYAN : "";
        fprintf(out, "%snote%s: ", ncol, rst);
        print_loc(out, &n->loc);
        fputc('\n', out);
        fprintf(out, "    %s\n", n->message ? n->message : "");
    }
}

static void emit(spl_diag_t *d, spl_diag_level_t level, spl_diag_kind_t kind,
                 const spl_dbg_node_t *loc, const char *fmt, va_list ap) {
    va_list ap2;
    va_copy(ap2, ap);
    int need = vsnprintf(NULL, 0, fmt, ap);
    if (need < 0)
        need = 0;

    char *msg = malloc((usize)need + 1);
    if (msg)
        vsnprintf(msg, (usize)need + 1, fmt, ap2);
    va_end(ap2);

    spl_diag_record_t rec = {0};
    rec.level = level;
    rec.stage = d->stage;
    rec.kind = kind;
    if (loc)
        rec.loc = *loc;
    rec.message = msg;
    rec.ctx = ctx_snapshot(d);
    vec_init(rec.notes);

    if (level == SPL_DIAG_LEVEL_ERROR)
        d->n_errors++;
    else if (level == SPL_DIAG_LEVEL_FATAL) {
        d->n_errors++;
        d->n_fatals++;
    } else if (level == SPL_DIAG_LEVEL_WARN)
        d->n_warns++;

    if (d->fout) {
        print_record(d, &rec);
        d->num_printed++;
    }

    if (d->collect)
        vec_push(d->records, rec);
    else
        record_free(&rec);
}

void spl_diag_note(spl_diag_t *d, spl_diag_kind_t kind, const spl_dbg_node_t *loc, const char *fmt,
                   ...) {
    va_list ap;
    va_start(ap, fmt);
    emit(d, SPL_DIAG_LEVEL_NOTE, kind, loc, fmt, ap);
    va_end(ap);
}

void spl_diag_warn(spl_diag_t *d, spl_diag_kind_t kind, const spl_dbg_node_t *loc, const char *fmt,
                   ...) {
    va_list ap;
    va_start(ap, fmt);
    emit(d, SPL_DIAG_LEVEL_WARN, kind, loc, fmt, ap);
    va_end(ap);
}

void spl_diag_error(spl_diag_t *d, spl_diag_kind_t kind, const spl_dbg_node_t *loc, const char *fmt,
                    ...) {
    va_list ap;
    va_start(ap, fmt);
    emit(d, SPL_DIAG_LEVEL_ERROR, kind, loc, fmt, ap);
    va_end(ap);
}

void spl_diag_fatal(spl_diag_t *d, spl_diag_kind_t kind, const spl_dbg_node_t *loc, const char *fmt,
                    ...) {
    va_list ap;
    va_start(ap, fmt);
    emit(d, SPL_DIAG_LEVEL_FATAL, kind, loc, fmt, ap);
    va_end(ap);
}

void spl_diag_attach_note(spl_diag_t *d, const spl_dbg_node_t *loc, const char *fmt, ...) {
    if (vec_size(d->records) == 0)
        return;
    spl_diag_record_t *last = &vec_at(d->records, vec_size(d->records) - 1);

    va_list ap;
    va_start(ap, fmt);
    int need = vsnprintf(NULL, 0, fmt, ap);
    va_end(ap);
    if (need < 0)
        need = 0;

    char *msg = malloc((usize)need + 1);
    if (msg) {
        va_start(ap, fmt);
        vsnprintf(msg, (usize)need + 1, fmt, ap);
        va_end(ap);
    }

    spl_diag_note_t n = {0};
    if (loc)
        n.loc = *loc;
    n.message = msg;
    vec_push(last->notes, n);
}

bool spl_diag_has_error(const spl_diag_t *d) { return d->n_errors > 0; }

bool spl_diag_has_fatal(const spl_diag_t *d) { return d->n_fatals > 0; }

usize spl_diag_error_count(const spl_diag_t *d) { return d->n_errors; }

usize spl_diag_warn_count(const spl_diag_t *d) { return d->n_warns; }

void spl_diag_summary(const spl_diag_t *d) {
    FILE *out = (FILE *)d->fout;
    if (!out)
        return;

    /* 把还没实时打印过的记录补打一遍。
       collect == false 时 records 是空的，num_printed 可能非零，
       用 records.size 兜底避免越界。 */
    usize printed_limit = d->num_printed;
    if (printed_limit > vec_size(d->records))
        printed_limit = vec_size(d->records);

    for (usize i = printed_limit; i < vec_size(d->records); ++i)
        print_record(d, &vec_at(d->records, i));

    /* 计数 */
    const char *ecol = d->color && d->n_errors > 0 ? ANSI_FG_RED : "";
    const char *rst = d->color ? ANSI_NONE : "";
    fprintf(out, "%s%zu error%s, %zu warning%s%s\n", ecol, d->n_errors, d->n_errors == 1 ? "" : "s",
            d->n_warns, d->n_warns == 1 ? "" : "s", rst);
}