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

#define MAPPER_NAME "MMC1"
#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;

    int chr_bank_offset[2];
    int chr_ram_lim;

    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];

    uint8_t chr_ram[];
} 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) {
    int n_banks = map->chr_rom_banks;
    if (n_banks <= 0) n_banks = 32;
    int banks[2] = {0};
    if (!(map->reg_control & 0b10000)) {
        banks[0] = (map->reg_chr_0 & 0b11110) % n_banks;
        banks[1] = banks[0] + 1;
    } else {
        banks[0] = map->reg_chr_0 % n_banks;
        banks[1] = map->reg_chr_1 % n_banks;
    }

    MAP_LOG("CHR: %d + %d", banks[0], banks[1]);

    map->chr_bank_offset[0] = banks[0] * nes_chr_page_size;
    map->chr_bank_offset[1] = banks[1] * nes_chr_page_size;

    map->chr_bank[0] = &map->chr[map->chr_bank_offset[0]];
    map->chr_bank[1] = &map->chr[map->chr_bank_offset[1]];
}

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) % map->prg_rom_banks;
    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(void* data) {
    mmc1_mapper* map = (mmc1_mapper*)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 void* mmc1_init(nes_mapper* nes_map, nes_cart* cart) {
    int chr_ram_size = (    cart->chr_rom_banks > 0 ?
                            0 : 128 * 1024);
    mmc1_mapper* map = calloc(1, sizeof(mmc1_mapper) +
                                 chr_ram_size);
    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;
        if (map->chr_rom_banks > 0) {
            map->chr_rom = cart->chr_rom;
            map->chr = map->chr_rom;
        } else {
            map->chr_rom = NULL;
            map->chr = map->chr_ram;
        }
        map->chr_ram_lim = 0;

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

        mmc1_reset(map);
    }
    return map;
}

static void mmc1_done(void* data) {
    free(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(void* data, uint16_t addr) {
    uint8_t val = 0;
    if (addr >= nes_mem_rom_start) {
        val = *(mmc1_prg_addr((mmc1_mapper*)data, addr));
    } else if (addr >= nes_mem_wram_start) {
        val = *(mmc1_wram_addr((mmc1_mapper*)data, addr));
    }
    return val;
}

static void mmc1_write(void* data, uint16_t addr, uint8_t val) {
    mmc1_mapper* map = (mmc1_mapper*)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(void* data, uint16_t addr) {
    int page = (addr >> 12) & 1;
    addr &= 0xFFFU;
    return &((mmc1_mapper*)data)->chr_bank[page][addr];
}

static void mmc1_chr_write(void* data,
                           uint16_t addr, uint8_t val) {
    mmc1_mapper* map = (mmc1_mapper*)data;
    if (NULL == map->chr_rom) {
        int pos = map->chr_bank_offset[(addr >> 12) & 1] +
                  (addr & 0xFFFU);
        if (pos >= map->chr_ram_lim) map->chr_ram_lim = pos + 1;
        *(mmc1_chr_addr(data, addr)) = val;
    }
}

static uint8_t* mmc1_vram_addr(void* data, 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*)data)->vram_bank[page][loc];
}


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

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


/* Save State */

static inline int mmc1_chr_ram_size(const mmc1_mapper* map) {
    return (map->chr_rom_banks <= 0 ? (128 * 1024) : 0);
}

static int mmc1_state_size(const void* _map) {
    mmc1_mapper* map = (mmc1_mapper*)_map;

    return (    map->chr_ram_lim +
                (void*)map +
                sizeof(*map) -
                (void*)&(map->reg_shift));
}

static int mmc1_state_read(void* _map, const void* data,
                           int data_len) {
    mmc1_mapper* map = (mmc1_mapper*)_map;

    int base_size = mmc1_state_size(map) - map->chr_ram_lim;
    int size = base_size + mmc1_chr_ram_size(map);
    if (size > data_len) size = data_len;
    if (data_len - base_size > map->chr_ram_lim) {
        map->chr_ram_lim = data_len - base_size;
    }
    memcpy(&(map->reg_shift), data, size);

    mmc1_update_prg(map);
    mmc1_update_chr(map);
    mmc1_update_vram(map);

    return size;
}

static int mmc1_state_write(const void* _map, void* data,
                            int data_len) {
    mmc1_mapper* map = (mmc1_mapper*)_map;
    int size = mmc1_state_size(_map);
    if (size > data_len) size = data_len;
    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,
};