#include <time.h>
#include <assert.h>
#include <stdio.h>
#include <stdbool.h>
#include <libswscale/swscale.h>
#include <libavutil/pixfmt.h>
#include <libavutil/pixdesc.h>
#include <libavutil/imgutils.h>

#define MINSIZE 8
#define MAXSIZE 128
#define SWS_ALIGN 16

typedef uint8_t *ImagePlane;

typedef struct
{
	ImagePlane plane[4];
	int stride[4];
	int pix_fmt;
	const const AVPixFmtDescriptor *pix_fmt_desc;
	int w, h;
}
Image;

Image *alloc_image(enum AVPixelFormat pix_fmt, int w, int h)
{
	Image *img = av_malloc(sizeof(Image));
	assert(img);

	img->w = w;
	img->h = h;
	img->pix_fmt = pix_fmt;
	img->pix_fmt_desc = &av_pix_fmt_descriptors[pix_fmt];
	assert(img->pix_fmt_desc);

	int sz = av_image_alloc(img->plane, img->stride, w, h, pix_fmt, SWS_ALIGN);
	assert(sz > 0);
	assert(img->plane[0]);
	memset(img->plane[0], 0, sz);

	return img;
}

void free_image(Image *img)
{
	av_freep(&img->plane[0]);
	av_free(img);
}

void fill_image(Image *img)
{
	int i, x, y;
	for(i = 0; i < img->pix_fmt_desc->nb_components; ++i)
	{
		int thisW = img->w;
		int thisH = img->h;
		uint16_t comp = rand() & ((1 << (img->pix_fmt_desc->comp[i].depth_minus1 + 1)) - 1);
		if(i > 0 && i < ((img->pix_fmt_desc->nb_components - 1) | 1) && !(img->pix_fmt_desc->flags & PIX_FMT_RGB))
		{
			thisW = -((-thisW) >> img->pix_fmt_desc->log2_chroma_w);
			thisH = -((-thisH) >> img->pix_fmt_desc->log2_chroma_h);
		}
		uint16_t buf[thisW];
		for(x = 0; x < thisW; ++x)
			buf[x] = comp;
		for(y = 0; y < thisH; ++y)
			av_write_image_line(buf, img->plane, img->stride, img->pix_fmt_desc, 0, y, i, thisW);
	}
}

bool is_solid_color(Image *img, int ignore_x, int ignore_y, int ignore_w, int ignore_h)
{
	int i, x, y;
	for(i = 0; i < img->pix_fmt_desc->nb_components; ++i)
	{
		int val = -1;
		int thisW = img->w;
		int thisH = img->h;
		int thisIgnoreX = ignore_x;
		int thisIgnoreY = ignore_y;
		int thisIgnoreW = ignore_w;
		int thisIgnoreH = ignore_h;
		if(i > 0 && i < ((img->pix_fmt_desc->nb_components - 1) | 1) && !(img->pix_fmt_desc->flags & PIX_FMT_RGB))
		{
			thisW = -((-thisW) >> img->pix_fmt_desc->log2_chroma_w);
			thisH = -((-thisH) >> img->pix_fmt_desc->log2_chroma_h);
			thisIgnoreX >>= img->pix_fmt_desc->log2_chroma_w;
			thisIgnoreY >>= img->pix_fmt_desc->log2_chroma_h;
			thisIgnoreW = -((-thisIgnoreW) >> img->pix_fmt_desc->log2_chroma_w);
			thisIgnoreH = -((-thisIgnoreH) >> img->pix_fmt_desc->log2_chroma_h);
		}
		uint16_t buf[thisW];
		for(y = 0; y < thisH; ++y)
		{
			av_read_image_line(buf, (const uint8_t **) img->plane, img->stride, img->pix_fmt_desc, 0, y, i, thisW, 0);
			for(x = 0; x < thisW; ++x)
			{
				if(y >= thisIgnoreY && y < thisIgnoreY + thisIgnoreH)
					if(x >= thisIgnoreX && x < thisIgnoreX + thisIgnoreW)
						continue;
				if(val == -1)
					val = buf[x];
				if(val != buf[x])
				{
					printf(" (failed in comp=%d x=%d y=%d expected=%d got=%d)", i, x, y, val, buf[x]);
					return 0;
				}
			}
		}
	}
	return 1;
}

