#define _GNU_SOURCE

#ifdef _WIN32
#define NOMINMAX
#undef NOCRYPT
// For mingw-w64, use _aligned_malloc instead of aligned_alloc
#define aligned_alloc(align, size) _aligned_malloc(size, align)
#define aligned_free(ptr) _aligned_free(ptr)
#else
#include <sys/mman.h>
#include <sys/stat.h>        /* For mode constants */
#include <fcntl.h>           /* For O_* constants */
#define aligned_free(ptr) free(ptr)
#endif

#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <stdbool.h>
#include <stddef.h>
#include <math.h>
#include <string.h>
#include <SDL2/SDL.h>

static size_t state_dump_count;
static FILE *state_dump_file;

#define BUFFER_WIDTH 256
#define BUFFER_HEIGHT 240
#define WINDOW_WIDTH (256 * 3 + 256 * 2)
#define WINDOW_HEIGHT (240 * 3)

#define PRG_ROM_SIZE	(2 * 1024 * 1024)
#define CHR_ROM_SIZE	(1 * 1024 * 1024)
// #define PRG_ROM_SIZE	(512 * 1024)
// #define CHR_ROM_SIZE	(256 * 1024)

#define PIXELS_SIZE	(256 * 240)
#define RAM_SIZE		0x1000		// 0x800 in reality, but for aligned alloc it must be the size of the alignment (4096)
#define SRAM_SIZE		0x2000
#define CIRAM_SIZE	0x1000
#define CHR_RAM_SIZE	0x4000

static uint32_t buffer[BUFFER_WIDTH * BUFFER_HEIGHT] __attribute__((section(".bss"), aligned(4096)));
static uint32_t display_buffer[BUFFER_WIDTH * BUFFER_HEIGHT] __attribute__((section(".bss"), aligned(4096)));

static void audio_callback(int16_t *data, size_t frames) { }

#define FRAME_INTERVAL_NS (1000000000ULL / 60.0988)

#define DEBUG_PRINT(fmt, ...) printf(fmt, ##__VA_ARGS__)
#ifdef _WIN32
#include "win32_timer.c"
#else
#include "linux_timer.c"
#endif

// #include "audio.c"
#include "incbin.h"

#ifdef BENCHMARK
// Embed the ROM for benchmarking to eliminate file I/O overhead
// Uncomment the ROM you want to benchmark:
// INCBIN_BYTES(benchmark_rom, "data/Life Force (USA).nes");
// INCBIN_BYTES(benchmark_rom, "data/0001/Metroid (U) [!].nes");
INCBIN_BYTES(benchmark_rom, "data/0000/Super Mario Bros. (World) (HVC-SM).nes");
// INCBIN_BYTES(benchmark_rom, "data/0003/Gradius (USA).nes");
#endif

static int32_t frames;		// debug information
// static int32_t tas_frame;

// #include "helpers/battletoads_tas.h"	// REMOVE ME
// #include "helpers/smb_tas.h"	// REMOVE ME

// NES core
#include "mknes_mapper.h"
#include "mknes.h"
#include "mknes_apu.c"
#include "mknes_ppu.c"
#include "mknes_ppu_registers.c"
#include "mknes_memory.c"
#include "mknes_cpu.c"
#include "mknes_ines2.c"
#include "mknes_mapper.c"

#ifdef BENCHMARK
#include "mknes_bench.c"
#endif

// static void dump_state(struct nes_state *state) {
// 	size_t state_size = offsetof(struct nes_state, ram);

// 	if(!state_dump_file) {
// 		state_dump_file = fopen("state_dump.bin", "wb");
// 		if(!state_dump_file) {
// 			fprintf(stderr, "Failed to open state_dump.bin for writing\n");
// 			return;
// 		}
// 	}

// 	fwrite(state, 1, state_size, state_dump_file);
// 	state_dump_count++;
// }

