#include "wrap.h"
#include <stdlib.h>

typedef struct {
    int runIndex;
    int start, end;
    float width;
    int isSpace;
} Atom;

typedef struct {
    int atomStart, atomEnd; /* [atomStart, atomEnd) into the atoms array */
} Unit;

void wrap_result_init(WrapResult *wr) {
    wr->lines = NULL;
    wr->count = 0;
    wr->cap = 0;
}

void wrap_result_free(WrapResult *wr) {
    for (int i = 0; i < wr->count; i++) free(wr->lines[i].segs);
    free(wr->lines);
    wr->lines = NULL;
    wr->count = 0;
    wr->cap = 0;
}

static WrapLine *push_line(WrapResult *wr) {
    if (wr->count >= wr->cap) {
        wr->cap = wr->cap > 0 ? wr->cap * 2 : 4;
        wr->lines = realloc(wr->lines, (size_t)wr->cap * sizeof(WrapLine));
    }
    WrapLine *l = &wr->lines[wr->count++];
    l->segs = NULL;
    l->count = 0;
    l->cap = 0;
    return l;
}

static void push_seg(WrapLine *l, int runIndex, int start, int end, float width) {
    if (l->count >= l->cap) {
        l->cap = l->cap > 0 ? l->cap * 2 : 4;
        l->segs = realloc(l->segs, (size_t)l->cap * sizeof(WrapSegment));
    }
    WrapSegment *s = &l->segs[l->count++];
    s->runIndex = runIndex;
    s->start = start;
    s->end = end;
    s->width = width;
}

static void finalize_line(WrapResult *out, const Atom *atoms, int atomStart, int atomEnd) {
    WrapLine *line = push_line(out);
    if (atomEnd <= atomStart) return; /* empty line, no segments */
    int segRun = atoms[atomStart].runIndex;
    int segStart = atoms[atomStart].start;
    int segEnd = atoms[atomStart].end;
    float segWidth = atoms[atomStart].width;
    for (int i = atomStart + 1; i < atomEnd; i++) {
        if (atoms[i].runIndex == segRun) {
            segEnd = atoms[i].end;
            segWidth += atoms[i].width;
        } else {
            push_seg(line, segRun, segStart, segEnd, segWidth);
            segRun = atoms[i].runIndex;
            segStart = atoms[i].start;
            segEnd = atoms[i].end;
            segWidth = atoms[i].width;
        }
    }
    push_seg(line, segRun, segStart, segEnd, segWidth);
}

void wrap_layout(const char *text, const WrapRunSpan *runs, int runCount, float maxWidth,
                  WrapMeasureFn measure, void *measureCtx, WrapResult *out) {
    wrap_result_init(out);

    /* Step A: split each run's text into word/space atoms. */
    Atom *atoms = NULL;
    int atomCount = 0, atomCap = 0;
    for (int r = 0; r < runCount; r++) {
        int i = runs[r].start;
        int end = runs[r].end;
        while (i < end) {
            int isSpace = (text[i] == ' ');
            int j = i + 1;
            while (j < end && (text[j] == ' ') == isSpace) j++;
            if (atomCount >= atomCap) {
                atomCap = atomCap > 0 ? atomCap * 2 : 16;
                atoms = realloc(atoms, (size_t)atomCap * sizeof(Atom));
            }
            Atom *a = &atoms[atomCount++];
            a->runIndex = r;
            a->start = i;
            a->end = j;
            a->isSpace = isSpace;
            a->width = measure(measureCtx, r, text, i, j - i);
            i = j;
        }
    }

    if (atomCount == 0) {
        finalize_line(out, atoms, 0, 0);
        free(atoms);
        return;
    }

    /* Step B: group atoms into unbreakable units (word plus its trailing space, if any). */
    Unit *units = NULL;
    int unitCount = 0, unitCap = 0;
    int a = 0;
    while (a < atomCount) {
        int unitStart = a;
        while (a < atomCount) {
            int wasSpace = atoms[a].isSpace;
            a++;
            if (wasSpace) break;
        }
        if (unitCount >= unitCap) {
            unitCap = unitCap > 0 ? unitCap * 2 : 16;
            units = realloc(units, (size_t)unitCap * sizeof(Unit));
        }
        units[unitCount].atomStart = unitStart;
        units[unitCount].atomEnd = a;
        unitCount++;
    }

    /* Step C: greedily pack units onto lines. */
    int lineStartUnit = 0;
    float lineWidth = 0.0f;
    for (int u = 0; u < unitCount; u++) {
        float unitWidth = 0.0f;
        for (int i = units[u].atomStart; i < units[u].atomEnd; i++) unitWidth += atoms[i].width;

        if (u > lineStartUnit && lineWidth + unitWidth > maxWidth) {
            finalize_line(out, atoms, units[lineStartUnit].atomStart, units[u - 1].atomEnd);
            lineStartUnit = u;
            lineWidth = 0.0f;
        }
        lineWidth += unitWidth;
    }
    finalize_line(out, atoms, units[lineStartUnit].atomStart, units[unitCount - 1].atomEnd);

    free(units);
    free(atoms);
}