#include "arm11/acl.h"
#include "arm11/cheat.h"
#include "arm11/fmt.h"
#include "util.h"
#include <stddef.h>
#include <stdlib.h>
#include <string.h>

typedef u32 instruction_t;
typedef instruction_t* CodeLocation;
#define INSTR_SIZE  sizeof(instruction_t)
#define INSTR_LEN(memsize)   ((memsize)/INSTR_SIZE)
#define CodeAtLocation(p) (*(p))

#define MAKE_ENT(hole,key)  (((key)<<16) | (hole))
#define ENT_KEY(id)         ((id)>>16)
#define ENT_HOLE(id)        ((id)&0xffff)

typedef struct {
    u32  chtId;             // acl_chtid_t
    u16  entLen;            // entArr.length
    acl_index_t *entArr;    // 下标表示hole，值表示key *改为分两段，前段key，后段overwrite value
} CurrentCheat;

static CurrentCheat setting = {0, 0, NULL};
#define CCHT_OVERWRITE(c) (&c.entArr[c.entLen])

cheat_error_t init_current_cheat( u32 id, u16 len )
{
    if( setting.entArr ) fini_current_cheat();
    if( id == 0 ) return CCHT_OK;

    setting.chtId = id;
    setting.entLen = len;
    if( len )
    {
        acl_index_t *p = (acl_index_t*)malloc( len * sizeof(acl_index_t) * 2 );
        if( p == NULL ) return CCHT_NO_MEM;
        memset(p, 0, len*sizeof(acl_index_t)*2);
        setting.entArr = p;
    }
    return CCHT_OK;
}

cheat_error_t put_current_cheat( acl_entryid_t entid )
{
    if( setting.chtId == 0 ) return CCHT_NOT_INIT;
    u16 index = ENT_HOLE(entid);
    u16 option = ENT_KEY(entid);
    if( index == 0 ) return CCHT_INVALID;

    if( index <= setting.entLen )
    {
        setting.entArr[index-1] = option;
        return CCHT_OK;
    }
    else return CCHT_INVALID;
}

cheat_error_t overwrite_current_cheat( acl_entryid_t id, u16 value )
{
    if( setting.chtId == 0 ) return CCHT_NOT_INIT;
    
    u16 index = ENT_HOLE(id);
    if( index == 0 ) return CCHT_INVALID;

    if( index <= setting.entLen )
    {
        u16 *valarr = CCHT_OVERWRITE(setting);
        valarr[index-1] = value;
        return CCHT_OK;
    }
    else return CCHT_INVALID;
}

cheat_error_t info_current_cheat( u32 *id, u32 *len )
{
    if( setting.chtId == 0 ) return CCHT_NOT_INIT;
    if( id != NULL ) *id = setting.chtId;
    if( len != NULL ) *len = setting.entLen;
    return CCHT_OK;
}

cheat_error_t get_current_cheat( u32 index, acl_entryid_t *id )
{
    if( setting.chtId == 0 ) return CCHT_NOT_INIT;
    if( index < setting.entLen )
    {
        acl_index_t val = setting.entArr[index];
        if( id != NULL ) *id = MAKE_ENT( index+1, val );
        return CCHT_OK;
    }
    else return CCHT_INVALID;
}

cheat_error_t include_current_cheat( acl_entryid_t id )
{
    if( setting.chtId == 0 ) return CCHT_NOT_INIT;

    int index = ENT_HOLE(id);
    if( index == 0 ) return CCHT_INVALID;
    else if( index < setting.entLen )
    {
        acl_index_t val = setting.entArr[index-1];
        if( val == 0 ) return CCHT_INVALID;
        else if( val == ENT_KEY(id) ) return CCHT_OK;
        else return CCHT_INVALID;
    }
    else return CCHT_INVALID;
}

cheat_error_t fini_current_cheat()
{
    if( setting.entArr ) free( setting.entArr );
    setting.entArr = NULL;
    setting.entLen = 0;
    setting.chtId = 0;
    return CCHT_OK;
}

// *******************************************
//  code for patch rom with cheat instruction
// *******************************************

#define FIT_SPACE_RESERVED  80
#define ROM_LOC             ((CodeLocation)0x20000000u)

#define GBACPU_PREFETCH  2
#define GBACPU_PREFETCH_BYTE (INSTR_SIZE*GBACPU_PREFETCH)

#define GBA_KEYCODE(k)      (0x3ff & (~(k)))
#define GBA_CART_ADDR       (0x8000000u)

