const std = @import("std");

const contract = @import("contract");
const _cpu = @import("cpu");

const Bus = @import("bus.zig");
const Cartridge = @import("cartridge.zig");

const controller = @import("controller.zig");
pub const Button = controller.Button;

// Adjust this enum spelling to whatever your CPU package
// currently exports for the Ricoh 2A03.
const Cpu = _cpu.m6502.Cpu(
    Bus,
    .ricoh2a03,
    .cycle,
);

const Nes = @This();
pub const spec = contract.SystemSpec{
    .name = "Nintendo Entertainment System",

    .video_outputs = &.{
        .{
            .max_width = 256,
            .max_height = 240,

            .format = .rgb565,

            .aspect_ratio = .{
                .numerator = 4,
                .denominator = 3,
            },

            .refresh_rate = .{
                .numerator = 60,
                .denominator = 1,
            },
        },
    },

    .audio_outputs = &.{
        .{
            .sample_rate = 44_100,
            .channels = 1,
            .format = .i16,
        },
    },

    .input_devices = &.{
        .{
            .button_count = 8,
        },
        .{
            .button_count = 8,
        },
    },

    .storage_devices = &.{
        .{
            .name = "Battery PRG RAM",
            .min_size = 0,
            .max_size = 8192,
            .persistent = true,
            .removable = false,
            .writable = true,
        },
    },

    .save_state = .{
        .max_size = 32 * 1024,
    },
};

bus: Bus = undefined,
cpu: Cpu = undefined,

/// Initialize in place.
///
/// `rom` must remain valid for the lifetime of this NES
/// because Cartridge stores slices into it.
pub fn init(
    self: *Nes,
    rom: []const u8,
) Cartridge.Error!void {
    self.bus =
        try Bus.init(rom);

    self.cpu =
        Cpu.init(&self.bus);

    self.reset();
}

pub fn reset(
    self: *Nes,
) void {
    self.bus.reset();
    self.cpu.reset();
}

pub fn step(
    self: *Nes,
) ?u8 {
    const cycles = self.cpu.step() orelse return null;
    self.serviceOamDma();
    return cycles;
}

pub fn runFrame(
    self: *Nes,
) void {
    self.bus.ppu.frame_ready = false;
    self.bus.apu.sample_count = 0;

    while (!self.bus.ppu.frame_ready) {
        _ = self.step() orelse
            return;
    }
}

pub const conformance = @import("conformance.zig");

test {
    _ = conformance;
}

comptime {
    contract.validateSystem(Nes);
}

pub fn videoFrame(
    self: *const Nes,
    output: usize,
) contract.VideoFrame {
    std.debug.assert(output == 0);

    return .{
        .data = std.mem.sliceAsBytes(
            self.bus.ppu.framebuffer[0..],
        ),

        .width = 256,
        .height = 240,

        .pitch = 256 * @sizeOf(u16),

        .format = .rgb565,

        .frame_number = self.bus.ppu.frame_number,
    };
}

pub fn audioBuffer(
    self: *const Nes,
    output: usize,
) contract.AudioBuffer {
    std.debug.assert(output == 0);

    const samples = self.bus.apu.sample_buffer[0..self.bus.apu.sample_count];

    return .{
        .data = std.mem.sliceAsBytes(samples),
        .frames = self.bus.apu.sample_count,

        .sample_rate = 44_100,
        .channels = 1,
        .format = .i16,
    };
}

pub fn setInput(
    self: *Nes,
    device: usize,
    input: contract.DeviceInput,
) void {
    std.debug.assert(device < 2);

    self.bus.controllers[device]
        .setButtons(
        @truncate(input.buttons),
    );
}

pub fn storageView(
    self: *const Nes,
    slot: usize,
) contract.StorageView {
    std.debug.assert(slot == 0);
    return .{
        .data = &self.bus.cartridge.prg_ram,
        .generation = 1,
    };
}

pub fn loadStorage(
    self: *Nes,
    slot: usize,
    data: []const u8,
) contract.StorageLoadError!void {
    if (slot != 0) return error.InvalidSlot;
    if (data.len > self.bus.cartridge.prg_ram.len) return error.InvalidSize;
    @memcpy(self.bus.cartridge.prg_ram[0..data.len], data);
}

