#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 "iccard/aime.h"
#include "iccard/mifare.h"

#include "board/bpreader.h"

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

static bool flg = true;
static struct mifare card;
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_44_cmd();
static HRESULT bpreader_unk_52_cmd();
static HRESULT bpreader_unk_54_cmd();
static HRESULT bpreader_send_mifare_cmd();
static HRESULT bpreader_empty_resp();

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;
    }
    uint8_t sec1[16] = { 0x00, 0x02, 0x4E, 0x42, 0x47, 0x49, 0x43, 0x36, 0x3A, 0xE7, 0xB5, 0x59, 0x45, 0x64, 0x0F, 0x3D };

    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);

    HRESULT hr = aime_card_populate(&card, cfg->access_code_bytes, sizeof(cfg->access_code_bytes));
    //memcpy_s(card.sectors[0].blocks[1].bytes, sizeof(card.sectors[0].blocks[1].bytes), sec1, sizeof(sec1));
    
    if (FAILED(hr)) {
        return hr;
    }
    
    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 (!read_ct) {
            hr = crack_bpreader_request();
        } else {
            // We need a delay here in order to properly
            // "timeout" the DLL's read request.
            Sleep(1);
        }
        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
    
    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 Mifare\n");
            return bpreader_send_mifare_cmd();

        case 0x44:
            dprintf("Reader: Cmd 0x44\n");
            return bpreader_unk_44_cmd();

        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");
            return bpreader_unk_52_cmd();

        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(config->insert_card) & 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->access_code_bytes, 4);
        data_len = sizeof(data);
    
    
        resp_len = build_bpreader_response((uint8_t *)&data, data_len, &header, buff, sizeof(buff));
    }
    else {
        data_len = 3;
        uint8_t empty_buff[1] = {0};
        resp_len = build_bpreader_response((uint8_t *)&empty_buff, 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, 1, &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_unk_54_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;
    }
}

static HRESULT bpreader_send_mifare_cmd()
{
    struct bpreader_read_mifare_req req = { bp_uart.written.bytes[7], bp_uart.written.bytes[8], bp_uart.written.bytes[9] };
    struct bpreader_cmd_header header = { 0x00, 0x00, 0xFF, 0x03, 0xFD, 0xD5, 0x41 };
    uint8_t data[17] = { 0x00 };
    uint8_t buff[26] = { 0x00 };
    size_t resp_len = 0;

    switch (req.subcmd) {
        case MIFARE_CMD_READ:
            struct bpreader_read_mifare_resp_data resp;
            
            if (req.block > 3) {
                dprintf("Reader: MIFARE block number %d is out of range!\n", req.block);

                return E_FAIL;
            }

            dprintf("Reader: Mifare Read Block 0x%02X\n", req.block);
            resp.padding = 0;
            memcpy_s(resp.block, sizeof(resp.block), card.sectors[0].blocks[req.block].bytes, sizeof(card.sectors[0].blocks[req.block].bytes));
            header.data_len = 0x13;
            resp_len = build_bpreader_response((uint8_t *)&resp, 0x11, &header, buff, sizeof(buff));
            break;
        
        case MIFARE_CMD_AUTH_KEY_A:
            dprintf("Reader: Mifare Authentication with Key A\n");
            return iobuf_write(&bp_uart.readable, BPREADER_CMD_GO_NEXT, sizeof(BPREADER_CMD_GO_NEXT));
            break;
        
        case MIFARE_CMD_AUTH_KEY_B:
            dprintf("Reader: Mifare Authentication with Key B\n");
            resp_len = build_bpreader_response(data, 1, &header, buff, sizeof(buff));
            break;
        
        default:
            dprintf("Reader: Unknown Mifare Command 0x%02x\n", req.subcmd);
            return E_FAIL;
    }
    
    if (resp_len > 0) {
        return iobuf_write(&bp_uart.readable, buff, resp_len);
    
    } else {
        dprintf("Reader: No data to return in bpreader_send_mifare_cmd!\n");
        return E_FAIL;
    }
}

static HRESULT bpreader_unk_52_cmd()
{
    struct bpreader_cmd_header header = { 0x00, 0x00, 0xFF, 0x04, 0xFC, 0xD5, 0x53 };
    uint8_t buff[10] = { 0x00 };
    uint8_t bfr_data[2] = { 0x01, 0x00 };
    size_t resp_len = 0;

    resp_len = build_bpreader_response(bfr_data, 2, &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_unk_52_cmd!\n");
        return E_FAIL;
    }
}

static HRESULT bpreader_unk_44_cmd()
{
    struct bpreader_cmd_header header = { 0x00, 0x00, 0xFF, 0x04, 0xFC, 0xD5, 0x45 };
    uint8_t buff[10] = { 0x00 };
    uint8_t bfr_data[2] = { 0x01, 0x00 };
    size_t resp_len = 0;

    resp_len = build_bpreader_response(bfr_data, 2, &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_unk_44_cmd!\n");
        return E_FAIL;
    }
}