#include "render_ansi.h"
#include "ansi_style.h"
#include "block_scan.h"
#include "table_stats.h"
#include "fsm.h"
#include <stdio.h>
#include <string.h>

typedef struct {
    uint8_t bold;
    uint8_t italic;
    uint8_t code;
    uint8_t latex;
    uint8_t strike;
} ansi_inline_t;

typedef struct {
    uint8_t ul;
    uint8_t ol;
    uint8_t bq;
    uint8_t para;
    int     ol_num;
} ansi_env_t;

static void ansi_s(sashidown_ctx_t *ctx, const char *s)
{
    if (s && s[0])
        fsm_append_str(ctx, s);
}

static void ansi_close_lists(ansi_env_t *env, sashidown_ctx_t *ctx)
{
    (void)ctx;
    env->ul = 0;
    env->ol = 0;
    env->ol_num = 0;
}

static void ansi_close_para(ansi_env_t *env, sashidown_ctx_t *ctx)
{
    if (env->para) {
        fsm_append(ctx, '\n');
        env->para = 0;
    }
}

static void ansi_close_all(ansi_env_t *env, sashidown_ctx_t *ctx)
{
    ansi_close_lists(env, ctx);
    ansi_close_para(env, ctx);
    env->bq = 0;
}

static int ansi_render_link(sashidown_ctx_t *ctx, const ansi_theme_t *th,
                            const char *input, size_t len,
                            size_t *pos, int is_image)
{
    size_t i = *pos + 1;
    if (is_image) {
        if (i >= len || input[i] != '[')
            return 0;
        i++;
    }
    size_t text_start = i;
    i = block_scan_until_byte(input, len, i, ']');
    if (i >= len || input[i] != ']') return 0;
    size_t text_end = i;
    i++;
    if (i >= len || input[i] != '(') return 0;
    i++;
    size_t url_start = i;
    i = block_scan_until_byte(input, len, i, ')');
    if (i >= len) return 0;

    if (is_image) {
        ansi_s(ctx, th->image_on);
        fsm_append_lit(ctx, "[img:");
        fsm_append_len(ctx, input + text_start, text_end - text_start);
        fsm_append_lit(ctx, "](");
        ansi_s(ctx, th->url_on);
        fsm_append_len(ctx, input + url_start, i - url_start);
        ansi_s(ctx, th->url_off);
        fsm_append(ctx, ')');
        ansi_s(ctx, th->image_off);
    } else {
        ansi_s(ctx, th->link_on);
        fsm_append_len(ctx, input + text_start, text_end - text_start);
        ansi_s(ctx, th->link_off);
        fsm_append_lit(ctx, " (");
        ansi_s(ctx, th->url_on);
        fsm_append_len(ctx, input + url_start, i - url_start);
        ansi_s(ctx, th->url_off);
        fsm_append(ctx, ')');
    }
    *pos = i + 1;
    return 1;
}

