#include "compat.h"
#include "screen.h"
#include "unicode/grapheme.h"
#include <stdlib.h>
#include <string.h>

/*
 * A snapshot keeps one converted copy of each viewport row. An update
 * copies a row again only when the terminal reports a new key for it.
 * Cluster cells refer to an arena of [length, code points...] records that
 * the snapshot owns, so a renderer never reads the terminal's cluster table.
 */
struct term_snapshot {
    struct term_view view;
    uint32_t palette[TERM_PALETTE_SIZE];
    uint64_t source;                /* serial of the terminal it copied */

    struct term_cell *cells;        /* rows * cols */
    uint64_t *keys;                 /* per row */
    uint32_t *row_clusters;         /* arena records each row holds */
    struct term_span *selection;    /* per row */
    size_t cell_cap;
    int row_cap;

    uint32_t *arena;
    uint32_t arena_len, arena_cap, arena_dead;
};

struct term_snapshot *term_snapshot_new(void)
{
    struct term_snapshot *s = calloc(1, sizeof(*s));
    if (s)
        s->view.palette = s->palette;
    return s;
}

void term_snapshot_free(struct term_snapshot *s)
{
    if (!s)
        return;
    free(s->cells);
    free(s->keys);
    free(s->row_clusters);
    free(s->selection);
    free(s->arena);
    free(s);
}

/* Make room for cols x rows. A failed call keeps the old contents, so the
 * snapshot stays readable. */
static bool reserve(struct term_snapshot *s, int cols, int rows)
{
    size_t cells = (size_t)cols * (size_t)rows;
    if (cells > s->cell_cap) {
        struct term_cell *c = realloc(s->cells, cells * sizeof(*c));
        if (!c)
            return false;
        s->cells = c;
        s->cell_cap = cells;
    }
    if (rows > s->row_cap) {
        uint64_t *k = realloc(s->keys, (size_t)rows * sizeof(*k));
        if (k)
            s->keys = k;
        uint32_t *rc = realloc(s->row_clusters, (size_t)rows * sizeof(*rc));
        if (rc)
            s->row_clusters = rc;
        struct term_span *sel = realloc(s->selection, (size_t)rows * sizeof(*sel));
        if (sel)
            s->selection = sel;
        if (!k || !rc || !sel)
            return false;
        s->row_cap = rows;
    }
    return true;
}

/* Append a cluster record. The caller reserved room for it. */
static uint32_t arena_add(struct term_snapshot *s, const struct gstore *g, uint32_t handle)
{
    uint32_t cps[TERM_GRAPHEME_MAX_LEN];
    int n = term_gstore_expand(g, handle, cps, TERM_GRAPHEME_MAX_LEN);
    if (n <= 0)
        return ' ';
    uint32_t off = s->arena_len;
    s->arena[off] = (uint32_t)n;
    memcpy(&s->arena[off + 1], cps, (size_t)n * sizeof(*cps));
    s->arena_len += (uint32_t)n + 1;
    return TERM_TEXT_CLUSTER | off;
}

/* Arena words that converting src adds. Counts what arena_add() writes. */
static uint64_t row_words(const struct gstore *g, const struct cell *src, int cols)
{
    uint32_t cps[TERM_GRAPHEME_MAX_LEN];
    uint64_t words = 0;
    for (int x = 0; x < cols; x++) {
        uint32_t text = src[x].c;
        if (!term_grapheme_is_handle(text))
            continue;
        int n = term_gstore_expand(g, text, cps, TERM_GRAPHEME_MAX_LEN);
        if (n > 0)
            words += (uint64_t)n + 1;
    }
    return words;
}

/* Convert one stored row; returns the number of arena words it added. */
static uint32_t convert_row(struct term_snapshot *s, const struct gstore *g,
                            struct term_cell *dst,
                            const struct cell *src, int cols)
{
    uint32_t before = s->arena_len;
    for (int x = 0; x < cols; x++) {
        const struct cell *c = &src[x];
        uint16_t f = c->flags;
        uint32_t text = c->c;
        if (TERM_EXPECT(text < 0x20, 0))
            text = text ? ' ' : 0;
        else if (TERM_EXPECT(term_grapheme_is_handle(text), 0))
            text = arena_add(s, g, text);
        dst[x] = (struct term_cell){
            .text = text,
            .fg = c->fg,
            .bg = c->bg,
            .attrs = (uint16_t)(f & ~(CELL_WIDE | CELL_TRAIL | CELL_WRAP_PAD)),
            .width = (f & CELL_TRAIL) ? 0 : (f & CELL_WIDE) ? 2 : 1,
        };
    }
    return s->arena_len - before;
}

