#include <stdint.h>
#include <stdio.h>
#include <fcntl.h>
#include <sys/types.h>
#include <unistd.h>
#include <string.h>
#include <stdlib.h>
#include <assert.h>

#define MAX(x, y) (x) > (y) ? (x) : (y)

typedef struct irw_file
{
    FILE *f;
    uint16_t crc;
    uint32_t offset;
} IRW_FILE;

int print_hex_offset(const void *block, int count, int offset, uint32_t start);
int print_hex(const void *block, int count, uint32_t start);

static void update_crc16(uint16_t *crc, uint8_t byte)
{
    // CRC-16-CCITT, Initialize to 0xFFFF, No zero padding
    for (int i = 7; i >= 0; i--)
    {
        uint16_t xor_value = ((*crc >> 15) ^ ((byte >> i) & 1)) ? 0x1021 : 0;
        *crc = (*crc << 1) ^ xor_value;
    }
}

struct irw_file *irw_open(const char *filename, const char *mode)
{
    struct irw_file *f = malloc(sizeof(*f));
    memset(f, 0, sizeof(*f));
    f->f = fopen(filename, mode);
    return f;
}

int irw_readb(struct irw_file *f)
{
    int val = fgetc(f->f);
    if (val == EOF)
        return EOF;
    update_crc16(&f->crc, val);
    return val;
}

int irw_writeb(struct irw_file *f, int c) {
    update_crc16(&f->crc, c);
    return fputc(c, f->f);
}

struct Ice40Bitstream
{
    uint32_t offset;
    uint32_t current_bank;
    uint32_t current_width;
    uint32_t current_height;
    uint32_t current_offset;

    uint32_t cram_width;
    uint32_t cram_height;

    uint32_t bram_width;
    uint32_t bram_height;

    uint16_t crc_value;

    uint8_t warmboot;
    uint8_t nosleep;

    uint8_t frequency_range;
};


