#include "screencodec.h"

#include <stdlib.h>
#include <stdio.h>
#ifdef _WIN32
#include <WinSock.h>
#include <ipp.h>
#else
#include "SpBase.h"
#endif
#include <assert.h>

#include <zlib.h>


#define ZALLOC(size) zalloc(size)
#define ZCALLOC(num, size) czalloc(num, size)
#define MALLOC_T(type) (type*)malloc(sizeof(type))
#define ZALLOC_T(type) (type*)zalloc(sizeof(type))
#define CALLOC_T(num, type) (type*)calloc(num, sizeof(type))
#ifndef BS_WRITE_I2
#define BS_WRITE_I2(ptr, v) \
	do{\
		*(unsigned char*)ptr = v&0x00ff; \
		*((unsigned char*)ptr+1) = (v&0xff00) >> 8;\
	}while(0)\

#endif // !BS_WRITE_I2

static __inline int output(unsigned short *buf, int *size, unsigned int color, unsigned int count)
{
	int n = *size;
	assert(count);
	if (count == 1) {
		buf[n++] = color;
	} else if (count <= 0x7fff) {
		buf[n++] = color | 0x8000;
		buf[n++] = count | 0x8000;
	} else {
		buf[n++] = color | 0x8000;
		buf[n++] = count & 0x7fff;
		buf[n++] = 0x8000 | ((count & 0x3fff8000) >> 15);
	}
	*size = n;
	return 0;
}

static void rle_encode(int width, int height, const unsigned char *raw_buf, unsigned short *buf, int *size)
{
	const unsigned char *p = raw_buf;
	unsigned int last = -1, count = 0;
	int linesize = (width * 3 + 3) & 0xfffffffc;
	int len = 0;
	int i;
	for (i = 0; i < height; ++i) {
		const unsigned char *s = p;
		const unsigned char *e = p + width*3;
		while (p != e) {
			unsigned int b = *p++;
			unsigned int g = *p++;
			unsigned int r = *p++;
			unsigned int curr = ((r & 0xf8) << 7) | ((g&0xf8) << 2) | ((b & 0xf8) >> 3); // rgb555
			if (last != curr) {
				if (count) {
					output(buf, &len, last, count);
				}
				last = curr;
				count = 1;
			} else {
				count++;
			}
		}

		p = s + linesize;
	}
	output(buf, &len, last, count);
	*size = len;
}

SCREENCODEC_API(int) screencapture_encode(int width, int height, const void *raw_buf, void *buf, int *size)
{
	
	unsigned char *p;
	unsigned short *tmp_buf;
	int tmp_len;
	int rc;
	unsigned long src_len, dst_len;

	if (!buf) {
		if (size) {
			*size = width * height *3; // at most 
			return 0;
		}
		return -1;
	}

	if (!raw_buf || !buf)
		return -1;

	p = buf;

	BS_WRITE_I2(p, width);
	p += 2;
	BS_WRITE_I2(p, height);
	p += 2;

	tmp_len = width * height * 3;
	tmp_buf = (unsigned short*)malloc(tmp_len<<1);
	rle_encode(width, height, raw_buf, tmp_buf, &tmp_len);
	dst_len = *size - (p-(unsigned char*)buf);
	src_len = tmp_len << 1;
	rc = compress(p, &dst_len, (Bytef*)tmp_buf, src_len);
	free(tmp_buf);

	if (rc == 0) {
		p += dst_len;
		*size = p - (unsigned char *)buf;
	}

	return rc;
}

static int input(unsigned short **pp, unsigned short *end, unsigned int *color, unsigned int *count)
{
	unsigned short *p = *pp;
	if (p != end) {
		unsigned short t = *p++;
		if (t & 0x8000) {
			if (p != end) {
				unsigned short tt = *p++;
				if (tt & 0x8000) {
					*count = tt & 0x7fff;
					*color = t & 0x7fff;
				} else {
					if (p != end) {
						unsigned short ttt = *p++;
						if (ttt & 0x8000) {
							*color = t & 0x7fff;
							*count = (((unsigned int)ttt & 0x7fff) << 15) | ((unsigned int)tt & 0x7fff);
						} else {
							return -1;
						}
					} else {
						return -1;
					}
				}
			} else {
				return -1;
			}
		} else {
			*color = t;
			*count = 1;
		}
	} else {
		return -1;
	}
	*pp = p;
	return 0;
}

