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

#include "map.h"


typedef enum {
    mmc3_Flag_Horizontal   = 0b00000001,
    mmc3_Flag_IRQ_Enabled  = 0b00000010,
    mmc3_Flag_IRQ_Reload   = 0b00000100,
    mmc3_Flag_CHR_RAM      = 0b00001000,
    mmc3_Flag_Battery      = 0b00100000,
    mmc3_Flag_WRAM_Protect = 0b01000000,
    mmc3_Flag_WRAM_Enabled = 0b10000000,
} mmc3_Flag;

typedef enum {
    mmc3_Bank_Select_Reg    = 0b00000111,
    mmc3_Bank_Select_PRG    = 0b01000000,
    mmc3_Bank_Select_CHR    = 0b10000000,
} mmc3_Bank_Select;

typedef struct {
    uint8_t* prg_rom;
    int prg_rom_banks; // 16 KB / 8 KB = 2
    uint8_t* chr_rom;
    int chr_rom_banks; // 4 KB / 1 KB = 4

    nes_mapper* mapper;

    uint8_t* prg_bank[4]; // 8 KB
    uint8_t* chr_bank[8]; // 1 KB
    uint8_t* vram_bank[4];

    int chr_bank_offset[8];
    int chr_ram_lim;

    uint8_t r[8];
    uint8_t flags;
    uint8_t bank_select;
    uint8_t irq_count;
    uint8_t irq_latch;

    uint8_t wram[nes_mem_wram_size];
    uint8_t vram[2][nes_vram_page_size];
    uint8_t chr_ram[];
} mmc3_mapper;

static inline uint8_t* mmc3_prg_bank(mmc3_mapper* map, int bank) {
    return &map->prg_rom[(bank % map->prg_rom_banks) << 13];
}

static inline void mmc3_map_prg(mmc3_mapper* map,
                                int reg, int bank) {
    MAP_LOG("PRG ROM: 8k $%04x <- bank %d", reg << 13, bank);
    map->prg_bank[reg] = mmc3_prg_bank(map, bank);
}

static inline void mmc3_update_prg(mmc3_mapper* map,
                                   int reg, uint8_t bank) {
    if (reg == 7) {
        mmc3_map_prg(map, 1, bank);
    } else {
        if (!(map->bank_select & mmc3_Bank_Select_PRG)) {
            mmc3_map_prg(map, 0, bank);
        } else {
            mmc3_map_prg(map, 2, bank);
        }
    }
}

static inline uint8_t* mmc3_chr_bank(mmc3_mapper* map, int bank) {
    return &map->chr_rom[(bank % map->chr_rom_banks) << 10];
}

static inline void mmc3_map_2k_chr(mmc3_mapper* map,
                                   int reg, int bank) {
    bank &= ~1;
    if (bank >= map->chr_rom_banks) {
        MAP_LOG("CHR ROM OOB: %d > %d", bank, map->chr_rom_banks);
        bank = bank % map->chr_rom_banks;
    }
    MAP_LOG("CHR ROM: 2k $%04x <- bank %d + %d", reg << 10, bank, bank | 1);
    map->chr_bank[reg + 0] = mmc3_chr_bank(map, bank);
    map->chr_bank[reg + 1] = mmc3_chr_bank(map, bank | 1);
    map->chr_bank_offset[reg + 0] = map->chr_bank[reg + 0] - map->chr_rom;
    map->chr_bank_offset[reg + 1] = map->chr_bank[reg + 1] - map->chr_rom;
}

static inline void mmc3_map_1k_chr(mmc3_mapper* map,
                                   int reg, int bank) {
    if (bank >= map->chr_rom_banks) {
        MAP_LOG("CHR ROM OOB: %d > %d", bank, map->chr_rom_banks);
        bank = bank % map->chr_rom_banks;
    }
    MAP_LOG("CHR ROM: 1k $%04x <- bank %d", reg << 10, bank);
    map->chr_bank[reg] = mmc3_chr_bank(map, bank);
    map->chr_bank_offset[reg] = map->chr_bank[reg] - map->chr_rom;
}

static inline void mmc3_update_chr(mmc3_mapper* map,
                                   int reg, uint8_t bank) {
    if (!(map->bank_select & mmc3_Bank_Select_CHR)) {
        if (1 >= reg) {
            mmc3_map_2k_chr(map, reg * 2, bank);
        } else {
            mmc3_map_1k_chr(map, reg + 2, bank);
        }

    } else {
        if (1 >= reg) {
            mmc3_map_2k_chr(map, 4 + (reg * 2), bank);
        } else {
            mmc3_map_1k_chr(map, reg - 2, bank);
        }
    }
}

