#include <windows.h>
#include <stdint.h>
#include <stdbool.h>

#include "hook/iobuf.h"
#include "hook/iohook.h"

#include "hooklib/uart.h"
#include "hooklib/fdshark.h"

#include "util/dprintf.h"
#include "util/dump.h"
#include "util/hexstr.h"

#include "board/bpreader.h"

const uint8_t BPREADER_CMD_GO_NEXT[6] = { 0x00, 0x00, 0xFF, 0x00, 0xFF, 0x00 };

static bool flg = true;
static bool flg_second_pass = false;
static HRESULT bp_handle_irp(struct irp *irp);
static HRESULT bp_handle_irp_locked(struct irp *irp);
static HRESULT crack_bpreader_request();
static size_t build_bpreader_response(
    const uint8_t *data,
    const size_t len_data,
    struct bpreader_cmd_header *resp_header,
    uint8_t *response,
    const size_t resp_size);

static HRESULT bpreader_generic_cmd(uint32_t resp_code);
static HRESULT bpreader_poll_card_cmd();
static HRESULT bpreader_init_cmd();
static HRESULT bpreader_set_output_cmd();
static HRESULT bpreader_unk_08_cmd();
static HRESULT bpreader_unk_06_cmd();
static HRESULT bpreader_unk_0c_cmd();
static HRESULT bpreader_unk_54_cmd();

static struct bpreader_config *config;
static struct uart bp_uart;
static CRITICAL_SECTION bp_lock;
static uint8_t bp_written_bytes[520];
static uint8_t bp_readable_bytes[520];
static uint8_t last_cmd = 0;
static uint16_t read_ct = 0;

HRESULT bpreader_init(struct bpreader_config *cfg, uint16_t port)
{
    config = cfg;
    if (!config->enable) {
        return S_OK;
    }
    if (cfg->port > 0) {
        port = cfg->port;
    }

    uart_init(&bp_uart, port);
    bp_uart.written.bytes = bp_written_bytes;
    bp_uart.written.nbytes = sizeof(bp_written_bytes);
    bp_uart.readable.bytes = bp_readable_bytes;
    bp_uart.readable.nbytes = sizeof(bp_readable_bytes);
    InitializeCriticalSection(&bp_lock);
    
    dprintf("Reader: Init\n");
    return iohook_push_handler(bp_handle_irp);
}

static HRESULT bp_handle_irp(struct irp *irp)
{
    HRESULT hr;

    assert(irp != NULL);

    if (uart_match_irp(&bp_uart, irp)) {
        EnterCriticalSection(&bp_lock);
        hr = bp_handle_irp_locked(irp);
        LeaveCriticalSection(&bp_lock);
    }
    else {
        return iohook_invoke_next(irp);
    }

    return hr;
}