static int rle_decode(const void *src_buf, int size, int width, int height, unsigned char *buf)
{
	int rc = 0;
	int i = 0, j = 0;
	int linesize = (width * 3 + 3) & 0xfffffffc;
	int linegap = linesize - width * 3;
	unsigned char *c = (unsigned char*)buf;
	unsigned short *p = (unsigned short*)src_buf; 
	unsigned short *e = (unsigned short*)((const char*)src_buf + size);
	while (p != e) {
		unsigned int color, count;
		int r, g, b;
		rc = input(&p, e, &color, &count);
		if (rc != 0)
			break;
		b = (color & 0x1f) << 3;
		g = (color & 0x3e0) >> 2;
		r = (color & 0x7c00) >> 7;
		while (i < height && count > 0) {
			while (j < width && count > 0) {
				*c++ = b;
				*c++ = g;
				*c++ = r;
				j ++;
				count --;
			}
			if (j == width) {
				j = 0;
				i++;
				c += linegap;
			}
		}
	}
	return rc;
}


SCREENCODEC_API(int) screencapture_decode(int *width, int *height, const void *enc_buf, size_t enc_size, void *buf, int *size)
{
	const unsigned short *i2 = (const unsigned short *)enc_buf;
	int linesize;
	int linegap;
	unsigned char *tmp_buf;
	unsigned long tmp_len;
	int rc;

	*width = i2[0];
	*height = i2[1];
	linesize = (*width * 3 + 3) & 0xfffffffc;
	linegap = linesize - *width * 3;

	if (!buf) {
		*size = linesize * *height;
		return 0;
	}

	tmp_len = *width * *height * 3;
	tmp_buf = (unsigned char*)malloc(tmp_len);
	rc = uncompress(tmp_buf, &tmp_len, (Bytef*)&i2[2], enc_size-4);
	if (rc == 0) {
		rc = rle_decode(tmp_buf, tmp_len, *width, *height, buf);
	}
	free(tmp_buf);

	return rc;
}


#define SCREEN_CODEC_VER	1
#define SCREEN_CODEC_TAG	'RVC '

#define SCREEN_CODEC_TYPE_I	0
#define SCREEN_CODEC_TYPE_P	1

typedef struct screen_codec_hdr_t 
{
	unsigned int tag;
	unsigned int version : 3;
	unsigned int frame_type : 2;
	unsigned int compress : 1;
	unsigned int frame_id : 26;
	unsigned short width;
	unsigned short height;
	unsigned int size;
}screen_codec_hdr_t;

// 0: rle
// 1-254: spr
// 255: ref


struct screen_encoder_session_t
{
	BYTE *ref_frame;
	int width;
	int height;
	int seq_no;
};

static void img_rgb24_to_rgb8(int width, int height, const unsigned char *src, unsigned char *dst)
{
	const unsigned char *s = src;
	const unsigned char *e = s + width * height * 3;
	while (s != e) {
		*dst++ = ((s[0] & 0xC0) >> 6) | ((s[1] & 0xF0) >> 2) | (s[2] & 0xC0);
		s += 3;
	}
}

static void img_rgb8_to_rgb24(int width, int height, const unsigned char *src, unsigned char *dst)
{
	const unsigned char *s = src;
	const unsigned char *e = s + width * height;
	while (s != e) {
		unsigned int ch = *s;
		*dst++ = (ch & 3) << 6;
		*dst++ = (ch & 0x3c) << 2;
		*dst++ = (ch & 0xc0);
		s++;
	}
}

static int scan_rle_chunk(const unsigned char *s, const unsigned char *e)
{
	if (s == e) {
		return 0;
	} else {
		const unsigned char *ss = s;
		int ch = *ss;
		ss++;
		while (ss < e) {
			if (ch == *ss)
				ss++;
			else
				break;
		}
		return ss - s;
	}
}

static int scan_ref_chunk(const unsigned char *s, const unsigned char *e, const unsigned char *s_ref)
{
	if (s == e) {
		return 0;
	} else {
		const unsigned char *ss = s;
		const unsigned char *ss_ref = s_ref;
		while (ss < e) {
			if (*ss == *ss_ref) {
				ss++;
				ss_ref++;
			} else {
				break;
			}
		}
		return ss - s;
	}
}

typedef struct enc_output_ctx {
	int n;
	char spr[256];
}enc_output_ctx;

