#include <stdlib.h>
#include <string.h>

#include "map.h"


typedef struct {
    uint8_t* prg_rom;
    int prg_rom_banks;
    uint8_t* chr_rom;
    int chr_rom_banks;

    uint8_t battery;

    uint8_t* chr_bank[2];
    uint8_t* vram_bank[4];
    uint8_t* prg_bank[2];
    uint8_t* chr_ram;
    uint8_t* chr;

    uint8_t reg_shift;
    uint8_t reg_n_shift;
    uint8_t reg_control;
    uint8_t reg_chr_0;
    uint8_t reg_chr_1;
    uint8_t reg_prg;

    uint8_t vram[2][nes_vram_page_size];
    uint8_t wram[nes_mem_wram_size];
} mmc1_mapper;


static void mmc1_update_vram(mmc1_mapper* map) {
    // VRAM selection
    int nametable = (map->reg_control & 0b11);
    if (nametable == 0) {
        MAP_LOG("VRAM: 1-screen, lower bank");
        map->vram_bank[0] = map->vram_bank[1] =
            map->vram_bank[2] = map->vram_bank[3] =
                map->vram[0];
    } else if (nametable == 1) {
        MAP_LOG("VRAM: 1-screen, upper bank");
        map->vram_bank[0] = map->vram_bank[1] =
            map->vram_bank[2] = map->vram_bank[3] =
                map->vram[1];
    } else if (nametable == 2) {
        MAP_LOG("VRAM: Vertical mirroring");
        map->vram_bank[0] = map->vram_bank[2] = map->vram[0];
        map->vram_bank[1] = map->vram_bank[3] = map->vram[1];
    } else if (nametable == 3) {
        MAP_LOG("VRAM: Horizontal mirroring");
        map->vram_bank[0] = map->vram_bank[1] = map->vram[0];
        map->vram_bank[2] = map->vram_bank[3] = map->vram[1];
    }
}

static void mmc1_update_chr(mmc1_mapper* map) {
    // CHR RAM selection
    if (!(map->reg_control & 0b10000)) {
        int bank = (map->reg_chr_0 & 0b11110);
        MAP_LOG("CHR: 8 KB: %d + %d", bank, bank + 1);
        map->chr_bank[0] = &map->chr[   bank *
                                        nes_chr_page_size];
        map->chr_bank[1] = &map->chr[   (bank + 1) *
                                        nes_chr_page_size];
    } else {
        MAP_LOG("CHR: %d + %d", map->reg_chr_0, map->reg_chr_1);
        map->chr_bank[0] = &map->chr[   map->reg_chr_0 *
                                        nes_chr_page_size];
        map->chr_bank[1] = &map->chr[   map->reg_chr_1 *
                                        nes_chr_page_size];
    }
}

static void mmc1_update_prg(mmc1_mapper* map) {
    // PRG ROM selection
    int mode = (map->reg_control >> 2) & 3;
    int bank_0 = (map->reg_prg & 0b01111);
    int bank_1 = bank_0;
    if (!(mode & 0b10)) {
        bank_0 = (bank_0 & 0b01110);
        bank_1 = bank_0 + 1;
    } else if (mode == 2) {
        bank_0 = 0;
    } else if (mode == 3) {
        bank_1 = (map->prg_rom_banks - 1);
    }
    MAP_LOG("PRG: %d + %d", bank_0, bank_1);
    map->prg_bank[0] = &map->prg_rom[bank_0 * nes_prg_rom_page_size];
    map->prg_bank[1] = &map->prg_rom[bank_1 * nes_prg_rom_page_size];
}

static void mmc1_reset(nes_mapper* nes_map) {
    mmc1_mapper* map = (mmc1_mapper*)nes_map->data;
    map->reg_shift   = 0b10000;
    map->reg_control = 0b01100;
    map->reg_chr_0   = 0;
    map->reg_chr_1   = 0;
    map->reg_prg     = 0;
    mmc1_update_prg(map);
    mmc1_update_chr(map);
    mmc1_update_vram(map);
}

static int mmc1_init(nes_mapper* nes_map, nes_cart* cart) {
    mmc1_mapper* map = calloc(1, sizeof(mmc1_mapper));
    nes_map->data = map;
    if (NULL != map) {
        map->prg_rom       = cart->prg_rom;
        map->prg_rom_banks = cart->prg_rom_banks;
        map->chr_rom_banks = cart->chr_rom_banks / 2;
        if (map->chr_rom_banks > 0) {
            map->chr_rom = cart->chr_rom;
            map->chr_ram = NULL;
            map->chr = map->chr_rom;
        } else {
            map->chr_rom = NULL;
            map->chr_ram = calloc(32, nes_chr_page_size);
            map->chr = map->chr_ram;
        }

        map->battery = !!(cart->flags & Cart_Flag_Battery);

        mmc1_reset(nes_map);
    }
    return (NULL == map ? -1 : 0);
}

static void mmc1_done(nes_mapper* nes_map) {
    if (NULL != nes_map->data) {
        free(((mmc1_mapper*)nes_map->data)->chr_ram);
        free(nes_map->data);
    }
}

static inline uint8_t* mmc1_prg_addr(mmc1_mapper* map,
                                     uint16_t addr) {
    int bank = (addr >> 14) & 1;
    addr &= 0x3FFFU;
    return &map->prg_bank[bank][addr];
}

