#include <M5Unified.h>

/*
Copyright (c) 2026 Devine Lu Linvega

Permission to use, copy, modify, and distribute this software for any
purpose with or without fee is hereby granted, provided that the above
copyright notice and this permission notice appear in all copies.

THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
WITH REGARD TO THIS SOFTWARE.

The m5 Stick C Plus, requires m5stack(2.1.4)
*/

typedef int8_t Sint8;
typedef uint8_t Uint8;
typedef uint16_t Uint16;

static const uint8_t ROM_DATA[] = {
	0xa0, 0x4c, 0xfd, 0x80, 0x08, 0x37, 0xa0, 0x4c, 0xf3, 0x80, 0x0a, 0x37, 0xa0, 0xdc, 0xf2, 0x80,
	0x0c, 0x37, 0xa0, 0x01, 0x38, 0x80, 0x20, 0x37, 0x80, 0x22, 0x36, 0x26, 0xa0, 0x00, 0x20, 0x39,
	0x80, 0x22, 0x33, 0x80, 0x01, 0x3f, 0x80, 0x24, 0x36, 0x26, 0xa0, 0x00, 0x10, 0x39, 0x80, 0x2f,
	0x33, 0x80, 0x01, 0x3f, 0x60, 0x00, 0x40, 0x00, 0x80, 0x4f, 0x32, 0x9d, 0x20, 0x00, 0x04, 0xa0,
	0x38, 0x0f, 0x13, 0x26, 0xa0, 0x00, 0x00, 0x29, 0x20, 0x00, 0x04, 0xa0, 0x39, 0x03, 0x13, 0xa0,
	0x00, 0x01, 0x38, 0x80, 0x3a, 0x32, 0x9d, 0x20, 0x00, 0x04, 0xa0, 0x38, 0x0f, 0x13, 0x26, 0xa0,
	0x00, 0x00, 0x29, 0x20, 0x00, 0x04, 0xa0, 0x39, 0x03, 0x13, 0xa0, 0x00, 0x01, 0x38, 0x60, 0x00,
	0x01, 0x00, 0x80, 0x00, 0x60, 0x00, 0x08, 0x80, 0x16, 0x33, 0x80, 0x0d, 0x33, 0x80, 0x01, 0xa0,
	0x36, 0x26, 0x17, 0xa0, 0x01, 0x9a, 0x80, 0x2c, 0x37, 0xa0, 0x00, 0x00, 0x80, 0x28, 0x37, 0xa0,
	0x00, 0x00, 0x80, 0x2a, 0x37, 0x80, 0x2f, 0x97, 0x17, 0x6c, 0x00, 0x1f, 0x3f, 0x38, 0x38, 0x38,
	0x78, 0x7f, 0x00, 0xfe, 0xfe, 0x7e, 0x77, 0x77, 0xe3, 0xc3, 0x00, 0x0f, 0x1f, 0x3b, 0x7b, 0x77,
	0xe7, 0xc7, 0x00, 0xfc, 0xfe, 0x8f, 0x87, 0x07, 0x0e, 0xfc, 0x7f, 0x00, 0x00, 0x0f, 0xff, 0x7f,
	0x07, 0x00, 0x03, 0x01, 0x00, 0xff, 0xf0, 0xf8, 0xff, 0x00, 0x87, 0x00, 0x00, 0xff, 0x7f, 0x7f,
	0xff, 0x00, 0xf0, 0x00, 0x00, 0xe0, 0xfc, 0xfc, 0x80, 0x00 };

#define SCREEN_W 240
#define SCREEN_H 135

static Uint8 ram[0x10000];
static Uint8 dev[0x100];
static Uint8 stk[2][0x100];
static Uint8 ptr[2];
static uint16_t screen_palette[16];
static int scr_x1, scr_y1, scr_x2, scr_y2, scr_dirty;
static inline Uint16 peek2(const Uint8 *d) { return ((Uint16)d[0] << 8) | d[1]; }
static inline void poke2(Uint8 *d, Uint16 v) { d[0] = v >> 8; d[1] = v; }
#define TWOS(v) ((v) & 0x8000 ? (int)(v) - 0x10000 : (int)(v))