static HRESULT bp_handle_irp_locked(struct irp *irp)
{
    HRESULT hr;
    if (irp->op == IRP_OP_OPEN) {
        dprintf("Reader: Starting backend\n");
    }
    
    hr = uart_handle_irp(&bp_uart, irp);    
    if (FAILED(hr)) {
        return hr;
    }

#if 0
    if (irp->op == IRP_OP_WRITE) {
        dprintf("WRITE:\n");
        dump_iobuf(&bp_uart.written);
    }
#endif
    
    if (irp->op == IRP_OP_WRITE) {
        read_ct = 0;
        if (bp_uart.written.bytes[0] == 0x55) {
            dprintf("Reader: Hello\n");
            bp_uart.written.pos = 0; // consume the written buffer
        }
        else if (bp_uart.written.bytes[3] == 0x00) {
            bp_uart.written.pos = 0; // consume the written buffer
        }
        return hr;
    }
    
    else if (irp->op == IRP_OP_READ) {
        if (flg) {
            if (!read_ct || flg_second_pass) {
                hr = crack_bpreader_request();
            }
        } else {
            switch (bp_uart.written.bytes[3]) {
                case 0x02:
                    if (!read_ct) {                    
                        dprintf("Reader: Unknown 0x02\n");
                        uint8_t buff[] = { 0x00, 0x00, 0xFF, 0x05,
                            0xFB, 0xD5, 0x0D, 0x00, 0x06, 0x00, 0x18, 0x00 };
                        hr = iobuf_write(&bp_uart.readable, buff, sizeof(buff));
                    }
                    
                    bp_uart.written.pos = 0;
                    break;

                case 0x03:
                    if (!read_ct) {                    
                        dprintf("Reader: Unknown 0x03\n");
                        uint8_t buff[] = { 0x00, 0x00, 0xFF, 0x02, 
                            0xFE, 0xD5, 0x19, 0x12, 0x00 };
                        hr = iobuf_write(&bp_uart.readable, buff, sizeof(buff));
                    }
                    bp_uart.written.pos = 0;
                    break;
                
                case 0x04:
                    if (!read_ct && bp_uart.written.bytes[6] == 0x0E && last_cmd == 0x04) {
                        dprintf("Reader: Unknown second 0x04\n");
                        uint8_t buff[] = { 0x00, 0x00, 0xFF, 0x02,
                            0xFE, 0xD5, 0x0F, 0x1C, 0x00 };
                        hr = iobuf_write(&bp_uart.readable, buff, sizeof(buff));
                    }
                    else if (!read_ct) {  
                        dprintf("Reader: Unknown 0x04\n");
                        uint8_t buff[] = { 0x00, 0x00, 0xFF, 0x02,
                            0xFE, 0xD5, 0x33, 0xF8, 0x00 };
                        hr = iobuf_write(&bp_uart.readable, buff, sizeof(buff));
                    }
                    bp_uart.written.pos = 0;
                    break;
                
                case 0x06:
                    if (!read_ct) {  
                        dprintf("Reader: Unknown 0x06\n");
                        uint8_t buff[] = { 0x00, 0x00, 0xFF, 0x02,
                            0xFE, 0xD5, 0x33, 0xF8, 0x00 };
                        hr = iobuf_write(&bp_uart.readable, buff, sizeof(buff));
                    }
                    bp_uart.written.pos = 0;
                    break;
        
                case 0x05:
                    if (!read_ct) {  
                        dprintf("Reader: Unknown 0x05\n");
                        uint8_t buff[] = { 0x00, 0x00, 0xFF, 0x03,
                            0xFD, 0xD5, 0x09, 0x00, 0x22, 0x00 };
                        hr = iobuf_write(&bp_uart.readable, buff, sizeof(buff));
                    }
                    bp_uart.written.pos = 0;
                    break;

                case 0x0E:
                    if (!read_ct) {
                        dprintf("Reader: Unknown 0x0E\n");
                        uint8_t buff[] = { 0x00, 0x00, 0xff, 0x02,
                            0xfe, 0xd5, 0x33, 0xf8, 0x00 };
                        hr = iobuf_write(&bp_uart.readable, buff, sizeof(buff));
                    }
                    bp_uart.written.pos = 0;
                    break;
        
                case 0x12:
                    if (!read_ct) {  
                        dprintf("Reader: Unknown 0x12\n");
                        uint8_t buff[] = { 0x00, 0x00, 0xFF, 0x0a,
                            0xf6, 0xd5, 0x07, 0xff, 0x3f, 0x0e, 0xf1, 0xff, 0x3f, 0x0e, 0xf1, 0xaa, 0x00 };
                        hr = iobuf_write(&bp_uart.readable, buff, sizeof(buff));
                    }
                    bp_uart.written.pos = 0;
                    break;
                
                case 0x14:
                    if (!read_ct) {
                        dprintf("Reader: Unknown 0x14\n");
                        uint8_t buff[] = { 0x00, 0x00, 0xFF, 0x20, 0xE0, 0xD5, 0xA1, 
                            0x00, 0x1D, 0x07, // Unknown
                            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // IDm
                            0x00, 0x00, 0x01, // Unknown
                            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // IDm
                            0x00, 0x78, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,  // Unknown
                            0x00, 0x00 };
                        hr = iobuf_write(&bp_uart.readable, buff, sizeof(buff));
                    }
                    bp_uart.written.pos = 0;
                    break;
                
                case 0x18:
                    if (!read_ct) {
                        dprintf("Reader: Unknown 0x18\n");
                        uint8_t buff[] = { 0x00, 0x00, 0xFF, 0x0D,
                            0xF3, 0xD5, 0x07, 0xDC, 0xF4, 0x3F, 0x11, 0x4D, 0x85, 0x61, 0xF1, 0x26,
                            0x6A, 0x87, 0xC9, 0x00 };
                        hr = iobuf_write(&bp_uart.readable, buff, sizeof(buff));
                    }
                    bp_uart.written.pos = 0;
                    break;

                default:
                    dprintf("Reader: Unknown Command %02X\n", last_cmd);
                    dump_iobuf(&bp_uart.written);
                    break;
            }
        }   
        read_ct++;
    }

    if (FAILED(hr)) {
        dprintf("Reader: HR failed %lX\n", hr);
    }

#if 0    
    if (irp->op == IRP_OP_READ) {
        dprintf("READABLE:\n");
        dump_iobuf(&bp_uart.readable);
    }
#endif
    
    if (!flg_second_pass) {
        bp_uart.written.pos = 0; // consume the written buffer
    }
         
    return hr;
}