static void ansi_render_inline(sashidown_ctx_t *ctx, const ansi_theme_t *th,
                               const char *input, size_t len,
                               size_t start, size_t end)
{
    ansi_inline_t st = {0, 0, 0, 0, 0};
    size_t i = start;

    while (i < end) {
        if (st.latex) {
            size_t j = block_scan_until_byte(input, end, i, '$');
            if (j > i)
                fsm_append_len(ctx, input + i, j - i);
            if (j < end) {
                fsm_append(ctx, '$');
                ansi_s(ctx, th->latex_off);
                st.latex = 0;
                i = j + 1;
            } else {
                i = j;
            }
            continue;
        }

        if (st.code) {
            size_t j = block_scan_until_byte(input, end, i, '`');
            if (j > i)
                fsm_append_len(ctx, input + i, j - i);
            if (j < end) {
                ansi_s(ctx, th->code_off);
                st.code = 0;
                i = j + 1;
            } else {
                i = j;
            }
            continue;
        }

        unsigned char ch = (unsigned char)input[i];

        if ((ch & 0x80) == 0) {
            size_t j = block_scan_plain_ascii(input, end, i);
            if (j > i) {
                fsm_append_len(ctx, input + i, j - i);
                i = j;
                if (i >= end)
                    break;
                ch = (unsigned char)input[i];
            }
        } else {
            size_t j = block_scan_utf8_run(input, end, i);
            if (j > i) {
                fsm_append_len(ctx, input + i, j - i);
                i = j;
                continue;
            }
        }

        if (ch == '$' && i + 1 < end && input[i + 1] == '$') {
            size_t j = block_scan_until_2byte(input, end, i + 2, '$', '$');
            if (j < end)
                j += 2;
            ansi_s(ctx, th->latex_on);
            fsm_append_len(ctx, input + i, j - i);
            ansi_s(ctx, th->latex_off);
            i = j;
            continue;
        }

        if (ch == '$') {
            ansi_s(ctx, th->latex_on);
            fsm_append(ctx, '$');
            st.latex = 1;
            i++;
            continue;
        }

        if (ch == '!' && i + 1 < end && input[i + 1] == '[') {
            if (ansi_render_link(ctx, th, input, len, &i, 1))
                continue;
        }

        if (ch == '[') {
            if (ansi_render_link(ctx, th, input, len, &i, 0))
                continue;
        }

        if (ch == '\\' && i + 1 < end) {
            unsigned char next = (unsigned char)input[i + 1];
            if ((next >= 0x21 && next <= 0x2F) ||
                (next >= 0x3A && next <= 0x40) ||
                (next >= 0x5B && next <= 0x60) ||
                (next >= 0x7B && next <= 0x7E)) {
                fsm_append(ctx, (char)next);
                i += 2;
                continue;
            }
        }

        if ((ch == '*' || ch == '_') && i + 1 < end && input[i + 1] == (char)ch) {
            ansi_s(ctx, st.bold ? th->strong_off : th->strong_on);
            st.bold = !st.bold;
            i += 2;
            continue;
        }

        if (ch == '~' && i + 1 < end && input[i + 1] == '~') {
            ansi_s(ctx, st.strike ? th->strike_off : th->strike_on);
            st.strike = !st.strike;
            i += 2;
            continue;
        }

        if (ch == '*' || ch == '_') {
            ansi_s(ctx, st.italic ? th->emph_off : th->emph_on);
            st.italic = !st.italic;
            i++;
            continue;
        }

        if (ch == '`') {
            ansi_s(ctx, th->code_on);
            st.code = 1;
            i++;
            continue;
        }

        fsm_append(ctx, (char)ch);
        i++;
    }

    if (st.bold)   ansi_s(ctx, th->strong_off);
    if (st.italic) ansi_s(ctx, th->emph_off);
    if (st.code)   ansi_s(ctx, th->code_off);
    if (st.strike) ansi_s(ctx, th->strike_off);
    if (st.latex)  ansi_s(ctx, th->latex_off);
}

static void trim_cell(const char *input, size_t *start, size_t *end)
{
    while (*start < *end && input[*start] == ' ')
        (*start)++;
    while (*end > *start && input[*end - 1] == ' ')
        (*end)--;
}

static void ansi_table_row_at(sashidown_ctx_t *ctx, const ansi_theme_t *th,
                              const char *input, size_t len,
                              size_t pos, size_t le, int cols, int is_header)
{
    size_t i = pos;
    while (i < le && (input[i] == ' ' || input[i] == '\t')) i++;
    if (i < le && input[i] == '|')
        i++;

    for (int col = 0; col < cols; col++) {
        size_t cs = i;
        while (i < le && input[i] != '|') i++;
        size_t ce = i;
        trim_cell(input, &cs, &ce);
        if (col)
            ansi_s(ctx, th->sep);
        if (is_header)
            ansi_s(ctx, th->th_on);
        ansi_render_inline(ctx, th, input, len, cs, ce);
        if (is_header)
            ansi_s(ctx, th->th_off);
        if (i < le && input[i] == '|') i++;
    }
    fsm_append(ctx, '\n');
}

/* Scheme C (batch): header+sep commit, then body rows one-pass. Returns 1 if
 * a table was consumed (*pos_io advanced past body). */
static int ansi_try_consume_table(sashidown_ctx_t *ctx, const ansi_theme_t *th,
                                  ansi_env_t *env, const char *input, size_t len,
                                  size_t pos, size_t le, size_t *pos_io)
{
    int has_pipe = 0;
    int cols = table_count_cols(input + pos, le - pos, &has_pipe);
    if (!has_pipe)
        return 0;
    if (le >= len || input[le] != '\n')
        return 0;

    size_t next = le + 1;
    if (next >= len)
        return 0;
    size_t nle = block_line_end(ctx, input, len, next);
    if (!table_is_separator_row(input + next, nle - next)) {
        block_line_seek(ctx, pos);
        return 0;
    }

    ansi_close_all(env, ctx);
    ansi_table_row_at(ctx, th, input, len, pos, le, cols, 1);
    pos = block_skip_line(ctx, input, len, next);

    while (pos < len) {
        size_t ble = block_line_end(ctx, input, len, pos);
        int empty = (pos == ble || (pos + 1 == ble && input[pos] == '\r'));
        if (empty)
            break;
        size_t blen = ble - pos;
        if (!memchr(input + pos, '|', blen) ||
            table_is_separator_row(input + pos, blen))
            break;
        ansi_table_row_at(ctx, th, input, len, pos, ble, cols, 0);
        pos = block_skip_line(ctx, input, len, pos);
    }
    fsm_append(ctx, '\n');
    *pos_io = pos;
    return 1;
}