static int output_spr(enc_output_ctx *ctx, unsigned char *output)
{
	unsigned char *o = output;
	*o++ = (unsigned char)ctx->n;
	memcpy(o, ctx->spr, ctx->n);
	o += ctx->n;
	ctx->n = 0;
	return o - output;
}

static __inline int output_rle_chunk(enc_output_ctx *ctx, const unsigned char *chunk, int chunk_len, unsigned char *output)
{
	unsigned char *o = output;
	if (chunk_len == 1) {
		ctx->spr[ctx->n++] = *chunk;
		if (ctx->n == 254) {
			o += output_spr(ctx, o);
			return o - output;
		}
		return 0;
	} else {
		if (ctx->n > 0) {
			o += output_spr(ctx, o);
		}
		*o++ = 0; // rle start
		while (chunk_len >= 0x80) {
			*o++ = (unsigned char)(chunk_len | 0x80);
			chunk_len >>= 7;
		}
		*o++ = (unsigned char)chunk_len;
		*o++ = *chunk;
		return o - output;
	}
}

static __inline int output_ref_chunk(enc_output_ctx *ctx, int chunk_len, unsigned char *output)
{
	unsigned char *o = output;
	if (ctx->n > 0) {
		o += output_spr(ctx, o);
	}
	*o++ = 0xff; // ref start
	while (chunk_len >= 0x80) {
		*o++ = (unsigned char)(chunk_len | 0x80);
		chunk_len >>= 7;
	}
	*o++ = (unsigned char)chunk_len;
	return o - output;
}

static int encode(screen_encoder_session_t *session, const unsigned char* o, unsigned char *output)
{
	const unsigned char *s = o;
	const unsigned char *e = s + session->width * session->height;
	unsigned char *oo = output;
	enc_output_ctx ctx;
	int r0, r1;

	ctx.n = 0;

	for (;;) {
		if (session	->ref_frame) {
			r0 = scan_ref_chunk(s, e, session->ref_frame + (s - o));
		} else {
			r0 = 0;
		}
		r1 = scan_rle_chunk(s, e);
		if (r0 > r1) {
			if (r0 < 5) {
				oo += output_rle_chunk(&ctx, s, r1, oo);
				s += r1;
			} else {
				oo += output_ref_chunk(&ctx, r0, oo);
				s += r0;
			}
		} else if (r0 < r1) {
			oo += output_rle_chunk(&ctx, s, r1, oo);
			s += r1;
		} else {
			if (r1 == 0) {
				break; // finished
			}
			oo += output_rle_chunk(&ctx, s, r1, oo);
			s += r1;
		}
	}

	if (ctx.n > 0) {
		oo += output_spr(&ctx, oo);
	}

	return oo - output;
}

static int encode_frame(screen_encoder_session_t *session, const unsigned char* raw, unsigned char *output, int *size, int *type)
{
	unsigned char *rgb8_buf = (unsigned char*)malloc(session->width * session->height);
	img_rgb24_to_rgb8(session->width, session->height, raw, rgb8_buf);
	*size = encode(session, rgb8_buf, output);
	if (!session->ref_frame) {
		session->ref_frame = (BYTE*)malloc(session->width * session->height);
		memcpy(session->ref_frame, rgb8_buf, session->width * session->height);
		*type = SCREEN_CODEC_TYPE_I;
	} else {
		memcpy(session->ref_frame, rgb8_buf, session->width * session->height);
		*type = SCREEN_CODEC_TYPE_P;
	}
	free(rgb8_buf);
	return 0;
}

SCREENCODEC_API(int) screen_encoder_session_create(int width, int height, screen_encoder_session_t **p_session)
{
#ifdef _WIN32
	screen_encoder_session_t *session = MALLOC_T(screen_encoder_session_t);
#else
	screen_encoder_session_t *session = malloc(sizeof(screen_encoder_session_t));
#endif
	if (session) {
		session->width = width;
		session->height = height;
		session->ref_frame = NULL;
		session->seq_no = 0;
		*p_session = session;
		return 0;
	}
	return -1;
}

SCREENCODEC_API(void) screen_encoder_session_destroy(screen_encoder_session_t *session)
{
	if (session) {
		if (session->ref_frame) {
			free(session->ref_frame);
			session->ref_frame = NULL;
		}
		free(session);
	}
}