static HRESULT crack_bpreader_request() {
    struct bpreader_cmd_header header = { bp_uart.written.bytes[0], bp_uart.written.bytes[1], bp_uart.written.bytes[2],
        bp_uart.written.bytes[3], bp_uart.written.bytes[4], bp_uart.written.bytes[5], bp_uart.written.bytes[6], 
    };
    
    assert(header.padding0_00 == 0);
    assert(header.padding1_00 == 0);
    assert(header.padding2_ff == 0xFF);
    assert(header.d_identifier == 0xD4);
    last_cmd = header.cmd;
    
    switch(header.cmd) {
        case 0x06:
            dprintf("Reader: Cmd 0x06\n");
            return bpreader_unk_06_cmd();

        case 0x08:
            dprintf("Reader: Cmd 0x08\n");
            return bpreader_unk_08_cmd();
    
        case 0x12:
            dprintf("Reader: Cmd 0x12\n");
            return bpreader_generic_cmd(0x13);

        case 0x18:
            dprintf("Reader: Initialize\n");
            return bpreader_init_cmd();
        
        case 0x0C:
            dprintf("Reader: Cmd 0x0C\n");
            return bpreader_unk_0c_cmd();
        
        case 0x0E:
            dprintf("Reader: Set Output\n");
            return bpreader_generic_cmd(0x0F);
        
        case 0x32:
            dprintf("Reader: Cmd 0x32\n");
            return bpreader_generic_cmd(0x33);
        
        case 0x40:
            dprintf("Reader: Read Banapass\n"); // 01 30 00 -> Chip ID; 01 30 01 -> Thing after chip ID; 01 30 02 -> Access Code; 01 60 30 -> send key a
            break;

        case 0x4A:
            dprintf("Reader: Poll Card\n");
            return bpreader_poll_card_cmd();

        case 0xA0:
            dprintf("Reader: Read Felica\n");
            break;

        case 0x52:
            dprintf("Reader: Cmd 0x52\n");
            break;

        case 0x54:
            dprintf("Reader: Cmd 0x54\n");
            return bpreader_unk_54_cmd();
        
        default:
            dprintf("Reader: Unknown command 0x%02x\n", header.cmd);
    }

    dump_iobuf(&bp_uart.written);
    return (HRESULT)-1;
}

static size_t build_bpreader_response(
    const uint8_t *data,
    const size_t len_data,
    struct bpreader_cmd_header *resp_header,
    uint8_t *response,
    const size_t resp_size) 
{
    // TODO: Account for escape characters
    if (resp_size < 9) {
        dprintf("Reader: build_bpreader_response resp_size less then 9\n");
        return 0;
    }

    uint8_t full_resp_len = len_data + 2 + 7; // calculate full response length
    uint8_t final_checksum = 0; // initialize checksum
    
    resp_header->data_len = len_data + 2;
    resp_header->header_checksum = 0xFF - ((0xFF + len_data + 2) & 0xFF);
    
    memcpy_s(response, resp_size, resp_header, sizeof(*resp_header)); // Copy header
    if (len_data > 0) {        
        for (int i = 0; i < len_data; i++) {
            response[7+i] = data[i];
        } // copy data if there's any
    }
    
    for (int i = 0; i < full_resp_len; i++) {
        final_checksum = (final_checksum + response[i]) & 0xFF; // Calculate checksum
    }

    response[full_resp_len - 2] = 0xFF - final_checksum; // Assign checksum
    return full_resp_len;
}