#define SIZE_32M (32*1024*1024)

const instruction_t HOOKPOINT_INSTR[] = {
    0xe92d8000, // STMDB sp!, {pc}
    0xe51ff004, // LDR   pc, [pc, #-4]
    0
};

#define HP_INSTR_SIZE sizeof(HOOKPOINT_INSTR)
#define HP_INSTR_LEN INSTR_LEN(HP_INSTR_SIZE)
#define HP_WRAPPER_ADDR 2

const instruction_t IRQ_WRAPPER_INSTR[] = {
    0xe92d4000, // 0
    0xe3a0e402,
    0xe28ee701,
    0xe24ee004,
    0xe90e08ff,
    0xeb000000, // 5 - short jump
    0xe3a0e402,
    0xe28ee701,
    0xe24ee028,
    0xe89e08ff,
    0xe8bd4000, // 10
    0,          // irq
    0,
    0,
    0xe8bd8000
};

#define IW_INSTR_SIZE sizeof(IRQ_WRAPPER_INSTR)
#define IW_INSTR_LEN INSTR_LEN(IW_INSTR_SIZE)
#define IW_CALL_HANDLER 5
#define IW_ORIGNAL_IRQ 11

const instruction_t KEY_ONOFF_INSTR[] = {
    0xe3a00301, // mov     r0, #67108864   ; 0x4000000
    0xe5900130, // ldr     r0, [r0, #304]  ; 0x130
    0xe1df13bc, // ldrh    r1, [pc, #60]   ; 0x4c
    0xe1df23ba, // ldrh    r2, [pc, #58]   ; 0x4e
    0xe0000002, // and     r0, r0, r2
    0xe59f2034, // ldr     r2, [pc, #52]   ; 0x50
    0xe5d23000, // ldrb    r3, [r2]
    0xe1500001, // cmp     r0, r1
    0x0a000005, // beq     0x3c
    0xe2131010, // ands    r1, r3, #16     ; 0x10
    0x1203300f, // andne   r3, r3, #15     ; 0xf
    0x15c23000, // strneb  r3, [r2]
    0xe3530000, // cmp     r3, #0  ; 0x0
    0x1a000006, // bne     0x54
    0xe12fff1e, // bx      lr
    0xe2131010, // ands    r1, r3, #16     ; 0x10
    0x02233011, // eoreq   r3, r3, #17     ; 0x11
    0x05c23000, // streqb  r3, [r2]
    0xeafffff8, // b       0x30
    0x03ff0000, // mask-data and key-data
    0           // addr to store on/off flag
};

#define KOO_INSTR_SIZE sizeof(KEY_ONOFF_INSTR)
#define KOO_INSTR_LEN INSTR_LEN(KOO_INSTR_SIZE)
#define KOO_INSTR_KEYDATA 19
#define KOO_INSTR_MEMADDR 20

const instruction_t KEY_ENABLE_INSTR[] = {
    0xe3a00301, // mov     r0, #67108864   ; 0x4000000
    0xe5900130, // ldr     r0, [r0, #304]  ; 0x130
    0xe1df11b0, // ldrh    r1, [pc, #16]
    0xe1df20be, // ldrh    r2, [pc, #14]
    0xe0000002, // and     r0, r0, r2
    0xe1500001, // cmp     r0, r1
    0x0a000001, // beq     0x24
    0xe12fff1e, // bx      lr
    0x03ff0000  // mask-data and key-data
};

#define KEN_INSTR_SIZE sizeof(KEY_ENABLE_INSTR)
#define KEN_INSTR_LEN INSTR_LEN(KEN_INSTR_SIZE)
#define KEN_INSTR_KEYDATA 8

const instruction_t MEM_OVERWRITE_INSTR[] = {
    0xe28f0028, // add     r0, pc, #40     ; 0x28
    0xe4901004, // ldr     r1, [r0], #4
    0xe3510000, // cmp     r1, #0  ; 0x0
    0x0a000006, // beq     0x2c
    0xe4902004, // ldr     r2, [r0], #4
    0xe20230ff, // and     r3, r2, #255    ; 0xff
    0xe1a02422, // mov     r2, r2, lsr #8
    0xe4c13001, // strb    r3, [r1], #1
    0xe2522001, // subs    r2, r2, #1      ; 0x1
    0x1afffffc, // bne     0x1c
    0xeafffff5, // b       0x4
    0xe12fff1e  // bx      lr
};

