#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "test_image.h"

/* ------------------------------------------------------------------ */
/* BT.601 limited-range YCbCr → RGB  (for PPM output only)             */
/* ------------------------------------------------------------------ */

static void ycbcr_to_rgb(uint8_t y, uint8_t cb, uint8_t cr,
                          uint8_t *r, uint8_t *g, uint8_t *b)
{
	int yy  = (int)y  - 16;
	int cb_ = (int)cb - 128;
	int cr_ = (int)cr - 128;
	int rv  = (298 * yy + 409 * cr_          + 128) >> 8;
	int gv  = (298 * yy - 100 * cb_ - 208 * cr_ + 128) >> 8;
	int bv  = (298 * yy + 516 * cb_          + 128) >> 8;
	*r = (uint8_t)(rv < 0 ? 0 : rv > 255 ? 255 : rv);
	*g = (uint8_t)(gv < 0 ? 0 : gv > 255 ? 255 : gv);
	*b = (uint8_t)(bv < 0 ? 0 : bv > 255 ? 255 : bv);
}

/* ------------------------------------------------------------------ */
/* PPM output                                                           */
/* ------------------------------------------------------------------ */

static int write_ppm(const char *path, Test_Frame *f)
{
	FILE *fp = fopen(path, "wb");
	if (!fp) { perror(path); return -1; }

	fprintf(fp, "P6\n%d %d\n255\n", f->width, f->height);

	for (int row = 0; row < f->height; row++) {
		for (int col = 0; col < f->width; col++) {
			uint8_t r, g, b;

			if (f->fmt == TEST_FMT_BGRA) {
				const uint8_t *p = f->plane[0] + row * f->stride[0] + col * 4;
				b = p[0]; g = p[1]; r = p[2];
			} else {
				uint8_t luma = f->plane[0][row * f->stride[0] + col];
				uint8_t cb, cr;
				if (f->fmt == TEST_FMT_YUV420) {
					cb = f->plane[1][(row / 2) * f->stride[1] + col / 2];
					cr = f->plane[2][(row / 2) * f->stride[2] + col / 2];
				} else { /* YUV422 */
					cb = f->plane[1][row * f->stride[1] + col / 2];
					cr = f->plane[2][row * f->stride[2] + col / 2];
				}
				ycbcr_to_rgb(luma, cb, cr, &r, &g, &b);
			}

			uint8_t pix[3] = {r, g, b};
			fwrite(pix, 1, 3, fp);
		}
	}

	fclose(fp);
	return 0;
}

/* ------------------------------------------------------------------ */
/* Main                                                                 */
/* ------------------------------------------------------------------ */

static void usage(void)
{
	fprintf(stderr,
	    "usage: test_image_cli [--pattern bars|ramp|grid]\n"
	    "                      [--width N] [--height N]\n"
	    "                      [--format yuv420|yuv422|bgra]\n"
	    "                      --out FILE.ppm\n"
	    "\n"
	    "defaults: bars 1280x720 yuv420\n");
}

int main(int argc, char **argv)
{
	Test_Pattern pattern = TEST_PATTERN_BARS;
	Test_Fmt     fmt     = TEST_FMT_YUV420;
	int          width   = 1280;
	int          height  = 720;
	const char  *out     = NULL;

	for (int i = 1; i < argc; i++) {
		if (strcmp(argv[i], "--pattern") == 0 && i + 1 < argc) {
			i++;
			if      (strcmp(argv[i], "bars") == 0) { pattern = TEST_PATTERN_BARS; }
			else if (strcmp(argv[i], "ramp") == 0) { pattern = TEST_PATTERN_RAMP; }
			else if (strcmp(argv[i], "grid") == 0) { pattern = TEST_PATTERN_GRID; }
			else {
				fprintf(stderr, "unknown pattern: %s\n", argv[i]);
				usage(); return 1;
			}
		} else if (strcmp(argv[i], "--width") == 0 && i + 1 < argc) {
			width = atoi(argv[++i]);
		} else if (strcmp(argv[i], "--height") == 0 && i + 1 < argc) {
			height = atoi(argv[++i]);
		} else if (strcmp(argv[i], "--format") == 0 && i + 1 < argc) {
			i++;
			if      (strcmp(argv[i], "yuv420") == 0) { fmt = TEST_FMT_YUV420; }
			else if (strcmp(argv[i], "yuv422") == 0) { fmt = TEST_FMT_YUV422; }
			else if (strcmp(argv[i], "bgra")   == 0) { fmt = TEST_FMT_BGRA;   }
			else {
				fprintf(stderr, "unknown format: %s\n", argv[i]);
				usage(); return 1;
			}
		} else if (strcmp(argv[i], "--out") == 0 && i + 1 < argc) {
			out = argv[++i];
		} else {
			usage(); return 1;
		}
	}

	if (!out) {
		fprintf(stderr, "test_image_cli: --out FILE required\n");
		usage(); return 1;
	}
	if (width < 2 || height < 2) {
		fprintf(stderr, "test_image_cli: width and height must be >= 2\n");
		return 1;
	}

	Test_Frame *f = test_image_alloc(width, height, fmt);
	if (!f) {
		fprintf(stderr, "test_image_cli: allocation failed\n");
		return 1;
	}

	test_image_generate(f, pattern);

	int rc = write_ppm(out, f);
	test_image_free(f);
	if (rc != 0) { return 1; }

	const char *fmt_name = fmt == TEST_FMT_YUV420 ? "yuv420"
	                     : fmt == TEST_FMT_YUV422 ? "yuv422"
	                     :                          "bgra";
	const char *pat_name = pattern == TEST_PATTERN_BARS ? "bars"
	                     : pattern == TEST_PATTERN_RAMP ? "ramp"
	                     :                               "grid";

	printf("%s: %dx%d  format=%s  pattern=%s\n", out, width, height, fmt_name, pat_name);
	return 0;
}