static uint16_t nibble_to_rgb565(Uint8 r4, Uint8 g4, Uint8 b4)
{
	uint8_t r8 = (r4 << 4) | r4;
	uint8_t g8 = (g4 << 4) | g4;
	uint8_t b8 = (b4 << 4) | b4;
	return ((r8 & 0xF8) << 8) | ((g8 & 0xFC) << 3) | (b8 >> 3);
}

static unsigned int uxn_eval(Uint16 pc);

/*
@|System ------------------------------------------------------------ */

static int screen_vector = 0;

static void system_deo_colorize(void)
{
	unsigned int i, shift;
	uint16_t colors[4];
	for(i = 0, shift = 4; i < 4; ++i, shift ^= 4) {
		Uint8 r = (dev[0x8 + i/2] >> shift) & 0xf;
		Uint8 g = (dev[0xa + i/2] >> shift) & 0xf;
		Uint8 b = (dev[0xc + i/2] >> shift) & 0xf;
		colors[i] = nibble_to_rgb565(r, g, b);
	}
	for(i = 0; i < 16; i++)
		screen_palette[i] = colors[i >> 2 ? i >> 2 : i & 3];
	scr_dirty = 1;
	scr_x1 = scr_y1 = 0;
	scr_x2 = SCREEN_W;
	scr_y2 = SCREEN_H;
}

/*
@|Screen ------------------------------------------------------------ */

static Uint8 screen_layers[SCREEN_W * SCREEN_H];
static int rX, rY, rA, rMX, rMY, rMA, rML, rDX, rDY;
static const Uint8 alpha_lut[16] = {0,1,2,3,4,0,1,2,3,4,0,1,2,3,4,0};
static const Uint8 blend_lut[16][2][4] = {
	{{0,0,1,2},{0,0,4,8}},  {{0,1,2,3},{0,4,8,12}},
	{{0,2,3,1},{0,8,12,4}}, {{0,3,1,2},{0,12,4,8}},
	{{1,0,1,2},{4,0,4,8}},  {{1,1,2,3},{4,4,8,12}},
	{{1,2,3,1},{4,8,12,4}}, {{1,3,1,2},{4,12,4,8}},
	{{2,0,1,2},{8,0,4,8}},  {{2,1,2,3},{8,4,8,12}},
	{{2,2,3,1},{8,8,12,4}}, {{2,3,1,2},{8,12,4,8}},
	{{3,0,1,2},{12,0,4,8}}, {{3,1,2,3},{12,4,8,12}},
	{{3,2,3,1},{12,8,12,4}}, {{3,3,1,2},{12,12,4,8}}
};

static void screen_change(int x1, int y1, int x2, int y2)
{
	if(!scr_dirty) { scr_x1=x1; scr_y1=y1; scr_x2=x2; scr_y2=y2; scr_dirty=1; return; }
	if(x1 < scr_x1) scr_x1 = x1;
	if(y1 < scr_y1) scr_y1 = y1;
	if(x2 > scr_x2) scr_x2 = x2;
	if(y2 > scr_y2) scr_y2 = y2;
}

static void screen_deo_pixel(void)
{
	const int ctrl = dev[0x2e];
	const int hi = ctrl & 0x40;
	const Uint8 mask = hi ? 0x03 : 0x0c; 
	const Uint8 color = hi ? ((ctrl & 0x3) << 2) : (ctrl & 0x3);
	if(ctrl & 0x80) { 
		int px, py;
		const int x1 = (ctrl & 0x10) ? 0 : rX;
		const int x2 = (ctrl & 0x10) ? rX : SCREEN_W;
		const int y1 = (ctrl & 0x20) ? 0 : rY;
		const int y2 = (ctrl & 0x20) ? rY : SCREEN_H;
		for(py = y1; py < y2; py++)
			for(px = x1; px < x2; px++) {
				Uint8 *d = &screen_layers[py * SCREEN_W + px];
				*d = (*d & mask) | color;
			}
		screen_change(x1, y1, x2, y2);
	} else { 
		const int x = rX, y = rY;
		if(x >= 0 && x < SCREEN_W && y >= 0 && y < SCREEN_H) {
			Uint8 *d = &screen_layers[y * SCREEN_W + x];
			*d = (*d & mask) | color;
			screen_change(x, y, x+1, y+1);
		}
		if(rMX) rX++;
		if(rMY) rY++;
	}
}