static inline uint8_t* mmc1_wram_addr(mmc1_mapper* map,
                                      uint16_t addr) {
    return &(map->wram[addr & 0x1FFFU]);
}

static uint8_t mmc1_read(nes_mapper* map, uint16_t addr) {
    uint8_t val = 0;
    if (addr >= nes_mem_rom_start) {
        val = *(mmc1_prg_addr((mmc1_mapper*)map->data, addr));
    } else if (addr >= nes_mem_wram_start) {
        val = *(mmc1_wram_addr((mmc1_mapper*)map->data, addr));
    }
    return val;
}

static void mmc1_write(nes_mapper* nes_map,
                       uint16_t addr, uint8_t val) {
    mmc1_mapper* map = (mmc1_mapper*)nes_map->data;
    if (addr >= nes_mem_rom_start) {
        MAP_LOG("Write $%04x < %02x", addr, val);
        if (val & 0x80U) {
            MAP_LOG("Resetting");
            map->reg_shift = 0b10000;
            map->reg_control |= 0b01100;
            mmc1_update_prg(map);
        } else {
            // TODO: Handle consecutive-cycle writes?
            int done = (map->reg_shift & 1);
            map->reg_shift = (map->reg_shift >> 1) |
                             ((val & 1) << 4);
            if (done) {
                int reg = (addr >> 13) & 0b11;
                if (reg == 0) {
                    MAP_LOG("Control %02x", map->reg_shift);
                    map->reg_control = map->reg_shift;
                    mmc1_update_chr(map);
                    mmc1_update_vram(map);
                    mmc1_update_prg(map);
                } else if (reg == 1) {
                    MAP_LOG("CHR_0 %02x", map->reg_shift);
                    map->reg_chr_0 = map->reg_shift;
                    mmc1_update_chr(map);
                } else if (reg == 2) {
                    MAP_LOG("CHR_1 %02x", map->reg_shift);
                    map->reg_chr_1 = map->reg_shift;
                    mmc1_update_chr(map);
                } else {
                    MAP_LOG("PRG %02x", map->reg_shift);
                    map->reg_prg = map->reg_shift;
                    mmc1_update_prg(map);
                }
                map->reg_shift = 0b10000;
            }
        }

    } else if ( addr < nes_mem_rom_start &&
                addr >= nes_mem_wram_start) {
        *(mmc1_wram_addr(map, addr)) = val;
    }
}

static uint8_t* mmc1_chr_addr(nes_mapper* nes_map,
                              uint16_t addr) {
    int page = (addr >> 12) & 1;
    addr &= 0xFFFU;
    return &((mmc1_mapper*)nes_map->data)->chr_bank[page][addr];
}

static void mmc1_chr_write(nes_mapper* nes_map,
                           uint16_t addr, uint8_t val) {
    if (NULL != ((mmc1_mapper*)nes_map->data)->chr_ram) {
        *(mmc1_chr_addr(nes_map, addr)) = val;
    }
}

static uint8_t* mmc1_vram_addr(nes_mapper* nes_map,
                               uint16_t addr) {
    int page = (addr >> 10) & 3;
    int loc = addr & 0x3FFU;
//    MAP_LOG("VRAM $%04x -> %p", addr, ((mmc1_mapper*)nes_map->data)->vram_bank[page]);
    return &((mmc1_mapper*)nes_map->data)->vram_bank[page][loc];
}


static void* mmc1_sram(nes_mapper* nes_map) {
    mmc1_mapper* map = (mmc1_mapper*)nes_map->data;
    return (map->battery ? map->wram : NULL);
}

static int mmc1_sram_size(nes_mapper* nes_map) {
    mmc1_mapper* map = (mmc1_mapper*)nes_map->data;
    return (map->battery ? sizeof(map->wram) : 0);
}


/* Save State */

int mmc1_state_size(const nes_mapper* nes_map) {
    mmc1_mapper* map = (mmc1_mapper*)nes_map->data;
    return (    (void*)map +
                sizeof(*map) -
                (void*)&(map->reg_shift));
}

int mmc1_state_read(nes_mapper* nes_map, const void* data) {
    mmc1_mapper* map = (mmc1_mapper*)nes_map->data;
    int size = mmc1_state_size(nes_map);
    memcpy(&(map->reg_shift), data, size);
    return size;
}

int mmc1_state_write(const nes_mapper* nes_map, void* data) {
    mmc1_mapper* map = (mmc1_mapper*)nes_map->data;
    int size = mmc1_state_size(nes_map);
    memcpy(data, &(map->reg_shift), size);
    return size;
}


nes_mapper mapper_mmc1 = {
    .name = "MMC1",
    .init = mmc1_init,
    .reset = mmc1_reset,
    .done = mmc1_done,
    .read = mmc1_read,
    .write = mmc1_write,
    .chr_addr = mmc1_chr_addr,
    .vram_addr = mmc1_vram_addr,
    .chr_write = mmc1_chr_write,

    .sram_size = mmc1_sram_size,
    .sram = mmc1_sram,

    .state_size = mmc1_state_size,
    .state_read = mmc1_state_read,
    .state_write = mmc1_state_write,
};