static void ansi_render_codeblock(sashidown_ctx_t *ctx, const ansi_theme_t *th,
                                  const char *input, size_t len,
                                  size_t *pos, int fence_len, char fence_char)
{
    *pos = block_skip_line(ctx, input, len, *pos);
    ansi_s(ctx, th->fence_on);
    while (*pos < len) {
        int fl = 0;
        char fc = 0;
        if (block_is_fence(input, len, *pos, &fl, &fc) &&
            fc == fence_char && fl >= fence_len) {
            size_t is, ie;
            block_fence_info_range(input, len, *pos, fl, fc, &is, &ie);
            if (ie <= is) {
                *pos = block_skip_line(ctx, input, len, *pos);
                break;
            }
        }
        size_t le = block_line_end(ctx, input, len, *pos);
        fsm_append_lit(ctx, "  ");
        fsm_append_len(ctx, input + *pos, le - *pos);
        fsm_append(ctx, '\n');
        *pos = block_skip_line(ctx, input, len, *pos);
    }
    ansi_s(ctx, th->fence_off);
    fsm_append(ctx, '\n');
}

static void ansi_render_latex_block(sashidown_ctx_t *ctx, const ansi_theme_t *th,
                                    const char *input, size_t len,
                                    size_t *pos)
{
    size_t i = *pos + 2;
    size_t start = i;
    while (i + 1 < len) {
        if (input[i] == '$' && input[i + 1] == '$') {
            ansi_s(ctx, th->latex_on);
            fsm_append_lit(ctx, "$$");
            if (i > start)
                fsm_append_len(ctx, input + start, i - start);
            fsm_append_lit(ctx, "$$");
            ansi_s(ctx, th->latex_off);
            fsm_append(ctx, '\n');
            i += 2;
            if (i < len && input[i] == '\n') i++;
            *pos = i;
            return;
        }
        i++;
    }
    ansi_s(ctx, th->latex_on);
    fsm_append_lit(ctx, "$$");
    if (i > start)
        fsm_append_len(ctx, input + start, (i < len ? i : len) - start);
    fsm_append_lit(ctx, "$$");
    ansi_s(ctx, th->latex_off);
    fsm_append(ctx, '\n');
    *pos = len;
}