int ice40_parse(struct irw_file *f)
{
    uint32_t preamble = 0;
    uint8_t wakeup = 0;
    struct Ice40Bitstream bs;

    memset(&bs, 0, sizeof(bs));

    while (1)
    {
        int b = irw_readb(f);
        if (b == EOF)
            break;

        preamble = (preamble << 8) | b;
        if (preamble == 0x7eaa997e)
        {
            // printf("found preamble at %d\n", bs.offset);
            break;
        }
    }

    while (!wakeup)
    {
        int b = irw_readb(f);
        if (b == EOF)
        {
            // printf("reached end of file\n");
            break;
        }

        uint8_t cmd = b >> 4;
        uint8_t payload_len = b & 0xf;
        uint32_t payload = 0;
        uint8_t last0, last1;
        int i;
        for (i = 0; i < payload_len; i++)
        {
            payload = (payload << 8) | (irw_readb(f) & 0xff);
        }

        // printf("cmd: %02x  payload: %02x payload_len: %d\n", cmd, payload, payload_len);
        switch (cmd)
        {
        case 0:
            switch (payload)
            {
            case 1:
                printf("CRAM data (bank %d): %d x %d @ 0x%08x; %d bits = %d bytes\n",
                       bs.current_bank,
                       bs.current_width,
                       bs.current_height,
                       bs.current_offset,
                       bs.current_width * bs.current_height,
                       (bs.current_width * bs.current_height) / 8);
                bs.cram_width = MAX(bs.cram_width, bs.current_width);
                bs.cram_height = MAX(bs.cram_height, bs.current_height);
                for (i = 0; i < ((bs.current_width * bs.current_height) / 8); i++)
                {
                    (void)irw_readb(f);
                }
                last0 = irw_readb(f);
                last1 = irw_readb(f);
                if (last0 || last1)
                {
                    printf("expected 0x0000 after CRAM data, got %02x %02x\n", last0, last1);
                }
                break;
            case 3:
            {
                printf("BRAM data (bank %d): %d x %d @ 0x%08x; %d bits = %d bytes\n",
                       bs.current_bank,
                       bs.current_width,
                       bs.current_height,
                       bs.current_offset,
                       bs.current_width * bs.current_height,
                       (bs.current_width * bs.current_height) / 8);
                bs.bram_width = MAX(bs.bram_width, bs.current_width);
                bs.bram_height = MAX(bs.bram_height, bs.current_height);
                uint8_t try_1[(bs.current_height * bs.current_width) / 8];
                uint8_t try_2[(bs.current_height * bs.current_width) / 8];
                uint8_t try_3[(bs.current_height * bs.current_width) / 8];
                uint8_t try_4[(bs.current_height * bs.current_width) / 8];
                memset(try_1, 0, sizeof(try_1));
                memset(try_2, 0, sizeof(try_2));
                memset(try_3, 0, sizeof(try_3));
                memset(try_4, 0, sizeof(try_4));
                for (i = 0; i < ((bs.current_width * bs.current_height) / 8); i++)
                {
                    uint8_t byte = irw_readb(f);
                    // int j;
                    // for (j = 0; j < 8; j++)
                    // {
                    //     int x = (i * 8 + j) % bs.current_width;
                    //     int y = (i * 8 + j) / bs.current_width;// + bs.current_offset;
                    //     printf("x: %d  y: %d\n", x, y);
                    //     try_2[(x * bs.current_width + y) / 8] |= (1 << j);
                    // }
                    try_1[i] = byte;
                    // try_3[i] = reverse_u8(byte);
                    // try_4[i] = byte;
                }
                // for (i = 0; i < sizeof(try_4); i += 2)
                // {
                //     uint8_t t = try_4[i];
                //     try_4[i] = try_4[i + 1];
                //     try_4[i + 1] = t;
                // }
                printf("try 1:\n");
                print_hex(try_1, sizeof(try_1), (!!bs.current_offset) * (bs.current_width * bs.current_height) / 8);
                // printf("try 3:\n");
                // print_hex(try_3, sizeof(try_3), 0);
                // printf("try 4:\n");
                // print_hex(try_4, sizeof(try_4), 0);
                last0 = irw_readb(f);
                last1 = irw_readb(f);
                if (last0 || last1)
                {
                    printf("expected 0x0000 after BRAM data, got %02x %02x\n", last0, last1);
                }
                break;
            }
            case 5:
                // printf("resetting crc\n");
                f->crc = 0xffff;
                break;

            case 6:
                // printf("wakeup\n");
                wakeup = 1;
                break;

            default:
                printf("unrecognized command 0x%02x 0x%02x\n", cmd, payload);
                break;
            }
            break;
        case 1:
            bs.current_bank = payload;
            // printf("setting bank number to %d\n", bs.current_bank);
            break;
        case 2:
            printf("crc check (%04x == %04x)\n", f->crc, payload);
            break;
        case 5:
            switch (payload)
            {
            case 0:
                bs.frequency_range = 0;
                break;
            case 1:
                bs.frequency_range = 1;
                break;
            case 2:
                bs.frequency_range = 2;
                break;
            default:
                printf("unknown frequency range payload: %02x\n", payload);
                break;
            }
            break;
        case 6:
            bs.current_width = payload + 1;
            // printf("setting bank width to %d\n", bs.current_width);
            break;
        case 7:
            bs.current_height = payload;
            // printf("setting bank height to %d\n", bs.current_height);
            break;
        case 8:
            bs.current_offset = payload;
            // printf("setting bank offset to %d\n", bs.current_offset);
            break;
        case 9:
            switch (payload)
            {
            case 0:
                bs.warmboot = 0;
                bs.nosleep = 0;
                break;
            case 1:
                bs.warmboot = 0;
                bs.nosleep = 1;
                break;
            case 32:
                bs.warmboot = 1;
                bs.nosleep = 0;
                break;
            case 33:
                bs.warmboot = 1;
                bs.nosleep = 1;
                break;
            default:
                printf("unrecognized feature flags: %02x\n", payload);
                break;
            }
            break;
        default:
            printf("unrecognized command: %02x\n", cmd);
            break;
        }
    }

    // printf("read %d bytes\n", bs.offset);
    return 0;
}

struct mapping {
    uint32_t count;
    uint32_t *mappings;
};

static void add_mapping(struct mapping *mapping, uint32_t src, uint32_t dest) {
    mapping->mappings = realloc(mapping->mappings, (mapping->count+1) * 2 * sizeof(uint32_t));
    mapping->mappings[mapping->count * 2] = src;
    mapping->mappings[mapping->count * 2 + 1] = dest;
    printf("creating mapping %08x -> %08x\n", src, dest);
    mapping->count++;
}

