#include "sashidown.h"
#include "line_index.h"
#include "fsm.h"
#include "render.h"
#include <stdlib.h>
#include <string.h>

#define STREAM_INPUT_INIT      4096
#define STREAM_OUTPUT_INIT     8192
#define STREAM_OUTPUT_MAX      (64u * 1024u * 1024u)
#define STREAM_LINE_IX_INIT    256

struct sashidown_stream {
    char                  *input_buf;
    size_t                 input_len;
    size_t                 input_cap;
    size_t                 indexed_len;      /* bytes already in line_ix */
    size_t                 render_pos;
    size_t                 last_complete_nl; /* [0, last_complete_nl) ends on line boundary */
    sashidown_render_state_t  render_state;
    line_index_t           line_ix;
    char                  *output_buf;
    size_t                 output_cap;
    sashidown_ctx_t           render_ctx;
    sashidown_stream_cb       on_html;
    void                  *user;
    int                    finished;
    int                    aborted;
};

static int stream_grow_input(sashidown_stream_t *s, size_t need)
{
    if (need <= s->input_cap)
        return 0;
    size_t cap = s->input_cap ? s->input_cap : STREAM_INPUT_INIT;
    while (cap < need)
        cap *= 2;
    char *p = (char *)realloc(s->input_buf, cap);
    if (!p)
        return -1;
    s->input_buf = p;
    s->input_cap = cap;
    return 0;
}

static int stream_grow_output(sashidown_stream_t *s, size_t need)
{
    if (need <= s->output_cap)
        return 0;
    size_t cap = s->output_cap ? s->output_cap : STREAM_OUTPUT_INIT;
    while (cap < need)
        cap *= 2;
    if (cap > STREAM_OUTPUT_MAX)
        cap = STREAM_OUTPUT_MAX;
    if (cap <= s->output_cap)
        return -1;
    char *p = (char *)realloc(s->output_buf, cap);
    if (!p)
        return -1;
    s->output_buf = p;
    s->output_cap = cap;
    return 0;
}

static int stream_flush_html(sashidown_stream_t *s)
{
    sashidown_ctx_t *ctx = &s->render_ctx;
    if (ctx->output_len == 0)
        return 0;
    if (!s->on_html)
        return 0;
    int r = s->on_html(ctx->output, ctx->output_len, s->user);
    if (r != 0) {
        s->aborted = 1;
        return -1;
    }
    ctx->output_len = 0;
    return 0;
}

/* Sync line index through input_len; update last_complete_nl. */
static void stream_catchup_index(sashidown_stream_t *s)
{
    if (s->indexed_len >= s->input_len)
        return;

    if (s->indexed_len == 0) {
        line_index_build(s->input_buf, s->input_len, &s->line_ix);
        size_t last = 0;
        if (s->line_ix.valid && s->line_ix.n > 1) {
            last = (size_t)s->line_ix.off[s->line_ix.n - 1];
        } else {
            const char *p = s->input_buf;
            const char *end = s->input_buf + s->input_len;
            const char *hit;
            while ((hit = memchr(p, '\n', (size_t)(end - p))) != NULL) {
                last = (size_t)(hit - s->input_buf) + 1;
                p = hit + 1;
            }
        }
        if (last)
            s->last_complete_nl = last;
    } else {
        size_t nl = line_index_extend(&s->line_ix, s->input_buf,
                                      s->indexed_len, s->input_len);
        if (nl)
            s->last_complete_nl = nl;
    }
    s->indexed_len = s->input_len;
}

/* 1 = overflow (state restored), 0 = ok, -1 = error */
static int stream_render_attempt(sashidown_stream_t *s, size_t stop, int is_final)
{
    sashidown_ctx_t *ctx = &s->render_ctx;
    size_t saved_pos = s->render_pos;
    sashidown_render_state_t saved_state = s->render_state;

    ctx->output_len = 0;
    ctx->overflow = 0;
    fsm_write_begin(ctx);
    int r = sashidown_render_incremental(ctx, &s->render_state,
                                      s->input_buf, s->input_len,
                                      &s->render_pos, stop, is_final);
    fsm_write_end(ctx);

    if (ctx->overflow) {
        s->render_pos = saved_pos;
        s->render_state = saved_state;
        return 1;
    }
    if (r != 0)
        return -1;
    return stream_flush_html(s);
}

static size_t stream_soft_stop(sashidown_stream_t *s, size_t from, size_t target)
{
    if (from >= target)
        return target;

    /* Keep each attempt roughly within half the output buffer. */
    size_t soft = from + s->output_cap / 2;
    if (soft >= target)
        return target;

    const char *base = s->input_buf + from;
    size_t lim = soft - from;
    const char *p = base;
    const char *end = base + lim;
    const char *last = NULL;
    const char *hit;
    while ((hit = memchr(p, '\n', (size_t)(end - p))) != NULL) {
        last = hit;
        p = hit + 1;
    }
    if (last)
        return (size_t)(last - s->input_buf) + 1;

    /* No newline in soft window — take the next one so we still make progress. */
    hit = memchr(s->input_buf + from, '\n', target - from);
    if (hit)
        return (size_t)(hit - s->input_buf) + 1;
    return target;
}

static size_t stream_next_line_stop(sashidown_stream_t *s, size_t from, size_t limit)
{
    if (from >= limit)
        return limit;
    const char *hit = memchr(s->input_buf + from, '\n', limit - from);
    if (hit)
        return (size_t)(hit - s->input_buf) + 1;
    return limit;
}