fn serviceOamDma(
    self: *Nes,
) void {
    const page =
        self.bus.takeOamDmaPage() orelse return;

    const odd_cycle =
        (self.cpu.cycles & 1) != 0;

    // Mandatory DMA halt cycle.
    self.cpu.tick();

    // Alignment cycle when necessary.
    if (odd_cycle)
        self.cpu.tick();

    const base =
        @as(u16, page) << 8;

    for (0..256) |i| {
        // DMA read cycle.
        const value =
            self.cpu.readCycle(
                base |
                    @as(u16, @intCast(i)),
            );

        // DMA write cycle.
        self.bus.ppu.writeOamDma(value);
        self.cpu.tick();
    }
}

pub fn saveState(
    self: *const Nes,
    buffer: []u8,
) contract.state.Error!usize {
    var offset: usize = 0;

    if (buffer.len < 8) return error.BufferTooSmall;
    @memcpy(buffer[offset..][0..8], "6SOZNES1");
    offset += 8;

    // CPU
    if (offset + 15 > buffer.len) return error.BufferTooSmall;
    std.mem.writeInt(u16, buffer[offset..][0..2], self.cpu.registers.pc, .little);
    offset += 2;
    buffer[offset] = self.cpu.registers.a;
    offset += 1;
    buffer[offset] = self.cpu.registers.x;
    offset += 1;
    buffer[offset] = self.cpu.registers.y;
    offset += 1;
    buffer[offset] = self.cpu.registers.sp;
    offset += 1;
    buffer[offset] = @bitCast(self.cpu.registers.status);
    offset += 1;
    std.mem.writeInt(u64, buffer[offset..][0..8], self.cpu.cycles, .little);
    offset += 8;

    // Bus RAM & PRG RAM
    if (offset + self.bus.ram.len + self.bus.cartridge.prg_ram.len > buffer.len) return error.BufferTooSmall;
    @memcpy(buffer[offset..][0..self.bus.ram.len], &self.bus.ram);
    offset += self.bus.ram.len;
    @memcpy(buffer[offset..][0..self.bus.cartridge.prg_ram.len], &self.bus.cartridge.prg_ram);
    offset += self.bus.cartridge.prg_ram.len;

    if (self.bus.cartridge.chr_is_ram) {
        if (offset + self.bus.cartridge.chr_ram.len > buffer.len) return error.BufferTooSmall;
        @memcpy(buffer[offset..][0..self.bus.cartridge.chr_ram.len], &self.bus.cartridge.chr_ram);
        offset += self.bus.cartridge.chr_ram.len;
    }

    // PPU State
    const ppu_size = 1 + 1 + 1 + 1 + 2 + 2 + 1 + 1 + 2 + 2 + 4096 + 32 + 256;
    if (offset + ppu_size > buffer.len) return error.BufferTooSmall;

    buffer[offset] = self.bus.ppu.ctrl;
    offset += 1;
    buffer[offset] = self.bus.ppu.mask;
    offset += 1;
    buffer[offset] = self.bus.ppu.status;
    offset += 1;
    buffer[offset] = self.bus.ppu.oam_addr;
    offset += 1;
    std.mem.writeInt(u16, buffer[offset..][0..2], self.bus.ppu.vram_addr, .little);
    offset += 2;
    std.mem.writeInt(u16, buffer[offset..][0..2], self.bus.ppu.temp_addr, .little);
    offset += 2;
    buffer[offset] = self.bus.ppu.fine_x;
    offset += 1;
    buffer[offset] = if (self.bus.ppu.write_latch) 1 else 0;
    offset += 1;
    std.mem.writeInt(u16, buffer[offset..][0..2], self.bus.ppu.dot, .little);
    offset += 2;
    std.mem.writeInt(u16, buffer[offset..][0..2], self.bus.ppu.scanline, .little);
    offset += 2;

    @memcpy(buffer[offset..][0..4096], &self.bus.ppu.nametable);
    offset += 4096;
    @memcpy(buffer[offset..][0..32], &self.bus.ppu.palette);
    offset += 32;
    @memcpy(buffer[offset..][0..256], &self.bus.ppu.oam);
    offset += 256;

    return offset;
}