static uint32_t get_mapping(const struct mapping *mapping, uint32_t val) {
    int i;
    for (i = 0; i < mapping->count; i++) {
        if (mapping->mappings[i * 2] == val) {
            printf("found mapping for %08x: %08x\n", val, mapping->mappings[i * 2 + 1]);
            return mapping->mappings[i * 2 + 1];
        }
    }
    printf("couldn't find mapping for %08x\n", val);
    return -1;
}

uint8_t get_bit(uint32_t *field, uint32_t offset)
{
    // printf("offset&31: %d\n", offset & 31);
    // printf("offset/sizeof(*field): %d\n", offset >> 5);
    assert(offset < 65536);
    return !!(field[offset >> 5] & (1 << (offset & 31)));
}

void set_bit(uint32_t *field, uint32_t offset)
{
    assert(offset < 65536);
    field[offset >> 5] |= (1 << (offset & 31));
}

void clear_bit(uint32_t *field, uint32_t offset)
{
    assert(offset < 65536);
    field[offset >> 5] &= ~(1 << (offset & 31));
}

static uint32_t get_bit_offset(int x, int total_bits) {
    // return (8192 * (x & 7)) + (x >> 3);
    int bitshift = ffs(total_bits)-1;
    return ((x * 8192) % total_bits) + ((x*8192) >> bitshift);
}

uint32_t xorshift32(uint32_t x)
{
	/* Algorithm "xor" from p. 4 of Marsaglia, "Xorshift RNGs" */
	x = x ^ (x << 13);
	x = x ^ (x >> 17);
	x = x ^ (x << 5);
	return x;
}

uint32_t get_rand(uint32_t x) {
    uint32_t out = 0;
    int i;
    for (i = 0; i < 32; i++) {
        x = xorshift32(x);
        if ((x & 1) == 1)
            out = out | (1 << i);
    }
    return out;
}

static uint32_t fill_rand(uint32_t *bfr, int count) {
    int i;
    uint32_t last = 1;
    for (i = 0; i < count / 4; i++) {
        last = get_rand(last);
        bfr[i] = last;
    }
    return i;
}

uint32_t swap_u32(uint32_t word) {
    return (((word >> 24) & 0x000000ff)
             | ((word >> 8) & 0x0000ff00)
             | ((word << 8) & 0x00ff0000)
             | ((word << 24) & 0xff000000));
}