/* Render up to pass_stop (is_final as given). Grow/shrink on overflow.
 * Returns 0 on success, -1 on error. */
static int stream_render_window(sashidown_stream_t *s, size_t pass_stop, int is_final)
{
    sashidown_ctx_t *ctx = &s->render_ctx;
    for (;;) {
        int ar = stream_render_attempt(s, pass_stop, is_final);
        if (ar == 0)
            return 0;
        if (ar < 0)
            return -1;
        if (stream_grow_output(s, s->output_cap * 2) == 0) {
            ctx->output = s->output_buf;
            ctx->output_cap = s->output_cap;
            continue;
        }
        size_t one = stream_next_line_stop(s, s->render_pos, pass_stop);
        if (one <= s->render_pos || one >= pass_stop)
            return -1;
        pass_stop = one;
        is_final = 0;
    }
}

static int stream_render_pass(sashidown_stream_t *s, size_t stop, int is_final)
{
    if (s->aborted)
        return -1;
    if (!is_final && stop <= s->render_pos)
        return 0;

    sashidown_ctx_t *ctx = &s->render_ctx;
    ctx->input = s->input_buf;
    ctx->input_len = s->input_len;

    size_t target = is_final ? s->input_len : stop;

    /* Common case: one or a few lines already fit in the output buffer. */
    if (target - s->render_pos <= s->output_cap / 4) {
        int pass_final = is_final && (target >= s->input_len);
        size_t before = s->render_pos;
        if (stream_render_window(s, target, pass_final) != 0)
            return -1;
        if (s->render_pos > before || pass_final)
            return (is_final && !pass_final)
                       ? stream_render_window(s, s->input_len, 1)
                       : 0;
        if (!is_final)
            return 0; /* deferred (incomplete table, etc.) */
        /* fall through to spanned finalization */
    }

    while (s->render_pos < target) {
        size_t before = s->render_pos;
        size_t pass_stop = stream_soft_stop(s, s->render_pos, target);
        int pass_final = is_final && (pass_stop >= s->input_len);

        if (stream_render_window(s, pass_stop, pass_final) != 0)
            return -1;

        if (s->render_pos > before) {
            if (pass_final)
                return 0;
            continue;
        }

        /* No progress: soft window may have deferred a table. Retry full target. */
        if (pass_stop < target) {
            pass_final = is_final && (target >= s->input_len);
            if (stream_render_window(s, target, pass_final) != 0)
                return -1;
            if (s->render_pos > before) {
                if (pass_final)
                    return 0;
                continue;
            }
        }

        /* Still stuck — waiting for more input (incomplete table, etc.). */
        if (!is_final)
            return 0;
        return -1;
    }

    if (!is_final)
        return 0;

    /* EOF reached without a finalizing pass — close open blocks. */
    return stream_render_window(s, s->input_len, 1);
}

sashidown_stream_t *sashidown_stream_create(const sashidown_stream_opts_t *opts)
{
    if (!opts || !opts->on_html)
        return NULL;

    sashidown_stream_t *s = (sashidown_stream_t *)calloc(1, sizeof(*s));
    if (!s)
        return NULL;

    size_t out_cap = opts->output_buf_size ? opts->output_buf_size : STREAM_OUTPUT_INIT;
    if (out_cap > STREAM_OUTPUT_MAX)
        out_cap = STREAM_OUTPUT_MAX;

    s->output_cap = out_cap;
    s->output_buf = (char *)malloc(s->output_cap);
    if (!s->output_buf) {
        free(s);
        return NULL;
    }

    if (line_index_init_dynamic(&s->line_ix, STREAM_LINE_IX_INIT) != 0) {
        free(s->output_buf);
        free(s);
        return NULL;
    }

    s->on_html = opts->on_html;
    s->user = opts->user;
    sashidown_render_state_init(&s->render_state);

    s->render_ctx.output = s->output_buf;
    s->render_ctx.output_cap = s->output_cap;
    s->render_ctx.line_ix = &s->line_ix;

    return s;
}

int sashidown_stream_feed(sashidown_stream_t *s, const char *chunk, size_t len)
{
    if (!s || s->aborted)
        return -1;
    if (s->finished)
        return 0;
    if (!chunk && len > 0)
        return -1;
    if (len == 0)
        return 0;

    int has_nl = (memchr(chunk, '\n', len) != NULL);

    if (stream_grow_input(s, s->input_len + len) != 0)
        return -1;

    memcpy(s->input_buf + s->input_len, chunk, len);
    s->input_len += len;

    /* Tiny-chunk fast path: no newline → defer index + render. */
    if (!has_nl)
        return 0;

    stream_catchup_index(s);

    if (s->last_complete_nl <= s->render_pos)
        return 0;

    return stream_render_pass(s, s->last_complete_nl, 0);
}

int sashidown_stream_finish(sashidown_stream_t *s)
{
    if (!s || s->aborted)
        return -1;
    if (s->finished)
        return 0;

    stream_catchup_index(s);
    int r = stream_render_pass(s, s->input_len, 1);
    s->finished = 1;
    return r;
}

void sashidown_stream_destroy(sashidown_stream_t *s)
{
    if (!s)
        return;
    free(s->input_buf);
    free(s->output_buf);
    line_index_free(&s->line_ix);
    free(s);
}

size_t sashidown_stream_pending(const sashidown_stream_t *s)
{
    if (!s || s->input_len <= s->render_pos)
        return 0;
    return s->input_len - s->render_pos;
}