bool set_to_cropped_image(Image *cropped, Image *orig, int x, int y, int w, int h)
{
	if(x < 0)
		return 0;
	if(y < 0)
		return 0;
	if(x + w > orig->w)
		return 0;
	if(y + h > orig->h)
		return 0;
	if(w < MINSIZE)
		return 0;
	if(h < MINSIZE)
		return 0;
	if(!(orig->pix_fmt_desc->flags & PIX_FMT_RGB))
		if(orig->pix_fmt_desc->nb_components >= 3)
		{
			// we have chroma
			// so only allow subrectangles that take whole chroma pixels
			if(x & ((1 << orig->pix_fmt_desc->log2_chroma_w) - 1))
				return 0;
			if(y & ((1 << orig->pix_fmt_desc->log2_chroma_h) - 1))
				return 0;
			if(w & ((1 << orig->pix_fmt_desc->log2_chroma_w) - 1))
				return 0;
			if(h & ((1 << orig->pix_fmt_desc->log2_chroma_h) - 1))
				return 0;
		}

	memset(cropped, 0, sizeof(*cropped));
	cropped->pix_fmt = orig->pix_fmt;
	cropped->pix_fmt_desc = orig->pix_fmt_desc;
	cropped->w = w;
	cropped->h = h;

	int p, c;
	for(p = 0; p < 4; ++p)
	{
		bool componentExists = 0;
		int thisOffset = 0;
		int thisY = 0;
		for(c = 0; c < orig->pix_fmt_desc->nb_components; ++c)
		{
			if(orig->pix_fmt_desc->comp[c].plane != p)
				continue;
			int step = orig->pix_fmt_desc->comp[c].step_minus1 + 1;
			int thisThisX = x;
			int thisThisY = y;
			int thisThisW = w;
			int thisThisH = h;
			if(c > 0 && c < ((orig->pix_fmt_desc->nb_components - 1) | 1) && !(orig->pix_fmt_desc->flags & PIX_FMT_RGB))
			{
				thisThisX >>= orig->pix_fmt_desc->log2_chroma_w;
				thisThisY >>= orig->pix_fmt_desc->log2_chroma_h;
				thisThisW = -((-thisThisW) >> orig->pix_fmt_desc->log2_chroma_w);
				thisThisH = -((-thisThisH) >> orig->pix_fmt_desc->log2_chroma_w);
			}
			int thisThisOffset = thisThisX * step;
			if(orig->pix_fmt_desc->flags & PIX_FMT_BITSTREAM)
			{
				if(thisThisOffset % 8 || thisThisW * step % 8)
				{
					memset(cropped, 0, sizeof(*cropped));
					return 0;
				}
				thisThisOffset /= 8;
			}
			if(componentExists)
			{
				assert(thisThisY == thisY);
				assert(thisThisOffset == thisOffset);
			}
			else
			{
				thisY = thisThisY;
				thisOffset = thisThisOffset;
				componentExists = 1;
			}
		}
		if(!componentExists)
			continue;
		cropped->plane[p] = orig->plane[p] + orig->stride[p] * thisY + thisOffset;
		cropped->stride[p] = orig->stride[p];
	}

	return 1;
}

int main(int argc, char **argv)
{
	if(argc != 3)
	{
		printf("Usage: %s srcFormat dstFormat\n", argv[0]);
		return 1;
	}

	enum AVPixelFormat srcFormat = av_get_pix_fmt(argv[1]);
	assert(srcFormat >= 0);

	enum AVPixelFormat dstFormat = av_get_pix_fmt(argv[2]);
	assert(dstFormat >= 0);

	srand(31337);
	struct SwsContext *ctx = NULL;

	for(;;)
	{
		int w1 = MINSIZE + rand() % (MAXSIZE - MINSIZE + 1);
		int h1 = MINSIZE + rand() % (MAXSIZE - MINSIZE + 1);
		// w1 = ((w1 - 1) | (SWS_ALIGN - 1)) + 1;
		// h1 = ((h1 - 1) | (SWS_ALIGN - 1)) + 1;

		int w2 = MINSIZE + rand() % (MAXSIZE - MINSIZE + 1);
		int h2 = MINSIZE + rand() % (MAXSIZE - MINSIZE + 1);
		// w2 = ((w2 - 1) | (SWS_ALIGN - 1)) + 1;
		// h2 = ((h2 - 1) | (SWS_ALIGN - 1)) + 1;

		// now choose a random subsection of image 1
		int w3 = rand() % w1 + 1;
		int h3 = rand() % h1 + 1;
		int x3 = rand() % (w1 - w3 + 1);
		int y3 = rand() % (h1 - h3 + 1);

		printf("TEST: %dx%d -> %dx%d+%d+%d @ %dx%d:",
			w2, h2,
			w3, h3, x3, y3,
			w1, h1);
		fflush(stdout);

		Image *i1 = alloc_image(dstFormat, w1, h1);
		Image *i2 = alloc_image(srcFormat, w2, h2);

		fill_image(i1);
		assert(is_solid_color(i1, 0, 0, 0, 0));
		fill_image(i2);
		assert(is_solid_color(i2, 0, 0, 0, 0));

		// create an Image aliasing this
		Image i3;
		if(set_to_cropped_image(&i3, i1, x3, y3, w3, h3))
		{
			assert(is_solid_color(&i3, 0, 0, 0, 0));
			int swsFlags =
				(1 << (rand() % 11)) | // 0x1 to 0x400 (scaler selection)
				// SWS_PRINT_INFO |
				(rand() & 1 ? SWS_DIRECT_BGR : 0) |
				(rand() & 1 ? SWS_ACCURATE_RND : 0) |
				(rand() & 1 ? SWS_BITEXACT : 0);
			ctx = sws_getCachedContext(ctx,
				i2->w, i2->h, i2->pix_fmt,
				i3.w, i3.h, i3.pix_fmt,
				swsFlags, NULL,
				NULL, NULL);
			if(ctx)
			{
				printf(" swsFlags=0x%x", swsFlags);
				int scaleResult = sws_scale(ctx, (const uint8_t * const*) i2->plane, i2->stride, 0, i2->h, i3.plane, i3.stride);
				assert(scaleResult == i3.h);
				int bad = 0;
				if(!is_solid_color(&i3, 0, 0, 0, 0))
				{
					printf(" output image is not solid!");
					bad = 1;
				}
				if(!is_solid_color(i1, x3, y3, w3, h3))
				{
					printf(" wrote beyond bounds!");
					bad = 1;
				}
				if(bad)
					printf(" FAILED\n");
				else
					printf(" OK\n");
			}
			else
				printf(" could not get sws context, skipped\n");
		}
		else
			printf(" invalid target dimensions, skipped\n");

		free_image(i2);
		free_image(i1);
	}
}