static void screen_deo_sprite(void)
{
	int i, j, x = rX, y = rY;
	const int ctrl = dev[0x2f];
	const int flipx = ctrl & 0x10, dx = flipx ? -rDY : rDY;
	const int flipy = ctrl & 0x20, dy = flipy ? -rDX : rDX;
	const int row_start = flipx ? 0 : 7, row_delta = flipx ? 1 : -1;
	const int col_start = flipy ? 7 : 0, col_delta = flipy ? -1 : 1;
	const int layer = ctrl & 0x40, layer_mask = layer ? 0x3 : 0xc;
	const int mode_2bpp = ctrl >> 7, addr_2bpp = mode_2bpp ? rMA << 2 : rMA << 1;
	const int blend = ctrl & 0xf;
	const Uint8 opaque = alpha_lut[blend];
	const Uint8 *table = blend_lut[blend][layer >> 6];
	for(i = 0; i <= rML; i++, x += dx, y += dy, rA += addr_2bpp) {
		if(x < -8 || x >= SCREEN_W || y < -8 || y >= SCREEN_H) continue;
		const Uint8 *col = &ram[rA + col_start];
		for(j = 0; j < 8; j++, col += col_delta) {
			const int ch1 = *col, ch2 = mode_2bpp ? col[8] : 0;
			for(int k = 0, row = row_start; k < 8; k++, row += row_delta) {
				int px = x + k, py = y + j;
				if(px < 0 || px >= SCREEN_W || py < 0 || py >= SCREEN_H) continue;
				const int color = ((ch1 >> row) & 1) | (((ch2 >> row) & 1) << 1);
				if(opaque || color) {
					Uint8 *d = &screen_layers[py * SCREEN_W + px];
					*d = (*d & layer_mask) | table[color];
				}
			}
		}
	}
	{
		int x1 = flipx ? x : rX, x2 = flipx ? rX : x;
		int y1 = flipy ? y : rY, y2 = flipy ? rY : y;
		screen_change(x1-8, y1-8, x2+8, y2+8);
	}
	if(rMX) rX += flipx ? -rDX : rDX;
	if(rMY) rY += flipy ? -rDY : rDY;
}

static void screen_flush(void)
{
	if(!scr_dirty) return;
	int x1 = scr_x1 < 0 ? 0 : scr_x1;
	int y1 = scr_y1 < 0 ? 0 : scr_y1;
	int x2 = scr_x2 > SCREEN_W ? SCREEN_W : scr_x2;
	int y2 = scr_y2 > SCREEN_H ? SCREEN_H : scr_y2;
	for(int y = y1; y < y2; y++) {
		for(int x = x1; x < x2; x++) {
			Uint8 idx = screen_layers[y * SCREEN_W + x];
			M5.Lcd.drawPixel(x, y, screen_palette[idx & 0xf]);
		}
	}
	scr_x1 = scr_y1 = scr_x2 = scr_y2 = scr_dirty = 0;
}

static void screen_deo_auto(void)
{
	rMX = dev[0x26] & 0x1;
	rMY = dev[0x26] & 0x2;
	rMA = dev[0x26] & 0x4;
	rML = dev[0x26] >> 4;
	rDX = rMX << 3;
	rDY = rMY << 2;
}

static Uint8 screen_dei_rxhb(void) { return rX >> 8; }
static Uint8 screen_dei_rxlb(void) { return rX; }
static Uint8 screen_dei_ryhb(void) { return rY >> 8; }
static Uint8 screen_dei_rylb(void) { return rY; }
static Uint8 screen_dei_rahb(void) { return rA >> 8; }
static Uint8 screen_dei_ralb(void) { return rA; }

static void screen_deo_vector(void) { screen_vector = peek2(&dev[0x20]); }
static void screen_deo_x(void) { rX = TWOS(peek2(&dev[0x28])); }
static void screen_deo_y(void) { rY = TWOS(peek2(&dev[0x2a])); }
static void screen_deo_addr(void) { rA = peek2(&dev[0x2c]); }

/*
@|Controller -------------------------------------------------------- */

