#include <stdbool.h>
#include <unistd.h>
#include <ctype.h>

#define STB_IMAGE_WRITE_IMPLEMENTATION
#include "stb_image_write.h"

enum ltok {
	END,
	XOR, OR, AND, NOT,
	ADD, DIV, SUB, MUL, MOD, EXP,
	EQ, LT, GT, COUNT,
	DUP, SWAP,
	LEFT, RIGHT,
	LNOT,
	HEIGHT, WIDTH,
	X, Y, LIT, ERROR,
};

struct token {
	enum ltok tag;
	unsigned int data;
};

/* i just realized i should have gone with a design where i lex into an array
 * so i don't have to lex the program once per pixel but i don't want to
 * rewrite it now */
struct token
ntok(char **source) {
	if (!**source) return (struct token){0};
	enum ltok tag = LIT;
	unsigned int data = 0;
	while (**source == ' ') (*source)++;
	switch (*(*source)++) {
	case '^': tag = XOR; break;
	case '|': tag = OR; break;
	case '&': tag = AND; break;
	case '~': tag = NOT; break;
	case '!': tag = LNOT; break;
	case '+': tag = ADD; break;
	case '/': tag = DIV; break;
	case '-': tag = SUB; break;
	case '*': tag = MUL; break;
	case '%': tag = MOD; break;
	case 'e': tag = EXP; break;
	case '=': tag = EQ; break;
	case '<': tag = LT; break;
	case '>': tag = GT; break;
	case '#': tag = COUNT; break;
	case 'x': tag = X; break;
	case 'y': tag = Y; break;
	case 'd': tag = DUP; break;
	case 's': tag = SWAP; break;
	case 'l': tag = LEFT; break;
	case 'r': tag = RIGHT; break;
	case 'h': tag = HEIGHT; break;
	case 'w': tag = WIDTH; break;
	default:
		 (*source)--;
		 if (!isdigit(**source)) {
			 (*source)++;
			 return (struct token){ERROR, 0};
		 } data = strtoul(*source, source, 10);
	};
	return (struct token){tag, data};
};

unsigned width = 255;
unsigned height = 255;

unsigned int *
eval(char *program, unsigned int x, unsigned int y, unsigned int *stack)
{
	char *errortext = "stack underflow";
	char *p = program;
	unsigned int *sp = stack;
#define push(n) (*(sp++) = n)
	struct token t = { 0 };
	for (;;) {
		t = ntok(&p);
		if (!t.tag) break;
		switch (t.tag) {
		case ERROR:
			errortext = "unrecognised token";
			goto error;
		case X: push(x); continue;
		case Y: push(y); continue;
		case HEIGHT: push(height); continue;
		case WIDTH: push(width); continue;
		case LIT: push(t.data); continue;
		};
		if (--sp < stack) goto error;
		unsigned int b = *sp;
		switch (t.tag) {
		case NOT: push(~b); continue;
		case LNOT: push(!b); continue;
		/* shhhhhhhh it's fine portability is overrated */
		case COUNT: push(__builtin_popcount(b)); continue;
		case DUP: push(b); push(b); continue;
		};
		if (--sp < stack) goto error;
		unsigned int a = *sp;

		if (!b && (t.tag == MOD || t.tag == DIV)) {
			errortext = "division by zero";
			goto error;
		}
		switch (t.tag) {
		case XOR: push(a ^ b); continue;
		case OR: push(a | b); continue;
		case AND: push(a & b); continue;
		case ADD: push(a + b); continue;
		case DIV: push(a / b); continue;
		case SUB: push(a - b); continue;
		case MUL: push(a * b); continue;
		case MOD: push(a % b); continue;
		case EXP:;
			  unsigned int res = 1;
			  for (unsigned int i = 0; i < b; i++) res *= a;
			  push(res);
			  continue;
		case EQ: push(a == b); continue;
		case LT: push(a < b); continue;
		case GT: push(a > b); continue;
		case SWAP: push(b); push(a); continue;
		case LEFT: push(a << b); continue;
		case RIGHT: push(b >> b); continue;
		}
	}
	return sp;
error:
	fputs(program, stderr);
	fputc('\n', stderr);
	for (char *a = p - 1; a > program; a--) fputc(' ', stderr);
	fprintf(stderr, "\033[31;1m^--- %s\033[0m\n", errortext);
	exit(1);
}

void
usage(char *name)
{
	fprintf(stderr, "usage: %s [-g] [-c] [-i] [-w width] [-h height] [-s size] [-o out.png] program\n", name);
	exit(1);
}

void
write_image_stdout(void *context, void *data, int size)
{
	write(STDOUT_FILENO, data, size);
}

int
main(int ac, char **av)
{
	int opt;
	bool gradient = false;
	bool invert = false;
	bool color = false;
	char *filename = NULL;
	while ((opt = getopt(ac, av, "o:w:h:s:gci")) != -1) {
		switch (opt) {
		case 'w':
			width = atoi(optarg);
			break;
		case 'h':
			height = atoi(optarg);
			break;
		case 's':
			height = atoi(optarg);
			width = height;
			break;
		case 'o':
			filename = optarg;
			break;
		case 'g':
			gradient = true;
			break;
		case 'c':
			color = true;
			break;
		case 'i':
			invert = true;
			break;
		default:
			usage(av[0]);
		}
	}
	int channels = color ? 3 : 1;
	if (optind >= ac) usage(av[0]);
	char *program = av[optind];
	char *data = malloc(width * height * channels);
	/* stack cannot ever be larger than the program length, since there are
	* no instructions which grow the stack by more than one */
	unsigned int *stack = calloc(strlen(program), sizeof(*stack));
	for (unsigned y = 1; y < height; y++)
	for (unsigned x = 1; x < width; x++) {
		unsigned int *sp = eval(program, x, y, stack);
		if (sp - stack < channels) {
			fprintf(stderr, "insufficient values left on the stack\n");
			exit(1);
		}
		for (int i = 0; i < channels; i++) {
			data[(x + y * width) * channels + i] =
				gradient ? stack[i] : ((!!stack[i]) != invert) ? -1 : 0;
		}
	}
	free(stack);
	if (filename == NULL)
		stbi_write_png_to_func(&write_image_stdout, NULL, width, height, channels, data, 0);
	else stbi_write_png(filename, width, height, channels, data, 0);
	free(data);
	return 0;
}