static int sashidown_render_ansi_theme(sashidown_ctx_t *ctx, const char *input, size_t len,
                                    const ansi_theme_t *th)
{
    ansi_env_t env;
    memset(&env, 0, sizeof(env));
    size_t pos = 0;

    fsm_write_begin(ctx);

    while (pos < len) {
        size_t le = block_line_end(ctx, input, len, pos);
        if (pos == le || (pos + 1 == le && input[pos] == '\r')) {
            ansi_close_all(&env, ctx);
            if (env.bq)
                env.bq = 0;
            fsm_append(ctx, '\n');
            pos = block_skip_line(ctx, input, len, pos);
            continue;
        }

        if (memchr(input + pos, '|', le - pos) != NULL) {
            size_t newt = 0;
            if (ansi_try_consume_table(ctx, th, &env, input, len, pos, le, &newt)) {
                pos = newt;
                block_line_seek(ctx, pos);
                continue;
            }
        }

        line_info_t li = block_classify_line(ctx, input, len, pos, le);

        if (li.kind == LK_PARA) {
            line_info_t sli;
            size_t send = 0;
            int sr = block_try_setext(ctx, input, len, pos, le, len, 1, &sli, &send);
            if (sr > 0) {
                int lvl = sli.heading_level;
                if (lvl < 1) lvl = 1;
                if (lvl > 6) lvl = 6;
                ansi_close_all(&env, ctx);
                ansi_s(ctx, th->heading[lvl]);
                ansi_render_inline(ctx, th, input, len, sli.content_start, sli.content_end);
                ansi_s(ctx, th->heading_end);
                fsm_append(ctx, '\n');
                if (lvl <= 2)
                    fsm_append(ctx, '\n');
                pos = send;
                block_line_seek(ctx, pos);
                continue;
            }
        }

        switch (li.kind) {
        case LK_EMPTY:
            ansi_close_all(&env, ctx);
            if (env.bq)
                env.bq = 0;
            fsm_append(ctx, '\n');
            pos = block_skip_line(ctx, input, len, pos);
            break;

        case LK_HR:
            ansi_close_all(&env, ctx);
            fsm_append_lit(ctx, "---\n\n");
            pos = block_skip_line(ctx, input, len, pos);
            break;

        case LK_FENCE:
            ansi_close_all(&env, ctx);
            ansi_render_codeblock(ctx, th, input, len, &pos, li.fence_len, li.fence_char);
            break;

        case LK_LATEX_LINE:
            ansi_close_all(&env, ctx);
            ansi_s(ctx, th->latex_on);
            fsm_append_lit(ctx, "$$");
            fsm_append_len(ctx, input + li.content_start,
                           li.content_end - li.content_start);
            fsm_append_lit(ctx, "$$");
            ansi_s(ctx, th->latex_off);
            fsm_append(ctx, '\n');
            pos = block_skip_line(ctx, input, len, pos);
            break;

        case LK_LATEX_BLOCK:
            ansi_close_all(&env, ctx);
            ansi_render_latex_block(ctx, th, input, len, &pos);
            break;

        case LK_HEADING: {
            int lvl = li.heading_level;
            if (lvl < 1) lvl = 1;
            if (lvl > 6) lvl = 6;
            ansi_close_all(&env, ctx);
            ansi_s(ctx, th->heading[lvl]);
            ansi_render_inline(ctx, th, input, len, li.content_start, le);
            ansi_s(ctx, th->heading_end);
            fsm_append(ctx, '\n');
            if (lvl <= 2)
                fsm_append(ctx, '\n');
            pos = block_skip_line(ctx, input, len, pos);
            break;
        }

        case LK_BQ:
            ansi_close_lists(&env, ctx);
            env.bq = 1;
            ansi_s(ctx, th->bq_prefix);
            ansi_render_inline(ctx, th, input, len, li.content_start, le);
            fsm_append(ctx, '\n');
            pos = block_skip_line(ctx, input, len, pos);
            break;

        case LK_UL:
            if (env.ol) { env.ol = 0; env.ol_num = 0; }
            env.ul = 1;
            ansi_s(ctx, th->bullet);
            if (li.task == 1)
                fsm_append_lit(ctx, "[ ] ");
            else if (li.task == 2)
                fsm_append_lit(ctx, "[x] ");
            ansi_render_inline(ctx, th, input, len, li.content_start, le);
            ansi_s(ctx, th->bullet_end);
            fsm_append(ctx, '\n');
            pos = block_skip_line(ctx, input, len, pos);
            break;

        case LK_OL: {
            int num = 0;
            size_t cs = li.content_start;
            block_is_ol(input, len, pos, &cs, &num);
            if (env.ul) env.ul = 0;
            env.ol = 1;
            char buf[32];
            int n = snprintf(buf, sizeof(buf), "%d. ", num > 0 ? num : ++env.ol_num);
            ansi_s(ctx, th->ol_on);
            if (n > 0)
                fsm_append_len(ctx, buf, (size_t)n);
            ansi_s(ctx, th->ol_off);
            ansi_render_inline(ctx, th, input, len, li.content_start, le);
            fsm_append(ctx, '\n');
            pos = block_skip_line(ctx, input, len, pos);
            break;
        }

        case LK_PARA:
        default:
            if (env.bq) env.bq = 0;
            ansi_close_lists(&env, ctx);
            if (!env.para) {
                env.para = 1;
            } else {
                fsm_append(ctx, '\n');
            }
            ansi_render_inline(ctx, th, input, len, pos, le);
            pos = block_skip_line(ctx, input, len, pos);
            if (pos >= len) {
                ansi_close_para(&env, ctx);
            } else {
                size_t nle = block_line_end(ctx, input, len, pos);
                line_info_t nli = block_classify_line(ctx, input, len, pos, nle);
                if (nli.kind != LK_PARA)
                    ansi_close_para(&env, ctx);
            }
            break;
        }

        if (ctx->overflow) {
            fsm_write_end(ctx);
            return -1;
        }
    }

    ansi_close_all(&env, ctx);
    fsm_write_end(ctx);
    return ctx->overflow ? -1 : 0;
}

int sashidown_render_ansi(sashidown_ctx_t *ctx, const char *input, size_t len)
{
    return sashidown_render_ansi_theme(ctx, input, len, ansi_theme_color());
}

int sashidown_render_ansi_plain(sashidown_ctx_t *ctx, const char *input, size_t len)
{
    return sashidown_render_ansi_theme(ctx, input, len, ansi_theme_plain());
}