// 1. Read ROM file into `input`
// 2. For each pair of bytes of input, 
// uint16_t mappings[65536];
int ice40_patch(struct irw_file *f, struct irw_file *rom, struct irw_file *o)
{
    uint32_t preamble = 0;
    uint8_t wakeup = 0;
    struct Ice40Bitstream bs;
    uint32_t input_rom[2048];
    uint32_t input_rand[2048];
    uint32_t output_rand[2048];
    uint32_t output_rom[2048];
    uint8_t *i8 = (uint8_t *)input_rom;
    uint16_t *ora16 = (uint16_t *)output_rand;
    uint16_t *oro16 = (uint16_t *)output_rom;
    unsigned int ora_ptr;
    unsigned int ora_matches;
    int b;
    // struct mapping mapping;

    // memset(&mapping, 0, sizeof(mapping));
    memset(&bs, 0, sizeof(bs));

    int input_ptr;

    // Read the ROM into a source buffer
    memset(input_rom, 0, sizeof(input_rom));
    input_ptr = 0;
    while ((b = irw_readb(rom)) != EOF)
        i8[input_ptr++] = b;
    printf("read %d bytes from rom\n", input_ptr);

    // Generate our reference pattern
    memset(input_rand, 0, sizeof(input_rand));
    fill_rand(input_rand, sizeof(input_rand));

    for (input_ptr = 0; input_ptr < sizeof(input_rom)/4; input_ptr++) {
        // input_rand[input_ptr] = swap_u32(input_rand[input_ptr]);
        // input_rom[input_ptr] = swap_u32(input_rom[input_ptr]);
    }

    // Spray the reference pattern and ROM like they would exist in the FPGA
    for (input_ptr = 0; input_ptr < sizeof(input_rom) * 8; input_ptr++) {
        int bit;
        bit = get_bit(input_rand, get_bit_offset(input_ptr, sizeof(input_rand)*8));
        if (bit)
            set_bit(output_rand, input_ptr);
        else
            clear_bit(output_rand, input_ptr);

        bit = get_bit(input_rom, get_bit_offset(input_ptr, sizeof(input_rom)*8));
        if (bit)
            set_bit(output_rom, input_ptr);
        else
            clear_bit(output_rom, input_ptr);
    }

    // // Finally, build a mapping from the random data to the replacement rom data.
    // uint32_t collisions = 0;
    // for (input_ptr = 0; input_ptr < 4096; input_ptr += 1) {
    //     // uint32_t ra1 = (output_rand[input_ptr + 0] & 0xffff0000)
    //     //                | ((output_rand[input_ptr + 1] >> 16) & 0x0000ffff);
    //     // uint32_t ra2 = (output_rand[input_ptr + 1] & 0xffff0000)
    //     //                | ((output_rand[input_ptr + 0] >> 16) & 0x0000ffff);
    //     // add_mapping(&mapping, ra1, output_rom[input_ptr + 0]);
    //     // add_mapping(&mapping, ra2, output_rom[input_ptr + 1]);
    //     if (mappings[ora16[input_ptr]]) {
    //         collisions++;
    //         printf("mapping position %d not empty\n", input_ptr);
    //     }
    //     mappings[ora16[input_ptr]] = oro16[input_ptr];
    // }
    // if (collisions)
    //     printf("had %d collisions\n", collisions);

    while (1)
    {
        b = irw_readb(f);
        if (b == EOF)
            break;
        irw_writeb(o, b);

        preamble = (preamble << 8) | b;
        if (preamble == 0x7eaa997e)
        {
            // printf("found preamble at %d\n", bs.offset);
            break;
        }
    }

    while (!wakeup)
    {
        int b = irw_readb(f);
        if (b == EOF)
        {
            // printf("reached end of file\n");
            break;
        }
        irw_writeb(o, b);

        uint8_t cmd = b >> 4;
        uint8_t payload_len = b & 0xf;
        uint32_t payload = 0;
        uint8_t last0, last1;
        int i;
        for (i = 0; i < payload_len; i++)
        {
            b = irw_readb(f);
            payload = (payload << 8) | (b & 0xff);

            // Don't write the CRC16 out, since we'll do that later on.
            if (cmd != 2)
                irw_writeb(o, b);
        }

        switch (cmd)
        {
        case 0:
            switch (payload)
            {
            case 1:
                printf("CRAM data (bank %d): %d x %d @ 0x%08x; %d bits = %d bytes\n",
                       bs.current_bank,
                       bs.current_width,
                       bs.current_height,
                       bs.current_offset,
                       bs.current_width * bs.current_height,
                       (bs.current_width * bs.current_height) / 8);
                bs.cram_width = MAX(bs.cram_width, bs.current_width);
                bs.cram_height = MAX(bs.cram_height, bs.current_height);
                for (i = 0; i < ((bs.current_width * bs.current_height) / 8); i++)
                {
                    irw_writeb(o, irw_readb(f));
                }
                last0 = irw_readb(f);
                last1 = irw_readb(f);
                if (last0 || last1)
                {
                    printf("expected 0x0000 after CRAM data, got %02x %02x\n", last0, last1);
                }
                irw_writeb(o, last0);
                irw_writeb(o, last1);
                break;
            case 3:
                printf("BRAM data (bank %d): %d x %d @ 0x%08x; %d bits = %d bytes\n",
                       bs.current_bank,
                       bs.current_width,
                       bs.current_height,
                       bs.current_offset,
                       bs.current_width * bs.current_height,
                       (bs.current_width * bs.current_height) / 8);
                bs.bram_width = MAX(bs.bram_width, bs.current_width);
                bs.bram_height = MAX(bs.bram_height, bs.current_height);
                for (i = 0; i < ((bs.current_width * bs.current_height) / 8); i += 2)
                {
                    uint16_t word = 
                        // ((irw_readb(f) << 24) & 0xff000000)
                        // |
                        // ((irw_readb(f) << 16) & 0x00ff0000)
                        // |
                        ((irw_readb(f) << 8) & 0x0000ff00)
                        |
                        ((irw_readb(f) << 0) & 0x000000ff)
                        ;
                    int words_per_block = 0;
                    if (bs.current_width == 160)
                        words_per_block = 5;
                    else if (bs.current_width == 80)
                        words_per_block = 3;
                    else
                        printf("unrecognized width: %d\n", bs.current_width);
                    if (ora_matches >= words_per_block) {
                        ora_matches = 0;
                        ora_ptr = ora_ptr + 10;
                    }
                    int i;
                    int found = 0;
                    for (i = 0; i < 16; i++) {
                        if (ora16[i + ora_ptr] == word) {
                            found = 1;
                            printf("found %04x @ %d\n", word, i + ora_ptr);
                            word = oro16[i + ora_ptr];
                            ora_matches++;
                            break;
                        }
                    }
                    if (!found) {
                        printf("couldn't find word %04x\n", word);
                    }
                    // printf("%04x -> %04x\n", word, mapping[word]);
                    // word = get_mapping(&mapping, word);
                    // irw_writeb(o, word >> 24);
                    // irw_writeb(o, word >> 16);
                    // word = mappings[word];
                    irw_writeb(o, word >> 8);
                    irw_writeb(o, word);
                }
                last0 = irw_readb(f);
                last1 = irw_readb(f);
                if (last0 || last1)
                {
                    printf("expected 0x0000 after BRAM data, got %02x %02x\n", last0, last1);
                }
                irw_writeb(o, last0);
                irw_writeb(o, last1);
                break;

            // Reset CRC
            case 5:
                f->crc = 0xffff;
                o->crc = 0xffff;
                break;

            // Wakeup
            case 6:
                wakeup = 1;
                break;

            default:
                printf("unrecognized command 0x%02x 0x%02x\n", cmd, payload);
                break;
            }
            break;

        // Set current bank
        case 1:
            bs.current_bank = payload;
            ora_ptr = 0;
            ora_matches = 0;
            // printf("setting bank number to %d\n", bs.current_bank);
            break;

        // Validate CRC16
        case 2:
            printf("crc check (%04x == %04x)\n", f->crc, 0);
            uint16_t crc16 = o->crc;
            irw_writeb(o, crc16 >> 8);
            irw_writeb(o, crc16);
            break;

        // Set frequency range
        case 5:
            switch (payload)
            {
            case 0:
                bs.frequency_range = 0;
                break;
            case 1:
                bs.frequency_range = 1;
                break;
            case 2:
                bs.frequency_range = 2;
                break;
            default:
                printf("unknown frequency range payload: %02x\n", payload);
                break;
            }
            break;

        // Set current width
        case 6:
            bs.current_width = payload + 1;
            break;

        // Set current height
        case 7:
            bs.current_height = payload;
            break;

        // Set current ofset
        case 8:
            bs.current_offset = payload;
            break;

        // Set flags
        case 9:
            switch (payload)
            {
            case 0:
                bs.warmboot = 0;
                bs.nosleep = 0;
                break;
            case 1:
                bs.warmboot = 0;
                bs.nosleep = 1;
                break;
            case 32:
                bs.warmboot = 1;
                bs.nosleep = 0;
                break;
            case 33:
                bs.warmboot = 1;
                bs.nosleep = 1;
                break;
            default:
                printf("unrecognized feature flags: %02x\n", payload);
                break;
            }
            break;

        default:
            printf("unrecognized command: %02x\n", cmd);
            break;
        }
    }

    // Padding
    irw_writeb(o, 0);
    return 0;
}

int main(int argc, char **argv)
{

    struct irw_file *input = irw_open("top.bin", "r");
    struct irw_file *output = irw_open("new.bin", "w");
    struct irw_file *rom = irw_open("rom.bin", "r");

    // int i;
    // for (i = 0; i < 32; i++) {
    //     printf("bit %d: %d\n", i, get_bit_offset(i, 8192*8));
    // }

    if (!input)
    {
        perror("couldn't open top.bin");
        return 1;
    }

    if (!output) {
        perror("couldn't open new.bin");
        return 2;
    }

    // struct irw_file *mem = irw_open("memtest/memtest.bin", "r");
    // ice40_parse(mem);
    ice40_patch(input, rom, output);
    return 0;
}