pub fn loadState(
    self: *Nes,
    buffer: []const u8,
) contract.state.Error!void {
    var offset: usize = 0;

    if (buffer.len < 8) return error.InvalidState;
    if (!std.mem.eql(u8, buffer[0..8], "6SOZNES1")) return error.InvalidState;
    offset += 8;

    // CPU
    if (offset + 15 > buffer.len) return error.InvalidState;
    self.cpu.registers.pc = std.mem.readInt(u16, buffer[offset..][0..2], .little);
    offset += 2;
    self.cpu.registers.a = buffer[offset];
    offset += 1;
    self.cpu.registers.x = buffer[offset];
    offset += 1;
    self.cpu.registers.y = buffer[offset];
    offset += 1;
    self.cpu.registers.sp = buffer[offset];
    offset += 1;
    self.cpu.registers.status = @bitCast(buffer[offset]);
    offset += 1;
    self.cpu.cycles = std.mem.readInt(u64, buffer[offset..][0..8], .little);
    offset += 8;

    // Bus RAM & PRG RAM
    if (offset + self.bus.ram.len + self.bus.cartridge.prg_ram.len > buffer.len) return error.InvalidState;
    @memcpy(&self.bus.ram, buffer[offset..][0..self.bus.ram.len]);
    offset += self.bus.ram.len;
    @memcpy(&self.bus.cartridge.prg_ram, buffer[offset..][0..self.bus.cartridge.prg_ram.len]);
    offset += self.bus.cartridge.prg_ram.len;

    if (self.bus.cartridge.chr_is_ram) {
        if (offset + self.bus.cartridge.chr_ram.len > buffer.len) return error.InvalidState;
        @memcpy(&self.bus.cartridge.chr_ram, buffer[offset..][0..self.bus.cartridge.chr_ram.len]);
        offset += self.bus.cartridge.chr_ram.len;
    }

    // PPU State
    const ppu_size = 1 + 1 + 1 + 1 + 2 + 2 + 1 + 1 + 2 + 2 + 4096 + 32 + 256;
    if (offset + ppu_size > buffer.len) return error.InvalidState;

    self.bus.ppu.ctrl = buffer[offset];
    offset += 1;
    self.bus.ppu.mask = buffer[offset];
    offset += 1;
    self.bus.ppu.status = buffer[offset];
    offset += 1;
    self.bus.ppu.oam_addr = buffer[offset];
    offset += 1;
    self.bus.ppu.vram_addr = std.mem.readInt(u16, buffer[offset..][0..2], .little);
    offset += 2;
    self.bus.ppu.temp_addr = std.mem.readInt(u16, buffer[offset..][0..2], .little);
    offset += 2;
    self.bus.ppu.fine_x = @truncate(buffer[offset]);
    offset += 1;
    self.bus.ppu.write_latch = buffer[offset] != 0;
    offset += 1;
    self.bus.ppu.dot = std.mem.readInt(u16, buffer[offset..][0..2], .little);
    offset += 2;
    self.bus.ppu.scanline = std.mem.readInt(u16, buffer[offset..][0..2], .little);
    offset += 2;

    @memcpy(&self.bus.ppu.nametable, buffer[offset..][0..4096]);
    offset += 4096;
    @memcpy(&self.bus.ppu.palette, buffer[offset..][0..32]);
    offset += 32;
    @memcpy(&self.bus.ppu.oam, buffer[offset..][0..256]);
    offset += 256;
}

test "NES System - Save State Serialization Roundtrip" {
    // Construct dummy iNES NROM header + 16K PRG + 8K CHR
    var rom: [16 + 0x4000 + 0x2000]u8 = [_]u8{0} ** (16 + 0x4000 + 0x2000);
    @memcpy(rom[0..4], "NES\x1a");
    rom[4] = 1; // 16K PRG
    rom[5] = 1; // 8K CHR

    var nes: Nes = undefined;
    try nes.init(&rom);

    // Mutate state
    nes.cpu.registers.pc = 0x1234;
    nes.cpu.registers.a = 0x42;
    nes.bus.ram[0x05] = 0xAA;
    nes.bus.ppu.vram_addr = 0x2050;

    var save_buf: [32768]u8 = undefined;
    const bytes_written = try nes.saveState(&save_buf);

    // Reset nes
    nes.reset();
    try std.testing.expect(nes.cpu.registers.pc != 0x1234);

    // Load state
    try nes.loadState(save_buf[0..bytes_written]);
    try std.testing.expectEqual(@as(u16, 0x1234), nes.cpu.registers.pc);
    try std.testing.expectEqual(@as(u8, 0x42), nes.cpu.registers.a);
    try std.testing.expectEqual(@as(u8, 0xAA), nes.bus.ram[0x05]);
    try std.testing.expectEqual(@as(u16, 0x2050), nes.bus.ppu.vram_addr);
}

comptime {
    contract.validateSystem(Nes);
}