#include <list>
#include <hash_map.h>
#include "romhack.h"

using namespace std;

list<u32> ambiguous;
typedef list<sig_t*> siglist_t;
static hash_map<u32, siglist_t> first_ops;

typedef enum {
    NOT_FOUND,
    FOUND,
    AMBIGUOUS,
} sig_return_t;

void make_first_ops(sigmap_t &sigs)
{
    sigmap_t::iterator i;

    for (i = sigs.begin(); i != sigs.end(); ++i) {
        first_ops[i->second->first_op & 0xffff0000].push_front(i->second);
    }
}

sig_return_t check_signature(u32 *code, sig_t *sig)
{
    u32 crc = 0;
    u32 op;
    int i;
    list<rel_t*>::iterator p;
    u32 *start = code;
    u8 not_sure = 0;
   
    p = sig->relocs.begin();

    for (i = 0; i < sig->num_ops; ++i) {
        op = *code++;
        if (p != sig->relocs.end() && i == (*p)->index) {
            sym_t *s = (*p)->sym;
            switch ((*p)->type) {
            case R_MIPS_HI16:
                if (s) {
                    if (s->found) {
                        if ((op & 0x0000ffff) != (s->hi + (*p)->hi_offset)) {
                            return NOT_FOUND;
                        }
                    } else {
                        not_sure = 1;
                    }
                }
                op &= 0xffff0000;
                break;
            case R_MIPS_LO16:
                if (s) {
                    if (s->found) {
                        if ((op & 0x0000ffff) != (s->lo + (*p)->lo_offset)) {
                            return NOT_FOUND;
                        }
                    } else {
                        not_sure = 1;
                    }
                }
                op &= 0xffff0000;
                break;
            case R_MIPS_26:
                if (s) {
                    if (s->found) {
                        if ((op & 0x03ffffff) != ( (((s->hi<<16)|s->lo) & 0x0ffffffc)>>2)) {
                            return NOT_FOUND;
                        }
                    } else {
                        not_sure = 1;
                    }
                }
                op &= 0xfc000000;
                break;
            default:
                fprintf(stderr, "Invalid relocation type: %d\n", (*p)->type);
                break;
            }
            ++p;
        } else {
            if (((op >> 26) & 0x3f) == 2) { /* Jump */
                op &= 0xfc000000;
            }
        }
        crc = crc32(crc, (u8*)&op, sizeof(op));
        if (i == 3 && crc != sig->partial_crc)
            return NOT_FOUND;
    }

    if (crc == sig->crc) {
        if (sig->ambiguous && not_sure)
            return AMBIGUOUS;
        code = start;
        p = sig->relocs.begin();
        for (i = 0; i < sig->num_ops; ++i) {
            op = *code++;
            if (p != sig->relocs.end() && i == (*p)->index) {
                sym_t *s = (*p)->sym;
                if (s && !s->found) {
                    switch ((*p)->type) {
                    case R_MIPS_HI16:
                        s->hi = (op-(*p)->hi_offset) & 0xffff;
                        s->found = s->lo != 0;
                        break;
                    case R_MIPS_LO16:
                        s->lo = (op-(*p)->lo_offset) & 0xffff;
                        s->found = s->hi != 0;
                        break;
                    case R_MIPS_26:
                        s->lo = ((op - (*p)->hi_offset) << 2) & 0xffff;
                        s->hi = (((op - (*p)->hi_offset) & 0x03ffffff) >> 14) | 0x8000;
                        fprintf(stderr, "%s set to 0x%04x%04x\n", s->symbol, s->hi, s->lo);
                        s->found = 1;
                        break;
                    }
                }
                ++p;
            }
        }
        return FOUND;
    } else {
        return NOT_FOUND;
    }
}

void delete_signature(sig_t *s)
{
    first_ops[s->first_op & 0xffff0000].remove(s);
}

u32 check_address(u8* rom, u32 *rom_end, u32 *addr, u8 rescan)
{
    siglist_t::iterator sl;
    hash_map<u32, siglist_t>::iterator i=first_ops.find(*addr & 0xffff0000);
    u32 ret = 1;
    if (i != first_ops.end()) {
        for (sl = (i->second).begin(); sl != (i->second).end(); sl++) {
            if (((*addr) & (*sl)->first_mask) != (*sl)->first_op)
                continue;
            if ((rom_end - addr) >= (*sl)->num_ops) {
                sig_return_t r = check_signature(addr, *sl);
                if (r == FOUND) {
                    fprintf(stderr, "%s found @ 0x%08x\n", (*sl)->symbol, (u8*)addr - rom);
                    (*sl)->offsets.push_back((u8*)addr - rom);
                    //delete_signature(*sl);
                    ret = (*sl)->num_ops;
                    break;
                } else if (r == AMBIGUOUS) {
                    if (rescan) {
                        fprintf(stderr, "Still ambiguous match @ 0x%08x [%s]\n", (u8*)addr - rom, (*sl)->symbol);
                    } else {
                        fprintf(stderr, "Found ambiguous match @ 0x%08x [%s]\n", (u8*)addr - rom, (*sl)->symbol);
                        ambiguous.push_back((u32)((u8*)addr-rom));
                    }
                    ret = (*sl)->num_ops;

                    break;
                }
            }
        }
    }

    return ret;
}

void scan_rom(u8* rom, u32 rom_len, sigmap_t &sigs)
{
    u32 *p = (u32 *) rom;
    u32 *rom_end = (u32 *)(rom + rom_len);

    make_first_ops(sigs);

    fprintf(stderr, "Scanning ROM for functions\n");
    while (p < rom_end) {
        p += check_address(rom, rom_end, p, 0);
    }
    fprintf(stderr, "done.\n\n");

    fprintf(stderr, "Re-scanning ambiguous functions\n");
    for (list<u32>::const_iterator i=ambiguous.begin(); i != ambiguous.end(); ++i) {
        check_address(rom, rom_end, (u32 *)(rom+*i), 1);
    }
    fprintf(stderr, "done.\n\n");
}