#define MO_INSTR_SIZE sizeof(MEM_OVERWRITE_INSTR)
#define MO_INSTR_LEN INSTR_LEN(MO_INSTR_SIZE)

static int end_of_rom( CodeLocation addr, u32 size )
{
    CodeLocation prom = addr + INSTR_LEN( size ) - 1;
    while( addr < prom && (CodeAtLocation(prom) == 0x00000000 || CodeAtLocation(prom) == 0xffffffff) ) 
    {
        prom--;
    }
    return INSTR_SIZE * (prom + 2 - addr); // consider the first 0x00000000/0xffffffff as in-using
}

#define MASK_PC0( m ) (pc[0]&m)
#define MASK_PC1( m ) (pc[1]&m)
#define MASK_PC2( m ) (pc[2]&m)
#define MASK_PC3( m ) (pc[3]&m)
#define MASK_PC4( m ) (pc[4]&m)
#define MASK_PC5( m ) (pc[5]&m)
#define MASK_PC6( m ) (pc[6]&m)
#define MASK_PC7( m ) (pc[7]&m)
#define MASK_PC8( m ) (pc[8]&m)
#define MASK_PC9( m ) (pc[9]&m)

#define IRQ_HANDLER_POINTER (0x03007ffc)    //see gbatek `BIOS Interrupt handling'
#define CORE_HOOKPOINT_RANGE (1<<9)
#define CORE_HOOKPOINT_NEARBY (1<<8)

static int rom_search_hookpoint( CodeLocation addr, int addrlen, CodeLocation hookpoint[MAX_HOOKPOINT] )
{
    CodeLocation mark[MAX_HOOKPOINT];
    memset( mark, 0, sizeof(mark) );
    int hookpoint_idx = 0, mark_idx = 0;
    for( int i=0; i < addrlen; ++i )
    {
        CodeLocation pc = addr + i;
        if( *pc == IRQ_HANDLER_POINTER && mark_idx < MAX_HOOKPOINT ) 
            mark[mark_idx++] = pc;

        if( MASK_PC0(0XFFFF0FFF) == 0XE3A00301 &&
            MASK_PC1(0XFFF00FFF) == 0XE2800C02 &&
            MASK_PC2(0XFFF00FFF) == 0XE5D00008 &&
            MASK_PC2(0XFFFF0000) != 0XE59F0000 )
            hookpoint[hookpoint_idx++] = pc;
        
        else 
        if( MASK_PC0(0XFFFF0FFF) == 0XE3A00301 &&
            MASK_PC1(0XFFF00FFF) == 0XE2800C02 &&
            MASK_PC2(0XFFF00FFF) == 0XE5900000 &&
            MASK_PC2(0XFFFF0000) != 0XE59F0000 )
            hookpoint[hookpoint_idx++] = pc;

        else 
        if( MASK_PC0(0XFFFF0000) == 0XE92D0000 &&
            MASK_PC1(0XFFFF0FFF) == 0XE3A00301 &&
            MASK_PC2(0XFFF00FFF) == 0XE5B00200 &&
            MASK_PC3(0XFFFF0000) != 0XE59F0000 )
            hookpoint[hookpoint_idx++] = pc;

        else 
        if( MASK_PC0(0XFFFF0FFF) == 0XE3A00640 &&
            MASK_PC1(0XFFF00FFF) == 0XE5B00200 &&
            MASK_PC2(0XFFF00000) == 0XE1D00000 &&
            MASK_PC5(0XFFFF0000) != 0XE59F0000 &&
            MASK_PC6(0XFFFF0000) != 0XE59F0000 &&
            MASK_PC7(0XFFFF0000) != 0XE59F0000 )
            hookpoint[hookpoint_idx++] = pc;

        else
        if( MASK_PC0(0XFFFF0FFF) == 0XE3A00301 &&
            MASK_PC1(0XFFF00FFF) == 0XE5B00200 &&
            MASK_PC2(0XFFF00FFF) == 0XE1D000B8 )
            hookpoint[hookpoint_idx++] = pc;

        else
        if( MASK_PC0(0XFFFF0000) == 0XE59F0000 &&
            MASK_PC1(0XFFF000FF) == 0XE5900000 &&
            MASK_PC2(0XFFFF0000) == 0XE1A00000 &&
            MASK_PC1(0XFFFF0000) != 0XE59F0000 &&
            MASK_PC3(0XFFFF0000) != 0XE59F0000 )
            hookpoint[hookpoint_idx++] = pc;

        if( hookpoint_idx >= MAX_HOOKPOINT ) break;
    }
    // test for core hookpoint
    CodeLocation core_hpt = NULL;
    for( int i=0; i < hookpoint_idx; ++i )
    {
        CodeLocation p = hookpoint[i];
        if( p - addr > CORE_HOOKPOINT_RANGE ) break;

        for( int j = 0; j < mark_idx; ++j )
        {
            CodeLocation q = mark[j];
            u32 d = p<q ? q-p : p-q;
            if( d < CORE_HOOKPOINT_NEARBY )
            {
                core_hpt = p;
                break;
            }
        }
        if( core_hpt != NULL ) break;
    }
    if( core_hpt != NULL )
    {
        // log
        hookpoint[0] = core_hpt;
        hookpoint_idx = 1;
    }
    return hookpoint_idx;
}