/* What one update does. Phase 1 fills it in and changes no snapshot field. */
struct plan {
    bool all;                   /* rebuild every row and reset the arena */
    int scrolled;               /* rows the content moved up */
    const uint64_t *keys;       /* key of each new row */
    const int *c0, *c1;         /* selected columns of each row */
};

/* Phase 2: write the plan into the snapshot. Nothing here can fail. */
static void apply(struct term_snapshot *s, const struct term *t, const struct plan *p)
{
    int cols = t->cols, rows = t->rows;
    int k = p->scrolled;
    if (p->all) {
        s->arena_len = 0;
        s->arena_dead = 0;
    }

    bool changed = p->all;
    if (k > 0) {
        for (int y = 0; y < k; y++)
            s->arena_dead += s->row_clusters[y];
        memmove(s->cells, s->cells + (size_t)k * cols,
                (size_t)(rows - k) * cols * sizeof(struct term_cell));
        memmove(s->row_clusters, s->row_clusters + k,
                (size_t)(rows - k) * sizeof(uint32_t));
        memmove(s->keys, s->keys + k,
                (size_t)(rows - k) * sizeof(uint64_t));
        memset(s->keys + (rows - k), 0, (size_t)k * sizeof(uint64_t));
        memset(s->row_clusters + (rows - k), 0, (size_t)k * sizeof(uint32_t));
    }

    for (int y = 0; y < rows; y++) {
        uint64_t key = p->keys[y];
        if (!p->all && s->keys[y] == key)
            continue;
        if (!p->all)
            s->arena_dead += s->row_clusters[y];
        const struct cell *src = term_screen_view_row(t, y);
        s->row_clusters[y] = convert_row(s, t->gstore, s->cells + (size_t)y * cols,
                                         src, cols);
        s->keys[y] = key;
        changed = true;
    }

    for (int y = 0; y < rows; y++) {
        struct term_span span = p->c1[y] >= p->c0[y] ? (struct term_span){ p->c0[y], p->c1[y] + 1 }
                                                     : (struct term_span){ 0 };
        if (p->all || span.begin != s->selection[y].begin || span.end != s->selection[y].end)
            changed = true;
        s->selection[y] = span;
    }

    struct term_cursor cursor = term_cursor(t);
    int view_offset = term_view_offset(t);
    cursor.pos.y += view_offset;
    if (cursor.pos.y >= rows)
        cursor.visible = false;
    const struct term_cursor *old = &s->view.cursor;
    if (cursor.pos.x != old->pos.x || cursor.pos.y != old->pos.y ||
        cursor.style != old->style || cursor.visible != old->visible ||
        cursor.blink != old->blink ||
        s->view.scroll_offset != view_offset ||
        s->view.reverse_video != t->reverse_video)
        changed = true;
    if (p->all || s->view.palette_version != t->palette.version) {
        memcpy(s->palette, t->palette.colors, sizeof(s->palette));
        s->view.palette_version = t->palette.version;
        changed = true;
    }

    s->source = t->serial;
    s->view.size = (struct term_vector2i){ .x = cols, .y = rows };
    s->view.cursor = cursor;
    s->view.scroll_offset = view_offset;
    s->view.scrolled = k;
    s->view.reverse_video = t->reverse_video;
    if (changed || k)
        s->view.version++;
}

/* Phase 1 for one update. Allocates and checks everything that can fail,
 * then hands the plan to apply(). A false return leaves the content as it was. */
