#include "jpegframe.h"
#include <string.h>

enum { SEEK_FF, SEEK_D8, MARK_FF, MARK, LEN_HI, LEN_LO, SEG, ENTROPY, ENTROPY_FF };

void jpegframe_init(jpegframe_t *f, uint8_t *out, size_t out_cap)
{
    memset(f, 0, sizeof *f);
    f->out = out;
    f->out_cap = out_cap;
}

void jpegframe_reset(jpegframe_t *f)
{
    f->state = SEEK_FF;
    f->out_len = 0;
    f->skipped = 0;
}

static inline bool put(jpegframe_t *f, uint8_t b)
{
    if (f->out_len >= f->out_cap) return false;
    f->out[f->out_len++] = b;
    return true;
}

size_t jpegframe_feed(jpegframe_t *f, const uint8_t *in, size_t n, bool *complete, bool *overflow)
{
    *complete = false;
    *overflow = false;
    size_t i = 0;
    while (i < n) {
        uint8_t b = in[i++];
        switch (f->state) {
        case SEEK_FF:
            if (b == 0xFF) f->state = SEEK_D8; else f->skipped++;
            break;
        case SEEK_D8:
            if (b == 0xD8) { f->out_len = 0; put(f, 0xFF); put(f, 0xD8); f->state = MARK_FF; }
            else if (b != 0xFF) { f->state = SEEK_FF; f->skipped += 2; }
            break;
        case MARK_FF:
            if (!put(f, b)) goto overflow;
            if (b == 0xFF) f->state = MARK;
            else { f->state = SEEK_FF; }               // corrupt: resync
            break;
        case MARK:
            if (!put(f, b)) goto overflow;
            if (b == 0xFF) break;                       // fill byte
            f->marker = b;
            if (b == 0xD9) { *complete = true; f->state = SEEK_FF; return i; }
            if (b == 0xD8 || b == 0x01 || (b >= 0xD0 && b <= 0xD7)) f->state = MARK_FF;
            else f->state = LEN_HI;
            break;
        case LEN_HI:
            if (!put(f, b)) goto overflow;
            f->remaining = (uint32_t)b << 8;
            f->state = LEN_LO;
            break;
        case LEN_LO:
            if (!put(f, b)) goto overflow;
            f->remaining |= b;
            if (f->remaining < 2) { f->state = SEEK_FF; break; }
            f->remaining -= 2;
            f->state = f->remaining ? SEG : (f->marker == 0xDA ? ENTROPY : MARK_FF);
            break;
        case SEG: {
            // bulk copy the segment payload
            size_t take = n - (i - 1);
            if (take > f->remaining) take = f->remaining;
            if (f->out_len + take > f->out_cap) goto overflow;
            memcpy(f->out + f->out_len, in + i - 1, take);
            f->out_len += take;
            f->remaining -= take;
            i += take - 1;
            if (!f->remaining) f->state = f->marker == 0xDA ? ENTROPY : MARK_FF;
            break;
        }
        case ENTROPY:
            if (!put(f, b)) goto overflow;
            if (b == 0xFF) f->state = ENTROPY_FF;
            break;
        case ENTROPY_FF:
            if (!put(f, b)) goto overflow;
            if (b == 0x00 || (b >= 0xD0 && b <= 0xD7)) f->state = ENTROPY;   // stuffed byte / restart
            else if (b == 0xFF) break;
            else if (b == 0xD9) { *complete = true; f->state = SEEK_FF; return i; }
            else { f->marker = b; f->state = LEN_HI; }                         // DHT/DQT/SOS of a later scan
            break;
        }
    }
    return i;
overflow:
    *overflow = true;
    f->state = SEEK_FF;
    f->out_len = 0;
    return i;
}