static HRESULT bpreader_generic_cmd(uint32_t resp_code)
{
    //uint8_t buff[] = { 0x00, 0x00, 0xFF, 0x02, 0xFE, 0xD5, 0x19, 0x12, 0x00 };
    uint8_t buff[9] = { 0x00 };
    struct bpreader_cmd_header header = { 0x00, 0x00, 0xFF, 0x02, 0xFE, 0xD5, resp_code };
    size_t resp_len = 0;

    resp_len = build_bpreader_response((uint8_t *){0}, 0, &header, buff, sizeof(buff));
    
    if (resp_len > 0) {
        HRESULT hr = iobuf_write(&bp_uart.readable, buff, resp_len);
        //dump_iobuf(&bp_uart.readable);
        return hr;
    
    } else {
        dprintf("Reader: No data to return in bpreader_generic_cmd (0x%2X response code)!\n", resp_code);
        return E_FAIL;
    }
}

static HRESULT bpreader_init_cmd()
{
    //uint8_t buff[] = { 0x00, 0x00, 0xFF, 0x02, 0xFE, 0xD5, 0x19, 0x12, 0x00 };
    uint8_t buff[9] = { 0x00 };
    struct bpreader_cmd_header header = { 0x00, 0x00, 0xFF, 0x02, 0xFE, 0xD5, 0x19 };
    size_t resp_len = 0;

    resp_len = build_bpreader_response((uint8_t *){0}, 0, &header, buff, sizeof(buff));
    
    if (resp_len > 0) {
        HRESULT hr = iobuf_write(&bp_uart.readable, buff, resp_len);
        return hr;
    
    } else {
        dprintf("Reader: No data to return in bpreader_init_cmd!\n");
        return E_FAIL;
    }
}

static HRESULT bpreader_set_output_cmd()
{
    //uint8_t buff[] = { 0x00, 0x00, 0xFF, 0x02, 0xFE, 0xD5, 0x0F, 0x1C, 0x00 };
    uint8_t buff[9] = { 0x00 };
    struct bpreader_cmd_header header = { 0x00, 0x00, 0xFF, 0x02, 0xFE, 0xD5, 0x0F };
    size_t resp_len = 0;

    resp_len = build_bpreader_response((uint8_t *){0}, 0, &header, buff, sizeof(buff));
    
    if (resp_len > 0) {
        HRESULT hr = iobuf_write(&bp_uart.readable, buff, resp_len);
        return hr;
    
    } else {
        dprintf("Reader: No data to return in bpreader_set_output_cmd!\n");
        return E_FAIL;
    }
}

static HRESULT bpreader_poll_card_cmd()
{
    struct bpreader_cmd_header header = { 0x00, 0x00, 0xFF, 0x03, 0xFD, 0xD5, 0x4B };
    struct bpreader_poll_banapass_data data = { 0x00 };
    uint8_t buff[32] = { 0x00 };
    int data_len = 1;
    size_t resp_len = 0;
    
    if (GetAsyncKeyState(VK_RETURN) & 0x8000) {
        dprintf("Reader: Touch card %ls\n", config->access_code);
        
        data.card_present = 0x01;
        data.unknown1 = 0x01;
        data.atqa[0] = 0x00;
        data.atqa[1] = 0x04;
        data.sak = 0x08;
        data.unknown5 = 0x04;
        memcpy_s(data.serial_number, sizeof(data.serial_number), config->chip_id_bytes, 4);
        data_len = sizeof(data);
    }
    
    resp_len = build_bpreader_response((uint8_t *)&data, data_len, &header, buff, sizeof(buff));
    
    if (resp_len > 0) {
        HRESULT hr = iobuf_write(&bp_uart.readable, buff, resp_len);
        return hr;
    
    } else {
        dprintf("Reader: No data to return in bpreader_poll_card_cmd!\n");
        return E_FAIL;
    }
}