static int cht_calc_needsize( int mode, int numirq )
{
    int total = IW_INSTR_SIZE * numirq + MO_INSTR_SIZE;

    if( mode == CHEAT_MODE_ENABYKEY )
    {
        total += KEN_INSTR_SIZE;
    }
    else if( mode == CHEAT_MODE_KEYONOFF )
    {
        total += KOO_INSTR_SIZE;
    }

    // get the cheat data length
    for( int i=0; i < setting.entLen; ++i )
    {
        u16 val = setting.entArr[i];
        if( val != 0 ) 
        {
            acl_elemlen_t len;
            acl_select_entry( MAKE_ENT(1+i, val), &len );
            total += len*INSTR_SIZE;
        }
    }

    return total + 1*INSTR_SIZE;    // the zero-end marks for cheat data
}

static CodeLocation rom_fit_newsize( CodeLocation rom, int realend, int totalsize, u32 szrom, u32 *newsize )
{
    u32 size = realend + totalsize;
    
    if( size > SIZE_32M ) 
    {
    // get the biggest space that all bytes is 0
        CodeLocation found=NULL, current=NULL;
        int szfound=0, szcurrent=0;
        int reallen = INSTR_LEN( realend );
        instruction_t freeinstr = 0; // 0/0xffffffff will be free

#define START_FREE_BLOCK(n) {\
                                current = rom+i;\
                                freeinstr = n;\
                            }
#define END_FREE_BLOCK  {\
                            if( szcurrent > szfound )\
                            {\
                                found = current;\
                                szfound = szcurrent;\
                            }\
                            current = NULL;\
                            szcurrent = 0;\
                        }

        for( int i=0; i < reallen; ++i )
        {
            if( rom[i] == 0 )
            {
                if( szcurrent == 0 ) 
                {
                    START_FREE_BLOCK(0);
                }
                else 
                {
                    if( freeinstr == 0 )// same block
                        szcurrent += INSTR_SIZE;
                    else END_FREE_BLOCK;
                } 
            }
            else if( rom[i] == 0xffffffff )
            {
                if( szcurrent == 0 ) 
                {
                    START_FREE_BLOCK(0xffffffff);
                }
                else
                {
                    if( freeinstr == 0xffffffff )
                        szcurrent += INSTR_SIZE;
                    else END_FREE_BLOCK;
                }
            }
            else if( szcurrent > 0 ) END_FREE_BLOCK;
        }

        // the zero-filled space is big enough, 3 is a guess value
        if( szfound > totalsize * 3 )
        {
            return found + FIT_SPACE_RESERVED;
        }
        // cannot find a good block
        else return NULL;
    }

    // grow cart volumn
    else if( szrom < size || nextPow2(size) != nextPow2(szrom) )
    {
        u32 sznrom = nextPow2(size);
        u32 start = realend - INSTR_SIZE;
        memset( &rom[INSTR_LEN(start)], 0, sznrom-start );
        *newsize = sznrom;
        return rom + INSTR_LEN( sznrom - ( (totalsize+0x1f) & (~0xf) ) );
    }
    
    return rom + INSTR_LEN( realend );
}

static void rom_patch_hookpoint( CodeLocation start, CodeLocation hookpoint[MAX_HOOKPOINT], int hookcnt )
{
    u32 addr = GBA_CART_ADDR + (start-ROM_LOC)*INSTR_SIZE;
    for( int i=0; i < hookcnt; ++i )
    {
        CodeLocation hook_point = hookpoint[i];
        memcpy( hook_point, HOOKPOINT_INSTR, HP_INSTR_SIZE );
        hook_point[HP_WRAPPER_ADDR] = addr + IW_INSTR_SIZE * i;
    }
}