static int controller_vector = 0;

static void controller_down(Uint8 mask) {
	if (!mask) return;
	dev[0x82] |= mask;
	if (controller_vector) uxn_eval(controller_vector);
}

static void controller_up(Uint8 mask) {
	if (!mask) return;
	dev[0x82] &= ~mask;
	if (controller_vector) uxn_eval(controller_vector);
}

static void controller_deo_vector(void) {
	controller_vector = peek2(&dev[0x80]);
}

static inline Uint8 emu_dei(const Uint8 port) {
	switch (port) {
	case 0x28: return screen_dei_rxhb();
	case 0x29: return screen_dei_rxlb();
	case 0x2a: return screen_dei_ryhb();
	case 0x2b: return screen_dei_rylb();
	case 0x2c: return screen_dei_rahb();
	case 0x2d: return screen_dei_ralb();
	default: return dev[port];
	}
}

static inline void emu_deo(const Uint8 port, const Uint8 value) {
	dev[port] = value;
	switch (port) {
	case 0x08:
	case 0x09:
	case 0x0a:
	case 0x0b:
	case 0x0c:
	case 0x0d: system_deo_colorize(); break;
	case 0x21: screen_deo_vector(); break;
	case 0x26: screen_deo_auto(); break;
	case 0x28:
	case 0x29: screen_deo_x(); break;
	case 0x2a:
	case 0x2b: screen_deo_y(); break;
	case 0x2c:
	case 0x2d: screen_deo_addr(); break;
	case 0x2e: screen_deo_pixel(); break;
	case 0x2f: screen_deo_sprite(); break;
	case 0x81: controller_deo_vector(); break;
	}
}

/*
@|Core -------------------------------------------------------------- */

#define OPC(opc, A, B) {\
	case 0x00|opc: {const Uint8 d=0,r=0;A B} break;\
	case 0x20|opc: {const Uint8 d=1,r=0;A B} break;\
	case 0x40|opc: {const Uint8 d=0,r=1;A B} break;\
	case 0x60|opc: {const Uint8 d=1,r=1;A B} break;\
	case 0x80|opc: {const Uint8 d=0,r=0,k=ptr[0];A ptr[0]=k;B} break;\
	case 0xa0|opc: {const Uint8 d=1,r=0,k=ptr[0];A ptr[0]=k;B} break;\
	case 0xc0|opc: {const Uint8 d=0,r=1,k=ptr[1];A ptr[1]=k;B} break;\
	case 0xe0|opc: {const Uint8 d=1,r=1,k=ptr[1];A ptr[1]=k;B} break;}
#define DEC(m) stk[m][--ptr[m]]
#define INC(m) stk[m][ptr[m]++]
#define IMM a = ram[pc++] << 8, a |= ram[pc++];
#define MOV pc = d ? (Uint16)a : pc + (Sint8)a;
#define POx(o,m) o = DEC(r); if(m) o |= DEC(r) << 8;
#define PUx(i,m,s) if(m) c = (i), INC(s) = c >> 8, INC(s) = c; else INC(s) = i;
#define GOT(o) if(d) o[1] = DEC(r); o[0] = DEC(r);
#define PUT(i,s) INC(s) = i[0]; if(d) INC(s) = i[1];
#define DEO(o,v) emu_deo(o, v[0]); if(d) emu_deo(o + 1, v[1]);
#define DEI(i,v) v[0] = emu_dei(i); if(d) v[1] = emu_dei(i + 1); PUT(v,r)
#define POK(o,v,m) ram[o] = v[0]; if(d) ram[(o + 1) & m] = v[1];
#define PEK(i,v,m) v[0] = ram[i]; if(d) v[1] = ram[(i + 1) & m]; PUT(v,r)