static inline void mmc3_update_rom_mode(mmc3_mapper* map, int val) {
    uint8_t delta = (map->bank_select ^ val);

    map->bank_select = val;

    if (delta & mmc3_Bank_Select_PRG) {
        mmc3_map_prg(map, 1, map->r[7]);
        if (!(val & mmc3_Bank_Select_PRG)) {
            mmc3_map_prg(map, 0, map->r[6]);
            mmc3_map_prg(map, 2, map->prg_rom_banks - 2);
        } else {
            mmc3_map_prg(map, 0, map->prg_rom_banks - 2);
            mmc3_map_prg(map, 2, map->r[6]);
        }
    }

    if (delta & mmc3_Bank_Select_CHR) {
        mmc3_update_chr(map, 0, map->r[0]);
        mmc3_update_chr(map, 1, map->r[1]);
        mmc3_update_chr(map, 2, map->r[2]);
        mmc3_update_chr(map, 3, map->r[3]);
        mmc3_update_chr(map, 4, map->r[4]);
        mmc3_update_chr(map, 5, map->r[5]);
    }
}

static inline void mmc3_update_vram(mmc3_mapper* map) {
    if (!(map->flags & mmc3_Flag_Horizontal)) {
        // Vertical mirroring
        MAP_LOG("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 {
        // Horizontal mirroring
        MAP_LOG("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 mmc3_reset(void* data) {
    mmc3_mapper* map = (mmc3_mapper*)data;
    map->irq_count   = 0;
    map->irq_latch   = 0;
    mmc3_update_rom_mode(map, 0);
    mmc3_map_prg(map, 3, map->prg_rom_banks - 1);
    mmc3_update_vram(map);
}

static void* mmc3_init(nes_mapper* nes_map, nes_cart* cart) {
    int chr_ram_size = (    cart->chr_rom_banks <= 0 ?
                            (256 * 1024) : 0);
    mmc3_mapper* map = calloc(1, sizeof(mmc3_mapper) +
                                 chr_ram_size);
    if (NULL != map) {
        map->mapper = nes_map;
        map->flags = (cart->flags & Cart_Flag_Horizontal) ?
                     mmc3_Flag_Horizontal : 0;
        map->prg_rom       = cart->prg_rom;
        map->prg_rom_banks = cart->prg_rom_banks * 2;
        if (cart->chr_rom_banks <= 0) {
            map->chr_rom = map->chr_ram;
            map->chr_rom_banks = 256;
            map->flags |= mmc3_Flag_CHR_RAM;
        } else {
            map->chr_rom       = cart->chr_rom;
            map->chr_rom_banks = cart->chr_rom_banks * 4;
        }

        if (cart->flags & Cart_Flag_Battery) {
            map->flags |= mmc3_Flag_Battery;
        }

        map->bank_select = mmc3_Bank_Select_PRG |
                           mmc3_Bank_Select_CHR;

        mmc3_reset(map);
    }
    return map;
}

static void mmc3_done(void* data) {
    free(data);
}

static inline uint8_t* mmc3_prg_addr(mmc3_mapper* map,
                                     uint16_t addr) {
    return &(map->prg_bank[(addr >> 13) & 3][addr & 0x1FFFU]);
}

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

static uint8_t mmc3_read(void* data, uint16_t addr) {
    uint8_t* ptr = NULL;
    mmc3_mapper* map = (mmc3_mapper*)data;
    if (addr >= nes_mem_rom_start) {
        ptr = mmc3_prg_addr(map, addr);
    } else if ( addr >= nes_mem_wram_start &&
                (map->flags & mmc3_Flag_WRAM_Enabled)) {
        ptr = mmc3_wram_addr(map, addr);
//        MAP_LOG("WRAM: $%04x > %02x", addr, *ptr);
    }

    uint8_t val = (NULL == ptr ? 0 : *ptr);

//    MAP_LOG("$%04x -> %04lx > %02x", addr, ptr - map->prg_rom, val);

    return val;
}

static void mmc3_write(void* data,
                       uint16_t addr, uint8_t val) {
    mmc3_mapper* map = (mmc3_mapper*)data;

    if (addr >= nes_mem_rom_start) MAP_LOG("$%04x < %02x", addr, val);

    if (addr < nes_mem_wram_start) {
        // Nothing prior to WRAM

    } else if (addr < nes_mem_rom_start) {
        if (    (map->flags & mmc3_Flag_WRAM_Enabled) &&
                !(map->flags & mmc3_Flag_WRAM_Protect)) {
//            MAP_LOG("WRAM: $%04x < %02x", addr, val);
            *(mmc3_wram_addr(map, addr)) = val;
        }

    } else if (addr < 0xA000U) {
        if (addr & 1) {
            // Bank data
            int reg = (map->bank_select & mmc3_Bank_Select_Reg);
            if (reg >= 6) {
                mmc3_update_prg(map, reg, val);
            } else {
                mmc3_update_chr(map, reg, val);
            }
            map->r[reg] = val;
        } else {
            // Bank select
            mmc3_update_rom_mode(map, val);
        }

    } else if (addr < 0xC000U) {
        if (addr & 1) {
            MAP_LOG("WRAM %s, %s", (val & mmc3_Flag_WRAM_Enabled) ? "enabled" : "disabled", (val & mmc3_Flag_WRAM_Protect) ? "protected" : "writable");
            // WRAM protection
            map->flags &= ~(mmc3_Flag_WRAM_Enabled |
                            mmc3_Flag_WRAM_Protect);
            map->flags |= (val & (mmc3_Flag_WRAM_Enabled |
                                  mmc3_Flag_WRAM_Protect));
        } else {
            // Mirroring
            map->flags &= ~mmc3_Flag_Horizontal;
            map->flags |= (val & mmc3_Flag_Horizontal);
            mmc3_update_vram(map);
        }

    } else if (addr < 0xE000U) {
        if (addr & 1) {
            MAP_LOG("IRQ Reload");
            map->flags |= mmc3_Flag_IRQ_Reload;
//            map->irq_count = 0;
        } else {
            MAP_LOG("IRQ Latch: %d", val);
            map->irq_latch = val;
        }

    } else {
        MAP_LOG("IRQ %s", (addr & 1) ? "Enable" : "Disable");
        if (addr & 1) {
            map->flags |= mmc3_Flag_IRQ_Enabled;
        } else {
            map->flags &= ~mmc3_Flag_IRQ_Enabled;
            nes_map_trigger_irq(map->mapper, 0);
        }
    }
}

static void mmc3_scanline(void* data) {
    mmc3_mapper* map = (mmc3_mapper*)data;

    if (    map->irq_count <= 0 ||
            (map->flags & mmc3_Flag_IRQ_Reload)) {
        map->irq_count = map->irq_latch;
        map->flags &= ~mmc3_Flag_IRQ_Reload;
    } else {
        map->irq_count--;
    }

    if (    map->irq_count <= 0 &&
            (map->flags & mmc3_Flag_IRQ_Enabled)) {
        MAP_LOG("IRQ Trigger");
        nes_map_trigger_irq(map->mapper, 1);
        map->irq_count = 0;
    }

}

static uint8_t* mmc3_chr_addr(void* data,
                              uint16_t addr) {
    mmc3_mapper* map = (mmc3_mapper*)data;
    return &map->chr_bank[(addr >> 10) & 7][addr & 0x3FFU];
}

static uint8_t* mmc3_vram_addr(void* data,
                               uint16_t addr) {
    mmc3_mapper* map = (mmc3_mapper*)data;
    return &map->vram_bank[(addr >> 10) & 3][addr & 0x3FFU];
}

static void mmc3_chr_write(void* data,
                           uint16_t addr, uint8_t val) {
    mmc3_mapper* map = (mmc3_mapper*)data;
    if (map->flags & mmc3_Flag_CHR_RAM) {
        uint8_t* ptr = mmc3_chr_addr(data, addr);
        int pos = (ptr - map->chr_ram);
        if (pos >= map->chr_ram_lim) map->chr_ram_lim = pos + 1;
        *ptr = val;
    }
//    MAP_LOG("CHR ROM Write: $%04x < %02x\n", addr, val);
}


static void* mmc3_sram(void* data) {
    mmc3_mapper* map = (mmc3_mapper*)data;
    return (    (map->flags & mmc3_Flag_Battery) ?
                map->wram : NULL);
}

static int mmc3_sram_size(void* data) {
    mmc3_mapper* map = (mmc3_mapper*)data;
    return (    (map->flags & mmc3_Flag_Battery) ?
                sizeof(map->wram) : 0);
}


/* Save State */

static inline int mmc3_chr_ram_size(const mmc3_mapper* map) {
    return (    (map->flags & mmc3_Flag_CHR_RAM) ?
                (256 * 1024) : 0);
}

static int mmc3_state_size(const void* data) {
    const mmc3_mapper* map = (mmc3_mapper*)data;
    return (    (map->wram - map->r) +
                sizeof(map->wram) +
                sizeof(map->vram) +
                map->chr_ram_lim);
}

static int mmc3_state_read(void* _map, const void* data,
                           int data_len) {
    mmc3_mapper* map = (mmc3_mapper*)_map;

    int base_size = mmc3_state_size(map) - map->chr_ram_lim;
    int size = base_size + mmc3_chr_ram_size(map);
    if (size > data_len) size = data_len;
    map->chr_ram_lim = data_len - base_size;
    memcpy(map->r, data, size);

    uint8_t new_bank_select = map->bank_select;
    map->bank_select = ~new_bank_select;

    mmc3_update_rom_mode(map, new_bank_select);
    mmc3_update_vram(map);

    return size;
}

static int mmc3_state_write(const void* _map, void* data,
                            int data_len) {
    mmc3_mapper* map = (mmc3_mapper*)_map;
    int size = mmc3_state_size(_map);
    if (size > data_len) size = data_len;
    memcpy(data, map->r, size);
    return size;
}


nes_mapper mapper_mmc3 = {
    .name = "MMC3",

    .init = mmc3_init,
    .reset = mmc3_reset,
    .done = mmc3_done,
    .read = mmc3_read,
    .write = mmc3_write,
    .chr_addr = mmc3_chr_addr,
    .vram_addr = mmc3_vram_addr,
    .chr_write = mmc3_chr_write,
    .scanline = mmc3_scanline,

    .sram_size = mmc3_sram_size,
    .sram = mmc3_sram,

    .state_size = mmc3_state_size,
    .state_read = mmc3_state_read,
    .state_write = mmc3_state_write,
};