#include <u.h>
#include <libc.h>
#include <qoi.h>

enum
{
	Bufsz = 8192,
};

int
qoiread(QOIctx *ctx, u8int *in, int sz)
{
	u8int r, g, b, a;
	int n, m, isz;
	QOIpix p;

	if(sz < 1){
		if(ctx->o == nil || ctx->o != ctx->e){
			werrstr("eof before end");
			goto err;
		}
	}

	if(ctx->o == nil){
		ctx->o = ctx->pix;
		ctx->e = ctx->o + ctx->w * ctx->h * ctx->chan;
		ctx->prev = (QOIpix){0, 0, 0, 0xff};
	}else if(ctx->o == ctx->e)
		return 0;

	isz = sz;
	p = ctx->prev;
	while(sz > 0 && ctx->o != ctx->e){
		a = p.a;
		if(*in == 0xfe){ /* rgb */
			if(sz < 4)
				break;
			sz -= 4;
			b = *ctx->o++ = in[3];
			g = *ctx->o++ = in[2];
			r = *ctx->o++ = in[1];
			if(ctx->chan > 3)
				*ctx->o++ = a;
			in += 4;
		}else if(*in == 0xff){ /* rgba */
			if(sz < 5)
				break;
			sz -= 5;
			b = *ctx->o++ = in[3];
			g = *ctx->o++ = in[2];
			r = *ctx->o++ = in[1];
			a = in[4];
			if(ctx->chan > 3)
				*ctx->o++ = a;
			in += 5;
		}else if((*in & 0xc0) == 0x80){ /* luma */
			if(sz < 2)
				break;
			sz -= 2;
			g = -32 + (in[0] & 0x3f);
			r = g - 8 + (in[1] >> 4);
			b = g - 8 + (in[1] & 0x0f);
			b = *ctx->o++ = p.b + b;
			g = *ctx->o++ = p.g + g;
			r = *ctx->o++ = p.r + r;
			if(ctx->chan > 3)
				*ctx->o++ = a;
			in += 2;
		}else if((*in & 0xc0) == 0x40){ /* diff */
			sz--;
			b = *ctx->o++ = p.b - 2 + ((*in>>0) & 3);
			g = *ctx->o++ = p.g - 2 + ((*in>>2) & 3);
			r = *ctx->o++ = p.r - 2 + ((*in>>4) & 3);
			if(ctx->chan > 3)
				*ctx->o++ = a;
			in++;
		}else if((*in & 0xc0) == 0x00){ /* index */
			sz--;
			b = *ctx->o++ = ctx->cache[*in].b;
			g = *ctx->o++ = ctx->cache[*in].g;
			r = *ctx->o++ = ctx->cache[*in].r;
			a = ctx->cache[*in].a;
			if(ctx->chan > 3)
				*ctx->o++ = a;
			in++;
		}else{ /* run */
			sz--;
			if((n = (*in&0x3f) + 1) > (m = (ctx->e - ctx->o)/ctx->chan)){
				werrstr("run out of bounds: %d > %d", n, m);
				goto err;
			}
			if(ctx->chan > 3){
				b = p.b*a/0xff;
				g = p.g*a/0xff;
				r = p.r*a/0xff;
				do{
					*ctx->o++ = b;
					*ctx->o++ = g;
					*ctx->o++ = r;
					*ctx->o++ = a;
				}while(--n > 0);
			}else{
				do{
					*ctx->o++ = p.b;
					*ctx->o++ = p.g;
					*ctx->o++ = p.r;
				}while(--n > 0);
			}
			in++;
			continue;
		}
		p = (QOIpix){b, g, r, a};
		ctx->cache[(r*3 + g*5 + b*7 + a*11) & 63] = p;
		if(ctx->chan > 3){
			ctx->o[-2] = ctx->o[-2]*a/0xff;
			ctx->o[-3] = ctx->o[-3]*a/0xff;
			ctx->o[-4] = ctx->o[-4]*a/0xff;
		}
	}
	ctx->prev = p;

	if(ctx->o == ctx->e){
		n = sz;
		if(sz > 8-ctx->end)
			sz -= 8-ctx->end;
		else
			sz = 0;
		ctx->end += n - sz;
	}

	return isz - sz;
err:
	werrstr("qoiread: %r");
	return -1;
}

u8int *
qoireadfile(QOIctx *ctx, int fd)
{
	u8int b[14], *p;
	int n, r, sz;

	ctx->pix = nil;
	p = nil;
	if(readn(fd, b, sizeof(b)) != sizeof(b))
		goto err;

	if(memcmp(b, "qoif", 4) != 0){
		werrstr("invalid magic");
		goto err;
	}
	ctx->w = b[4]<<24 | b[5]<<16 | b[ 6]<<8 | b[ 7];
	ctx->h = b[8]<<24 | b[9]<<16 | b[10]<<8 | b[11];
	ctx->chan = b[12];
	ctx->space = b[13];
	memset(ctx->cache, 0, sizeof(ctx->cache));
	if(ctx->w > 16384 || ctx->h > 16384 || ctx->chan < 3 || ctx->chan > 4 || ctx->space > 1){
		werrstr("invalid header");
		goto err;
	}
	if((ctx->pix = malloc(ctx->w * ctx->h * ctx->chan)) == nil || (p = malloc(Bufsz)) == nil){
		werrstr("no memory");
		goto err;
	}
	ctx->o = ctx->e = nil;
	for(sz = 0;;){
		if((n = read(fd, p+sz, Bufsz-sz)) < 0)
			goto err;
		sz += n;
		if((r = qoiread(ctx, p, sz)) < 0)
			goto err;
		memmove(p, p+r, sz-r);
		sz -= r;
		if(n == 0){
			if(r != 0){
				werrstr("short read");
				goto err;
			}
			break;
		}
	}
	free(p);

	return ctx->pix;

err:
	werrstr("qoireadfile: %r");
	free(ctx->pix);
	free(p);
	return nil;
}