static unsigned int uxn_eval(Uint16 pc)
{
	unsigned int a, b, c;
	Uint16 x[2], y[2], z[2];
	for(;;)
	switch(ram[pc++]) {
	case 0x00: return 1;
	case 0x20: if(DEC(0)) { IMM pc += a; } else pc += 2; break;
	case 0x40: IMM pc += a; break;
	case 0x60: IMM PUx(pc, 1, 1) pc += a; break;
	case 0xa0: INC(0) = ram[pc++]; /* fall-through */
	case 0x80: INC(0) = ram[pc++]; break;
	case 0xe0: INC(1) = ram[pc++]; /* fall-through */
	case 0xc0: INC(1) = ram[pc++]; break;
	OPC(0x01,POx(a,d),PUx(a+1,d,r))
	OPC(0x02,ptr[r] -= 1+d;,{})
	OPC(0x03,GOT(x) ptr[r] -= 1+d;,PUT(x,r))
	OPC(0x04,GOT(x) GOT(y),PUT(x,r) PUT(y,r))
	OPC(0x05,GOT(x) GOT(y) GOT(z),PUT(y,r) PUT(x,r) PUT(z,r))
	OPC(0x06,GOT(x),PUT(x,r) PUT(x,r))
	OPC(0x07,GOT(x) GOT(y),PUT(y,r) PUT(x,r) PUT(y,r))
	OPC(0x08,POx(a,d) POx(b,d),PUx(b==a,0,r))
	OPC(0x09,POx(a,d) POx(b,d),PUx(b!=a,0,r))
	OPC(0x0a,POx(a,d) POx(b,d),PUx(b>a,0,r))
	OPC(0x0b,POx(a,d) POx(b,d),PUx(b<a,0,r))
	OPC(0x0c,POx(a,d),MOV)
	OPC(0x0d,POx(a,d) POx(b,0),if(b) MOV)
	OPC(0x0e,POx(a,d),PUx(pc,1,!r) MOV)
	OPC(0x0f,GOT(x),PUT(x,!r))
	OPC(0x10,POx(a,0),PEK(a,x,0xff))
	OPC(0x11,POx(a,0) GOT(y),POK(a,y,0xff))
	OPC(0x12,POx(a,0),PEK(pc+(Sint8)a,x,0xffff))
	OPC(0x13,POx(a,0) GOT(y),POK(pc+(Sint8)a,y,0xffff))
	OPC(0x14,POx(a,1),PEK(a,x,0xffff))
	OPC(0x15,POx(a,1) GOT(y),POK(a,y,0xffff))
	OPC(0x16,POx(a,0),DEI(a,x))
	OPC(0x17,POx(a,0) GOT(y),DEO(a,y))
	OPC(0x18,POx(a,d) POx(b,d),PUx(b+a,d,r))
	OPC(0x19,POx(a,d) POx(b,d),PUx(b-a,d,r))
	OPC(0x1a,POx(a,d) POx(b,d),PUx(b*a,d,r))
	OPC(0x1b,POx(a,d) POx(b,d),PUx(a?b/a:0,d,r))
	OPC(0x1c,POx(a,d) POx(b,d),PUx(b&a,d,r))
	OPC(0x1d,POx(a,d) POx(b,d),PUx(b|a,d,r))
	OPC(0x1e,POx(a,d) POx(b,d),PUx(b^a,d,r))
	OPC(0x1f,POx(a,0) POx(b,d),PUx(b>>(a&0xf)<<(a>>4),d,r))
	} return 0;
}

void setup(void) {
	M5.begin();
	M5.Lcd.setRotation(1);
	M5.Lcd.fillScreen(TFT_BLACK);
	/* Load */
	memcpy(&ram[0x100], ROM_DATA, sizeof(ROM_DATA));
	/* Set screen size */
	poke2(&dev[0x22], SCREEN_W);
	poke2(&dev[0x24], SCREEN_H);
	/* On Reset */
	uxn_eval(0x100);
}

static uint8_t btn_prev = 0;

void loop(void) {
	M5.update();
	/* Controller */
	uint8_t btn_now = 0;
	if (M5.BtnA.isPressed()) btn_now |= 0x01; /* A */
	if (M5.BtnB.isPressed()) btn_now |= 0x02; /* B */
	uint8_t changed = btn_prev ^ btn_now;
	if (changed) {
		uint8_t pressed = changed & btn_now;
		uint8_t released = changed & ~btn_now;
		if (pressed) controller_down(pressed);
		if (released) controller_up(released);
		btn_prev = btn_now;
	}
	/* Screen Update */
	if(screen_vector) uxn_eval(screen_vector);
	screen_flush();
	delay(16); /* 60fps!! */
}