#pragma optimize("", off)
SCREENCODEC_API(int) screen_encoder_session_encode(screen_encoder_session_t *session, const void *raw, void *buf, int *size)
{
	int rc;

	if (!session)
		return -1;
	if (!raw)
		return -1;

	if (!buf) {
		if (size) {
			*size = session->width * session->height * 3 + sizeof(screen_codec_hdr_t); // at most
			return 0;
		}
		return -1;
	}
	if (!size)
		return -1;

	rc = 0;

	{
		int len = 0;
		int type = 0;
		void *tmp_buf = malloc(session->width * session->height * 3 * 2);
		rc = encode_frame(session, raw, (unsigned char*)tmp_buf, &len, &type);
		if (rc == 0) {
			unsigned long dst_len;
			rc = compress((char*)buf+sizeof(screen_codec_hdr_t), &dst_len, tmp_buf, len);
			if (rc == 0) {
				if (dst_len < len) { // compress
					screen_codec_hdr_t *hdr = (screen_codec_hdr_t *)buf;
					hdr->tag	 = SCREEN_CODEC_TAG;
					hdr->version = SCREEN_CODEC_VER;
					hdr->frame_type = type;
					hdr->compress = 1;
					hdr->frame_id = session->seq_no++;
					hdr->width = session->width;
					hdr->height = session->height;
					hdr->size = dst_len;
					*size = hdr->size + sizeof(screen_codec_hdr_t);
				} else {
					screen_codec_hdr_t *hdr = (screen_codec_hdr_t *)buf;
					hdr->tag	 = SCREEN_CODEC_TAG;
					hdr->version = SCREEN_CODEC_VER;
					hdr->frame_type = type;
					hdr->compress = 0;
					hdr->frame_id = session->seq_no++;
					hdr->width = session->width;
					hdr->height = session->height;
					hdr->size = len;
					*size = hdr->size + sizeof(screen_codec_hdr_t);
					memcpy((char*)buf+sizeof(screen_codec_hdr_t), tmp_buf, len);
				}
			}
		}
		free(tmp_buf);
	}

	return rc;
}
#pragma optimize("", on)

struct screen_decoder_session_t
{
	BYTE *ref_frame;
	int width;
	int height;
	int seq_no;
};

static __inline int expand_spr_chunk(int spr_len, const unsigned char *spr, unsigned char *s, unsigned char *e)
{
	if (e-s >= spr_len) {
		memcpy(s, spr, spr_len);
		return spr_len;
	} else {
		return -1;
	}
}

static __inline int expand_rle_chunk(int u, int u_len, unsigned char *s, unsigned char *e)
{
	if ((e - s) >= u_len) {
		memset(s, u, u_len);
		return u_len;
	} else {
		return -1;
	}
}

static __inline int expand_ref_chunk(int len, unsigned char *s, unsigned char *e, unsigned char *s_ref)
{
	if ((e - s) >= len) {
		memcpy(s, s_ref, len);
		return len;
	} else {
		return -1;
	}
}

static int read_7bit_int(const unsigned char *__s, const unsigned char *e, int *cnt)
{
	int count = 0;
	int shift = 0;
	const unsigned char *s = __s;
	unsigned char b;
	if (s == e)
		return -1;
	do {
		if (shift == 5*7)
			return -1; // currupt
		b = *s++;
		count |= (b & 0x7f) << shift;
		shift += 7;
	} while(b >= 0x80 && s < e);
	*cnt = count;
	return s - __s;
}

static int read_chunk(screen_decoder_session_t *session, const unsigned char *__ss, const unsigned char *ee, int *prefix, int *count, int *color)
{
	if (__ss == ee) {
		return 0;
	} else {
		const unsigned char *ss = __ss;
		*prefix = *ss++;
		if (ss == ee)
			return -1;
		if (*prefix == 0) { // rle
			int n;
			int rc = read_7bit_int(ss, ee, &n);
			if (rc < 0)
				return -1; // error
			ss += rc;
			*count = n;
			*color = *ss++;
			return ss - __ss;
		} else if (*prefix == 0xff) { // ref
			int n;
			int rc = read_7bit_int(ss, ee, &n);
			if (rc < 0)
				return -1; // error
			ss += rc;
			*count = n;
			return ss - __ss;
		} else {
			*count = *prefix;
			return ss - __ss;
		}
	}
}