static HRESULT bpreader_unk_08_cmd() 
{
    struct bpreader_cmd_header header = { 0x00, 0x00, 0xFF, 0x03, 0xFD, 0xD5, 0x09 };
    uint8_t buff[10] = { 0x00 };
    uint8_t bfr_data[1] = { 0x00 };
    size_t resp_len = 0;

    resp_len = build_bpreader_response(bfr_data, sizeof(bfr_data), &header, buff, sizeof(buff));

    if (resp_len > 0) {
        HRESULT hr = iobuf_write(&bp_uart.readable, buff, resp_len);
        return hr;
    
    } else {
        dprintf("Reader: No data to return in bpreader_set_output_cmd!\n");
        return E_FAIL;
    }
}

static HRESULT bpreader_unk_54_cmd() 
{
    struct bpreader_cmd_header header = { 0x00, 0x00, 0xFF, 0x03, 0xFD, 0xD5, 0x55 };
    uint8_t buff[10] = { 0x00 };
    uint8_t bfr_data[1] = { 0x00 };
    size_t resp_len = 0;

    resp_len = build_bpreader_response(bfr_data, sizeof(bfr_data), &header, buff, sizeof(buff));

    if (resp_len > 0) {
        HRESULT hr = iobuf_write(&bp_uart.readable, buff, resp_len);
        return hr;
    
    } else {
        dprintf("Reader: No data to return in bpreader_set_output_cmd!\n");
        return E_FAIL;
    }
}

static HRESULT bpreader_unk_06_cmd() 
{
    //uint8_t hdr[] = { 0x00, 0x00, 0xFF, 0x0a, 0xf6, 0xd5, 0x07 } 
    //uint8_t data[] = { 0xff, 0x3f, 0x0e, 0xf1, 0xff, 0x3f, 0x0e, 0xf1, 0xaa, 0x00 };
    struct bpreader_cmd_header header = { 0x00, 0x00, 0xFF, 0x0a, 0xF6, 0xD5, 0x07 };
    uint8_t buff_a[17] = { 0x00 };
    uint8_t buff_b[20] = { 0x00 };
    uint8_t bfr_data_a[8] = { 0xff, 0x3f, 0x0e, 0xf1, 0xff, 0x3f, 0x0e, 0xf1 };
    uint8_t bfr_data_b[11] = { 0xDC, 0xF4, 0x3F, 0x11, 0x4D, 0x85, 0x61, 0xF1, 0x26, 0x6A, 0x87 };
    size_t resp_len = 0;
    uint8_t sample = bp_uart.written.bytes[8];

    switch (sample) {
        case 0x1c: resp_len = build_bpreader_response(bfr_data_a, sizeof(bfr_data_a), &header, buff_a, sizeof(buff_a)); break;
        default: resp_len = build_bpreader_response(bfr_data_b, sizeof(bfr_data_b), &header, buff_b, sizeof(buff_b)); break;
    }

    if (resp_len > 0) {
        switch (sample) {
        case 0x1c: return iobuf_write(&bp_uart.readable, buff_a, sizeof(buff_a));
        default: return iobuf_write(&bp_uart.readable, buff_b, sizeof(buff_b));
    }
        
    } else {
        dprintf("Reader: No data to return in bpreader_set_output_cmd!\n");
        return E_FAIL;
    }
}

static HRESULT bpreader_unk_0c_cmd()
{
    //uint8_t hdr[] = uint8_t buff[] = { 0x00, 0x00, 0xFF, 0x05, 0xFB, 0xD5, 0x0D };
    //uint8_t data[] = { 0x00, 0x06, 0x00, 0x18, 0x00 };
    struct bpreader_cmd_header header = { 0x00, 0x00, 0xFF, 0x0a, 0xF6, 0xD5, 0x0D };
    uint8_t buff[12] = { 0x00 };
    uint8_t bfr_data[3] = { 0x00, 0x06, 0x00 };
    size_t resp_len = 0;

    resp_len = build_bpreader_response(bfr_data, sizeof(bfr_data), &header, buff, sizeof(buff));

    if (resp_len > 0) {
        HRESULT hr = iobuf_write(&bp_uart.readable, buff, resp_len);
        return hr;
    
    } else {
        dprintf("Reader: No data to return in bpreader_set_output_cmd!\n");
        return E_FAIL;
    }
}