static bool term_snapshot_update_inner(struct term_snapshot *s, struct term *t)
{
    if (!s || !t)
        return false;
    int cols = t->cols, rows = t->rows;
    bool ok = false;
    bool resized = s->view.size.x != cols || s->view.size.y != rows;
    bool all = resized || s->source != t->serial;
    int scrolled = 0;
    uint64_t need = 0;

    if (!reserve(s, cols, rows))
        return false;
    /* Scratch for rows above 256: new keys, then the selection spans. */
    uint64_t stack_keys[256];
    int stack[2][256];
    uint64_t *new_keys = stack_keys;
    int *c0 = stack[0], *c1 = stack[1];
    void *heap = NULL;
    if (rows > 256) {
        heap = malloc((size_t)rows * (sizeof(*new_keys) + 2 * sizeof(*c0)));
        if (!heap)
            return false;
        new_keys = heap;
        c0 = (int *)(new_keys + rows);
        c1 = c0 + rows;
    }
    struct plan p = { .keys = new_keys, .c0 = c0, .c1 = c1 };

    /* Reclaim cluster records once most of the arena is garbage. */
    if (s->arena_dead > 4096 && s->arena_dead * 2 > s->arena_len)
        all = true;

    term_hist_prepare_rows(&t->hist, cols);
    if (!term_hist_view_sync(&t->view, &t->hist, cols))
        goto done;
    term_hist_view_refine(&t->view, &t->hist, cols, rows);
    for (int y = 0; y < rows; y++)
        new_keys[y] = term_screen_view_row_key(t, y);

    if (!all && rows > 1) {
        int cand_k = 0;
        /* Find candidate scroll offset k by checking where top rows moved. */
        uint64_t key0 = new_keys[0];
        if (key0) {
            for (int k = 1; k < rows; k++) {
                if (s->keys[k] == key0) {
                    cand_k = k;
                    break;
                }
            }
        }
        if (!cand_k && rows > 2) {
            uint64_t key1 = new_keys[1];
            if (key1) {
                for (int k = 1; 1 + k < rows; k++) {
                    if (s->keys[1 + k] == key1) {
                        cand_k = k;
                        break;
                    }
                }
            }
        }

        if (cand_k > 0) {
            int total = rows - cand_k;
            int threshold = total / 2;
            int matches = 0, mismatches = 0;
            for (int y = 0; y < total; y++) {
                if (new_keys[y] == s->keys[y + cand_k]) {
                    if (++matches > threshold) {
                        scrolled = cand_k;
                        break;
                    }
                } else {
                    if (++mismatches > threshold)
                        break;
                }
            }
        }
    }

    /* Arena words for the rows that apply() converts. Count them before
     * anything changes, then grow the arena once so apply() cannot fail. */
    need = all ? 0 : s->arena_len;
    for (int y = 0; y < rows; y++) {
        uint64_t kept = s->keys[y];
        if (scrolled)
            kept = y + scrolled < rows ? s->keys[y + scrolled] : 0;
        if (!all && kept == new_keys[y])
            continue;
        need += row_words(t->gstore, term_screen_view_row(t, y), cols);
    }
    /* Offsets must fit below TERM_TEXT_CLUSTER. */
    if (need > TERM_TEXT_CLUSTER)
        goto done;
    if (need > s->arena_cap) {
        uint32_t cap = term_grow_cap32(s->arena_cap ? s->arena_cap * 2 : 256, (uint32_t)need);
        uint32_t *a = realloc(s->arena, cap * sizeof(*a));
        if (!a)
            goto done;
        s->arena = a;
        s->arena_cap = cap;
    }

    if (!term_selection_view_spans(t, c0, c1))
        goto done;
    p.all = all;
    p.scrolled = scrolled;
    apply(s, t, &p);
    ok = true;

done:
    free(heap);
    return ok;
}

bool term_snapshot_update(struct term_snapshot *s, struct term *t)
{
    if (!s || !t)
        return false;
    term_write_begin(t);
    bool ok = term_snapshot_update_inner(s, t);
    term_write_end(t);
    return ok;
}

const struct term_view *term_snapshot_view(const struct term_snapshot *s)
{
    return s ? &s->view : NULL;
}

const struct term_cell *term_snapshot_row(const struct term_snapshot *s, int y)
{
    if (!s || y < 0 || y >= s->view.size.y)
        return NULL;
    return s->cells + (size_t)y * s->view.size.x;
}

uint64_t term_snapshot_row_key(const struct term_snapshot *s, int y)
{
    if (!s || y < 0 || y >= s->view.size.y)
        return 0;
    return s->keys[y];
}

struct term_span term_snapshot_row_selection(const struct term_snapshot *s, int y)
{
    if (!s || y < 0 || y >= s->view.size.y)
        return (struct term_span){ 0 };
    return s->selection[y];
}

int term_snapshot_cluster(const struct term_snapshot *s, uint32_t text, uint32_t *cps, int cap)
{
    if (!s)
        return 0;
    if (!(text & TERM_TEXT_CLUSTER)) {
        if (cps && cap > 0)
            cps[0] = text;
        return 1;
    }
    uint32_t off = text & ~TERM_TEXT_CLUSTER;
    if (off >= s->arena_len)
        return 0;
    int n = (int)s->arena[off];
    int copy = n < cap ? n : cap;
    if (cps && copy > 0)
        memcpy(cps, &s->arena[off + 1], (size_t)copy * sizeof(*cps));
    return n;
}