int main(int argc, char **argv) {
	setbuf(stdout, 0);
	init_opcode_lut();
	init_opcode_ud_lut();

	struct nes_state *nstate = aligned_alloc(4096, (sizeof(struct nes_state) + 4095) & ~4095);

#ifdef BENCHMARK
	// Run benchmark with configurable parameters
	uint32_t num_runs = 10;
	uint32_t frames_per_run = 0x1000;

	// Parse command line arguments
	for(int i = 1; i < argc; i++) {
		if(strcmp(argv[i], "-n") == 0 && i + 1 < argc) {
			num_runs = atoi(argv[i + 1]);
			i++;
		} else if(strcmp(argv[i], "-f") == 0 && i + 1 < argc) {
			frames_per_run = atoi(argv[i + 1]);
			i++;
		}
	}

	run_benchmark(nstate, num_runs, frames_per_run);
	return 0;
#else
	memset(nstate, 0, sizeof(struct nes_state));
	ppu_reset(nstate);
	// ines2_load(nstate, "data/0000/10-Yard Fight (USA, Europe).nes");
	// ines2_load(nstate, "data/0000/Balloon Fight (USA).nes");
	// ines2_load(nstate, "data/0000/Excitebike (Japan, USA).nes");
	// ines2_load(nstate, "data/0000/Ice Climber (USA, Europe, Korea).nes");
	// ines2_load(nstate, "data/0000/Kung Fu (Japan, USA).nes");
	// ines2_load(nstate, "data/0000/Super Mario Bros. (World) (HVC-SM).nes");
	// ines2_load(nstate, "data/Super Mario Bros. (W) (V1.0) [!].nes");
	// ines2_load(nstate, "data/Super Mario Bros. (JU) [!].nes");
	// ines2_load(nstate, "data/0000/Urban Champion (World).nes");
	// ines2_load(nstate, "data/0000/Wrecking Crew (World).nes");
	// ines2_load(nstate, "data/0000/scanline.nes");
	// ines2_load(nstate, "data/0000/Sayoonara!.NES");
	// ines2_load(nstate, "data/0000/raster_demos/RasterChromaLuma.NES");
	// ines2_load(nstate, "data/0000/raster_demos/RasterTest1.NES");
	// ines2_load(nstate, "data/0000/raster_demos/RasterTest2.NES");
	// ines2_load(nstate, "data/0000/raster_demos/RasterTest3.NES");
	// ines2_load(nstate, "data/0000/raster_demos/RasterTest3a.NES");
	// ines2_load(nstate, "data/0000/raster_demos/RasterTest3b.NES");
	// ines2_load(nstate, "data/0000/raster_demos/RasterTest3c.NES");
	// ines2_load(nstate, "data/0000/raster_demos/RasterTest3d.NES");
	// ines2_load(nstate, "data/0000/raster_demos/RasterTest3e.NES");
	// ines2_load(nstate, "data/0000/NEStress.NES");
	// ines2_load(nstate, "data/0000/Super Mario Bros. (World) (HVC-SM).zip");
	// ines2_load(nstate, "data/0042/Super Mario Bros. + Duck Hunt (USA).zip");
	// ines2_load(nstate, "data/0000/Xevious - The Avenger (USA).zip");
	// ines2_load(nstate, "data/tv.nes");

	// ines2_load(nstate, "data/Life Force (USA).nes");	// 2002

	// ines2_load(nstate, "data/0003/Flipull - An Exciting Cube Game (Japan) (En).zip");
	// ines2_load(nstate, "data/0003/Friday the 13th (USA).zip");
	// ines2_load(nstate, "data/0003/Ghostbusters (Japan).zip");
	// ines2_load(nstate, "data/0003/Gradius (USA).zip");

	// ines2_load(nstate, "data/0004/Mega Man 3 (USA).zip");
	// ines2_load(nstate, "data/0004/Mega Man 4 (USA).zip");
	// ines2_load(nstate, "data/0004/Mega Man 5 (USA).zip");
	// ines2_load(nstate, "data/0004/Mega Man 6 (USA).zip");
	// ines2_load(nstate, "data/0004/Super Mario Bros. 2 (USA).zip");

	ines2_load(nstate, "data/0005/Castlevania III - Dracula's Curse (USA).zip");
	// ines2_load(nstate, "data/0005/Metal Slader Glory (Japan).zip");
	// ines2_load(nstate, "data/0007/Battletoads (USA).zip");
	// ines2_load(nstate, "data/0007/Beetlejuice (USA).zip");
	// ines2_load(nstate, "data/0007/Cabal (USA).zip");

	// ines2_load(nstate, "data/000b/Baby Boomer (USA) (Unl).zip");
	// ines2_load(nstate, "data/000b/Captain Comic - The Adventure (USA) (Unl).zip");
	// ines2_load(nstate, "data/000b/King Neptune's Adventure (USA) (Unl).nes");

	// ines2_load(nstate, "data/2002/Attack Animal Gakuen (Japan).zip");
	// ines2_load(nstate, "data/2002/Ballblazer (Japan).zip");
	// ines2_load(nstate, "data/2002/Best of the Best - Championship Karate (USA).zip");

	// ines2_load(nstate, "data/0001/Kid Icarus (UE) (V1.1) [!].nes");
	// ines2_load(nstate, "data/0001/Metroid (U) [!].nes");
	// ines2_load(nstate, "data/0001/Legend of Zelda, The (U) (V1.1) [!].nes");

	// ines2_load(nstate, "data/Blaster Master (USA).zip");	// mapper 1
	// ines2_load(nstate, "AccuracyCoin.nes");	// mapper 1

	mapper_setup(nstate);
	cpu_reset(nstate);

	// SDL Setup
	SDL_Init(SDL_INIT_VIDEO);
	SDL_Window *window = SDL_CreateWindow( "mknes SDL", SDL_WINDOWPOS_CENTERED, SDL_WINDOWPOS_CENTERED, WINDOW_WIDTH, WINDOW_HEIGHT, SDL_WINDOW_SHOWN | SDL_WINDOW_RESIZABLE );
	SDL_Renderer *renderer = SDL_CreateRenderer(window, -1, SDL_RENDERER_ACCELERATED | SDL_RENDERER_PRESENTVSYNC);
	SDL_Texture *texture = SDL_CreateTexture( renderer, SDL_PIXELFORMAT_RGBA8888, SDL_TEXTUREACCESS_STREAMING, BUFFER_WIDTH, BUFFER_HEIGHT );

	SDL_Rect game_viewport;
	game_viewport.x = 0;
	game_viewport.y = 0;
	game_viewport.w = 256 * 3;
	game_viewport.h = 240 * 3;

	// Debug nametable textures (256x240 each)
	SDL_Texture *debug_nt_texture[4];
	uint32_t *debug_nt_buffer[4];
	for(int i = 0; i < 4; i++) {
		debug_nt_texture[i] = SDL_CreateTexture(renderer, SDL_PIXELFORMAT_RGBA8888, SDL_TEXTUREACCESS_STREAMING, 256, 240);
		debug_nt_buffer[i] = malloc(256 * 240 * sizeof(uint32_t));
	}

	SDL_Rect nt_viewport[4];
	for(int i = 0; i < 4; i++) {
		nt_viewport[i].x = 256 * 3 + (i % 2) * 256;
		nt_viewport[i].y = (i / 2) * 240;
		nt_viewport[i].w = 256;
		nt_viewport[i].h = 240;
	}

	timer_init();
	struct timer_handle *timer = timer_new(FRAME_INTERVAL_NS);

	uint8_t running = true;
	const uint8_t *keyboard_state = SDL_GetKeyboardState(NULL);
	uint32_t palette[] = { 0x000000ff, 0xffffffff, 0xaaaaaaff, 0x555555ff };

	while(running) {
		SDL_Event event;
		while(SDL_PollEvent(&event)) {
			if(event.type == SDL_QUIT) {
				running = false;
			}
			if(event.type == SDL_KEYDOWN && event.key.keysym.sym == SDLK_ESCAPE) {
				running = false;
			}
		}

		// Render all 4 nametables through mapper to see mirroring
		for(int nt_idx = 0; nt_idx < 4; nt_idx++) {
			uint32_t nt_base_addr = 0x2000 + nt_idx * 0x400;

			// Render 32x30 tiles
			for(uint32_t ty = 0; ty < 30; ty++) {
				for(uint32_t tx = 0; tx < 32; tx++) {
					uint32_t nt_tile_addr = nt_base_addr + ty * 32 + tx;
					uint8_t tile_index = nstate->mapper_function.ciram_read(nstate, nt_tile_addr & 0xfff);

					// Determine which pattern table to use based on PPU control register
					uint32_t pattern_table = (nstate->ppu.reg_ctrl & 0x10) ? 0x1000 : 0x0000;
					uint32_t chr_tile_addr = pattern_table + tile_index * 16;

					// Render 8x8 tile using mapper's CHR read function
					for(uint32_t py = 0; py < 8; py++) {
						uint8_t plane0 = nstate->mapper_function.chr_read(nstate, chr_tile_addr + py);
						uint8_t plane1 = nstate->mapper_function.chr_read(nstate, chr_tile_addr + py + 8);

						for(uint32_t px = 0; px < 8; px++) {
							uint8_t bit0 = (plane0 >> (7 - px)) & 1;
							uint8_t bit1 = (plane1 >> (7 - px)) & 1;
							uint8_t color_index = (bit1 << 1) | bit0;

							uint32_t color = palette[color_index];

							uint32_t screen_x = tx * 8 + px;
							uint32_t screen_y = ty * 8 + py;
							debug_nt_buffer[nt_idx][screen_y * 256 + screen_x] = color;
						}
					}
				}
			}

			SDL_UpdateTexture(debug_nt_texture[nt_idx], 0, debug_nt_buffer[nt_idx], 256 * sizeof(uint32_t));
		}

		// Joypad input
		uint8_t input = 0;
		if(keyboard_state[SDL_SCANCODE_X]) { input |= (1 << 0); }
		if(keyboard_state[SDL_SCANCODE_Z]) { input |= (1 << 1); }
		if(keyboard_state[SDL_SCANCODE_SPACE]) { input |= (1 << 2); }
		if(keyboard_state[SDL_SCANCODE_RETURN]) { input |= (1 << 3); }
		if(keyboard_state[SDL_SCANCODE_UP]) { input |= (1 << 4); }
		if(keyboard_state[SDL_SCANCODE_DOWN]) { input |= (1 << 5); }
		if(keyboard_state[SDL_SCANCODE_LEFT]) { input |= (1 << 6); }
		if(keyboard_state[SDL_SCANCODE_RIGHT]) { input |= (1 << 7); }
		nstate->ppu.input[0] = input;

		// Run NES emulation for one frame
		while(!nstate->ppu.frame_ready) {
			cpu_tick(nstate);
		}
		nstate->ppu.frame_ready = 0;
		// if(nstate->ppu.open_bus > 0) {
		// 	int32_t v = nstate->ppu.open_bus;
		// 	v -= 5;
		// 	if(v < 0) {
		// 		v = 0;
		// 	}
		// 	nstate->ppu.open_bus = v;
		// }
		// nstate->ppu.input[0] = tas_input[tas_frame++];

		frames++;

		// // Dump state every frame starting from 2400
		// if(frames >= 2400 && frames <= 3100) {
		// 	dump_state(nstate);
		// }

		// Convert NES pixels to display buffer
		uint32_t * restrict dst = display_buffer;
		uint8_t * restrict src = nstate->pixels;
		for(uint32_t y = 0; y < 240; ++y) {
			for(uint32_t x = 0; x < 256; ++x) {
				uint8_t val = *src++;
				if(val >= 64) val = 0;
				dst[x] = nes_palette[val];
			}
			dst += BUFFER_WIDTH;
		}

		// Update SDL texture and render
		SDL_UpdateTexture(texture, 0, display_buffer, BUFFER_WIDTH * sizeof(uint32_t));
		SDL_RenderClear(renderer);

		// Render game on left side
		SDL_RenderCopy(renderer, texture, 0, &game_viewport);

		// Render all 4 nametables on right side
		for(int i = 0; i < 4; i++) {
			SDL_RenderCopy(renderer, debug_nt_texture[i], 0, &nt_viewport[i]);
		}

		SDL_RenderPresent(renderer);

		if(!keyboard_state[SDL_SCANCODE_F]) {
			timer_wait(timer);
		}
	}

	printf("total frames: %6d  total cycles: %12llu\n", frames, (unsigned long long)nstate->cpu.cycles);
	// printf("state dumps created: %zu\n", state_dump_count);

	// if(state_dump_file) {
	// 	fclose(state_dump_file);
	// }

	timer_destroy(timer);
	for(int i = 0; i < 4; i++) {
		free(debug_nt_buffer[i]);
	}
	free(nstate);
	timer_shutdown();

	for(int i = 0; i < 4; i++) {
		SDL_DestroyTexture(debug_nt_texture[i]);
	}
	SDL_DestroyTexture(texture);
	SDL_DestroyRenderer(renderer);
	SDL_DestroyWindow(window);
	SDL_Quit();

	return 0;
#endif
}