static int decode(screen_decoder_session_t *session, const void *enc_buf, size_t enc_size, unsigned char *output)
{
	unsigned char *s = (unsigned char*)output;
	unsigned char *e = s + session->width * session->height;
	const unsigned char *ss = (unsigned char *)enc_buf;
	const unsigned char *ee = (unsigned char *)enc_buf	+ enc_size;
	int rc;
	int prefix;
	int color;
	int count;
	for (;;) {
		rc = read_chunk(session, ss, ee, &prefix, &count, &color);
		if (rc > 0) {
			if (prefix == 0) { // rle
				int t = expand_rle_chunk(color, count, s, e);
				if (t < 0)
					return -1;
				s += t;
				ss += rc;
			} else if (prefix == 0xff) { // ref
				int t;
				if (!session->ref_frame)
					return -1;
				t = expand_ref_chunk(count, s, e, session->ref_frame + (s - output));
				if (t < 0)
					return -1;
				s += t;
				ss += rc;
			} else {
				int t;
				ss += rc;
				t = expand_spr_chunk(count, ss, s, e);
				if (t < 0)
					return -1;
				ss += t;
				s += t;
			}
		} else if (rc == 0) {
			break; // finished
		} else {
			return -1; // error
		}
	}
	return ss == ee ? 0 : -1;
}

static int decode_frame(screen_decoder_session_t *session, const void *enc_buf, size_t enc_size, void *output)
{
	void *rgb8_buf = malloc(session->width * session->height);
	int rc = decode(session, enc_buf, enc_size, rgb8_buf);
	if (rc == 0) {
		img_rgb8_to_rgb24(session->width, session->height, (unsigned char*)rgb8_buf, (unsigned char*)output);
		if (!session->ref_frame)
			session->ref_frame = (BYTE*)malloc(session->width * session->height);
		memcpy(session->ref_frame, rgb8_buf, session->width * session->height);
	}
	free(rgb8_buf);
	return rc;
}

SCREENCODEC_API(int) screen_decoder_session_create(screen_decoder_session_t **p_session)
{
#ifdef _WIN32
	screen_decoder_session_t *session = MALLOC_T(screen_decoder_session_t);
#else
	screen_decoder_session_t *session = malloc(sizeof(screen_decoder_session_t));
#endif //_WIN32
	if (session) {
		session->width = 0;
		session->height = 0;
		session->ref_frame = NULL;
		session->seq_no = 0;
		*p_session = session;
		return 0;
	}
	return -1;
}

SCREENCODEC_API(void) screen_decoder_session_destroy(screen_decoder_session_t *session)
{
	if (session) {
		if (session->ref_frame) {
			free(session->ref_frame);
			session->ref_frame = NULL;
		}
		free(session);
	}
}

SCREENCODEC_API(int) screen_decoder_session_decode(screen_decoder_session_t *session, const void *enc_buf, size_t enc_size, int *width, int *height, void *buf, int *size)
{
	int rc;
	const screen_codec_hdr_t *hdr;

	if (!session)
		return -1;
	if (!enc_buf)
		return -1;
	if (enc_size <= sizeof(screen_codec_hdr_t))
		return -1;

	hdr = (const screen_codec_hdr_t*)enc_buf;

	if (hdr->tag != SCREEN_CODEC_TAG)
		return -1;
	if (hdr->version != SCREEN_CODEC_VER)
		return -1;
	if (session->ref_frame) {
		if (hdr->frame_id != session->seq_no+1)
			return -1;
		if (hdr->width != session->width)
			return -1;
		if (hdr->height != session->height)
			return -1;
	} else {
		session->width = hdr->width;
		session->height = hdr->height;
	}

	if (!buf) {
		if (width) {
			*width = hdr->width;
		}
		if (height) {
			*height = hdr->height;
		}
		if (size) {
			*size = hdr->width * hdr->height * 3;
			return 0;
		}
		return -1;
	}

	if (hdr->compress == 0) {
		rc = decode_frame(session, (const void *)(hdr+1), hdr->size, buf);
	} else {
		unsigned long dst_len = session->width * session->height * 3;
		void *tmp_buf = malloc(dst_len);
		rc = uncompress(tmp_buf, &dst_len, (const void*)(hdr+1), hdr->size);
		if (rc == 0) {
			rc = decode_frame(session, tmp_buf, dst_len, buf);
		}
		free(tmp_buf);
	}

	if (rc == 0) {
		session->seq_no = hdr->frame_id;
	}

	return rc;
}