static void rom_append_newirq( CodeLocation start, CodeLocation hookpoint[MAX_HOOKPOINT], int hookcnt )
{
    CodeLocation p = start;
    u32 offset = IW_INSTR_LEN - IW_CALL_HANDLER - GBACPU_PREFETCH;// 意思是从iw_call_handler到本函数结尾的长度，并扣除CPU偏移
    for( int i=0; i < hookcnt; ++i )
    {
        memcpy( p, IRQ_WRAPPER_INSTR, IW_INSTR_SIZE );
        memcpy( p+IW_ORIGNAL_IRQ, hookpoint[i], HP_INSTR_SIZE );
        p[IW_CALL_HANDLER] |= offset + (hookcnt-1-i) * IW_INSTR_LEN;
        p += IW_INSTR_LEN;
    }
}

static void rom_append_cheatproc( int mode, CodeLocation start, u16 bindkey, u32 storagemem )
{
    if( mode == CHEAT_MODE_KEYONOFF )
    {
        memcpy( start, KEY_ONOFF_INSTR, KOO_INSTR_SIZE );
        start[ KOO_INSTR_KEYDATA ] |= bindkey;
        start[ KOO_INSTR_MEMADDR ] = storagemem;
        start += KOO_INSTR_LEN;
    }
    else if( mode == CHEAT_MODE_ENABYKEY )
    {
        memcpy( start, KEY_ENABLE_INSTR, KEN_INSTR_SIZE );
        start[ KEN_INSTR_KEYDATA ] |= bindkey;
        start += KEN_INSTR_LEN;
    }

    memcpy( start, MEM_OVERWRITE_INSTR, MO_INSTR_SIZE );
    start += MO_INSTR_LEN;

    u16 *overwrite = CCHT_OVERWRITE(setting);
    for( int i=0; i < setting.entLen; ++i )
    {
        u16 val = setting.entArr[i];
        if( val != 0 )
        {
            acl_elemlen_t len;
            acl_select_entry( MAKE_ENT(i+1, val), &len );
            val = overwrite[i];
            for( int j=0; j < len; ++j )
            {
                if( val > 0 )
                {
                    if( (j < 4) && (j & 1) )
                    {
                        u8 *p = (u8*)start++;
                        acl_entry_get_armcode(j, (u32*)p);
                        *p = j>1 ? val >> 8 : val & 0xff;
                    }
                    else acl_entry_get_armcode(j, start++);
                }
                else acl_entry_get_armcode(j, start++);
            }
        }
    }
    *start++ = 0;
}

cheat_error_t apply_cheat( int mode, u32 szrom, hookpoint_analyzer *analyzer, u16 bindkey, u32 storagemem, u32 *outsize )
{
    // try ignore patch
    if( mode == CHEAT_MODE_DISABLED ) return CCHT_OK;
    if( setting.chtId == 0 ) return CCHT_OK;
    if( ACHTLIB_SUCCESS != acl_select_cheat_set( setting.chtId ) )
        return CCHT_NO_CHEAT;

    CodeLocation romdata = ROM_LOC;
    int realend = end_of_rom(romdata, szrom);

    // find hook point
    CodeLocation hookpoint[MAX_HOOKPOINT];
    memset( hookpoint, 0, sizeof(hookpoint) );
    int n_hookpoint = 0;
    if( analyzer!=NULL && analyzer->provider != NULL )
        n_hookpoint = analyzer->provider( analyzer->caller_data, hookpoint );
    if( n_hookpoint == 0 ) 
        n_hookpoint = rom_search_hookpoint( romdata, INSTR_LEN(realend) - 3, hookpoint ); // hook point need at least 3 instructions
    if( n_hookpoint == 0 ) return CCHT_NO_IRQ;

    // find free space to put new code
    int total_size = cht_calc_needsize( mode, n_hookpoint );
    CodeLocation page = rom_fit_newsize( romdata, realend, total_size, szrom, outsize );
    if( page == NULL ) return CCHT_NO_SPACE;

    // patching the rom
    rom_append_newirq( page, hookpoint, n_hookpoint );
    rom_patch_hookpoint( page, hookpoint, n_hookpoint );
    rom_append_cheatproc( mode, page + n_hookpoint * IW_INSTR_LEN, GBA_KEYCODE(bindkey), storagemem );
    return CCHT_OK;
}