#include <unicore-mx/stm32/gpio.h>
#include <unicore-mx/stm32/rcc.h>
#include <unicore-mx/stm32/f4/rcc.h>
#include <unicore-mx/stm32/usart.h>
#include <unicore-mx/cm3/nvic.h>
#include <unicore-mx/usbd/usbd.h>
#include <unicore-mx/usbd/class/msc.h>
#include <unicore-mx/stm32/otg_fs.h>
#include <wolfssl/wolfcrypt/settings.h>
#include <wolfssl/options.h>
#include <wolfssl/wolfcrypt/sha256.h>
#include <wolfssl/wolfcrypt/aes.h>
#include <wolfssl/wolfcrypt/chacha.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <stdio.h>
#include "system.h"
#include "systick.h"
#include "sdcard.h"

//#define USB_DEBUG
#ifdef USB_DEBUG
#define USB_DBG printf
#define USB_DBG_BIN printbin
#else
#define USB_DBG(...) do{}while (0)
#define USB_DBG_BIN(...) do{}while (0)
#endif

#define CRYPTO_BLOCK_SIZE 64
#define KEYSZ 32
#define IVSZ 16

#define SECTOR_SIZE 512
static const char VID[] = "0001";
static const char PID[] = "0002";
static const char PR[]  = "0003";
static const char iv_bytes[16] = {0x00, 0x00, 0x00, 0x00, 0xa8, 0xbb, 0x2a, 0x3e, 0x3c, 0xb7, 0x19, 0x27, 0x77, 0xbf, 0x3e, 0x60 };
static struct usbd_device *usbd_dev = NULL;
static uint8_t tmp_blk_rd[SECTOR_SIZE];
static uint8_t tmp_blk_wr[SECTOR_SIZE];
static uint32_t bridge_mutex = 1;

int bridge_lock(void)
{
    while(_mutex_lock(&bridge_mutex) == 0)
        ;
    asm volatile ("cpsid i");
    return 0;
}

int bridge_unlock(void)
{
    asm volatile ("cpsie i");
    return _mutex_unlock(&bridge_mutex);
}

static int bridge_read_block(const usbd_msc_backend *backend, uint32_t lba, void *copy_to)
{
    uint32_t sd_size = sdcard_getsize(NULL);
    if (lba > sd_size)
        return -1;
    sdcard_read(NULL, copy_to, lba);
    return 0;
}

static int bridge_write_block(const usbd_msc_backend *backend, uint32_t lba, const void *copy_from)
{
    uint32_t sd_size = sdcard_getsize(NULL);
    if (lba > sd_size)
        return -1;
    sdcard_write(NULL, copy_from, lba);
    return 0;
}

static int bridge_read_block_enc(const usbd_msc_backend *backend, uint32_t lba, void *copy_to)
{
    ChaCha cha;
    uint8_t crypto_tmp[CRYPTO_BLOCK_SIZE];
    uint32_t i;
    uint32_t sd_size = sdcard_getsize(NULL);
    volatile uint8_t *dst;
    uint8_t crypto_iv[16];
    if (lba > sd_size)
        return -1;
    if (lba < 2) {
        USB_DBG("SD: requested block %08x\r\n", lba);
    }
    /*
    if (lba == 0) {
        dst = copy_to;
        memcpy(copy_to, BootSector, SECTOR_SIZE);
        dst[SECTOR_SIZE - 2] = 0x55;
        dst[SECTOR_SIZE - 1] = 0xAA;
        return 0;
    }
    */
    memset(copy_to, 0, SECTOR_SIZE);
    dst = (uint8_t *)tmp_blk_rd;
    sdcard_read(NULL, dst, lba);
    wc_Chacha_SetKey(&cha, usecfs_getkey(), 32);
    memcpy(crypto_iv, iv_bytes, 16);
    *((uint32_t *)crypto_iv) = lba;
    wc_Chacha_SetIV(&cha, crypto_iv, 16);
    for (i = 0; i < SECTOR_SIZE / CRYPTO_BLOCK_SIZE; i++) {
        memcpy(crypto_tmp, dst + i * CRYPTO_BLOCK_SIZE, CRYPTO_BLOCK_SIZE);
        wc_Chacha_Process(&cha, dst + i * CRYPTO_BLOCK_SIZE, crypto_tmp, CRYPTO_BLOCK_SIZE);
    }
    if (lba < 2) {
        USB_DBG("Decrypted:\r\n");
        USB_DBG_BIN(dst, SECTOR_SIZE);
    }
    memcpy(copy_to, dst, SECTOR_SIZE);
    return 0;
}



static int bridge_write_block_enc(const usbd_msc_backend *backend, uint32_t lba, const void *copy_from)
{
    ChaCha cha;
    uint8_t crypto_tmp[CRYPTO_BLOCK_SIZE];
    uint32_t i;
    const uint8_t *src;
    uint8_t crypto_iv[16];
    uint32_t sd_size;

    bridge_lock();
    sd_size = sdcard_getsize(NULL);
    if (lba > sd_size) {
        printf("SD error: attempting to write to block %u out of range\r\n", lba);
        return -1;
    }
    /*
    if (lba == 0) {
        return 0;
    }
    */
    if (lba < 2) {
        USB_DBG("SD: provided block %u\r\n", lba);
        USB_DBG_BIN(copy_from, SECTOR_SIZE);
    }
    memcpy(tmp_blk_wr, copy_from, SECTOR_SIZE);
    src = (uint8_t *)tmp_blk_wr;
    wc_Chacha_SetKey(&cha, usecfs_getkey(), 32);
    memcpy(crypto_iv, iv_bytes, 16);
    *((uint32_t *)crypto_iv) = lba;
    wc_Chacha_SetIV(&cha, crypto_iv, 16);
    for (i = 0; i < SECTOR_SIZE / CRYPTO_BLOCK_SIZE; i++) {
        memcpy(crypto_tmp, src + i * CRYPTO_BLOCK_SIZE, CRYPTO_BLOCK_SIZE);
        wc_Chacha_Process(&cha, tmp_blk_wr + i * CRYPTO_BLOCK_SIZE, crypto_tmp, CRYPTO_BLOCK_SIZE);
    }
    sdcard_write(NULL, tmp_blk_wr, lba);
    bridge_unlock();
    return 0;
}


static uint8_t usb_buffer[SECTOR_SIZE];

struct usbd_msc_backend flash_backend = {
	.vendor_id = VID,
    .product_id = PID,
    .product_rev = PR,
    .block_count = 0,
    //.lock = bridge_lock,
    //.unlock = bridge_unlock,
    .read_block = bridge_read_block,
    .write_block = bridge_write_block,
    .block_size = SECTOR_SIZE,
    .scratch_buffer = usb_buffer
};

const usbd_backend *msc_target_usb_driver(void)
{
    return USBD_STM32_OTG_FS;
}

static const struct usb_string_descriptor string_lang_list = {
    .bLength = USB_DT_STRING_SIZE(1),
    .bDescriptorType = USB_DT_STRING,
    .wData = {
        USB_LANGID_ENGLISH_UNITED_STATES
    }
};

const struct usb_string_descriptor usb_string_manuf = {
    .bLength = USB_DT_STRING_SIZE(10),
    .bDescriptorType = USB_DT_STRING,
    /* danielinux */
    .wData = {
        0x0064, 0x0061, 0x006e, 0x0069, 0x0065, 0x006c, 0x0069, 0x006e, 0x0075, 0x0078
    }
};

static const struct usb_string_descriptor usb_string_name = {
    .bLength = USB_DT_STRING_SIZE(42),
    .bDescriptorType = USB_DT_STRING,
    /* "Encrypted SD Card in Secure Storage device" */
    .wData = {
        0x0045, 0x006e, 0x0063, 0x0072, 0x0079, 0x0070, 0x0074, 0x0065, 0x0064, 0x0020, 0x0053, 0x0044,
        0x0020, 0x0043, 0x0061, 0x0072, 0x0064, 0x0020, 0x0069, 0x006e, 0x0020, 0x0053, 0x0065, 0x0063,
        0x0075, 0x0072, 0x0065, 0x0020, 0x0053, 0x0074, 0x006f, 0x0072, 0x0061, 0x0067, 0x0065, 0x0020,
        0x0064, 0x0065, 0x0076, 0x0069, 0x0063, 0x0065
    }
};

static const struct usb_string_descriptor usb_serialn = {
    .bLength = USB_DT_STRING_SIZE(3),
    .bDescriptorType = USB_DT_STRING,
    /* 000 */
    .wData = {
        0x0030, 0x0030, 0x0030
    }
};

static const struct usb_string_descriptor **string_data[1] = {
    (const struct usb_string_descriptor *[]) {
        &usb_string_manuf,
        &usb_string_name,
        &usb_serialn
    }
};

static const struct usb_device_descriptor msc_dev_desc= {
    .bLength = USB_DT_DEVICE_SIZE,
    .bDescriptorType = USB_DT_DEVICE,
    .bcdUSB = 0x0110,
    .bDeviceClass = 0,
    .bDeviceSubClass = 0,
    .bDeviceProtocol = 0,
    .bMaxPacketSize0 = 64,
    .idVendor = 0x1d50,
    .idProduct = 0xDAD0,
    .bcdDevice = 0x0200,
    .iManufacturer = 1,
    .iProduct = 2,
    .iSerialNumber = 3,
    .bNumConfigurations = 1
};

static const struct __attribute__((packed)) {
    struct usb_config_descriptor config;
    struct usb_interface_descriptor msc_iface;
    struct usb_endpoint_descriptor msc_endp[2];
} msc_config = {
    .config = {
        .bLength = USB_DT_CONFIGURATION_SIZE,
        .bDescriptorType = USB_DT_CONFIGURATION,
        .wTotalLength = sizeof(msc_config),
        .bNumInterfaces = 1,
        .bConfigurationValue = 1,
        .iConfiguration = 0,
        .bmAttributes = 0x80,
        .bMaxPower = 0x32
    },

    .msc_iface = {
        .bLength = USB_DT_INTERFACE_SIZE,
        .bDescriptorType = USB_DT_INTERFACE,
        .bInterfaceNumber = 0,
        .bAlternateSetting = 0,
        .bNumEndpoints = 2,
        .bInterfaceClass = USB_CLASS_MSC,
        .bInterfaceSubClass = USB_MSC_SUBCLASS_SCSI,
        .bInterfaceProtocol = USB_MSC_PROTOCOL_BBB,
        .iInterface = 0
    },

    .msc_endp = {{
        .bLength = USB_DT_ENDPOINT_SIZE,
        .bDescriptorType = USB_DT_ENDPOINT,
        .bEndpointAddress = 0x01,
        .bmAttributes = USB_ENDPOINT_ATTR_BULK,
        .wMaxPacketSize = 64,
        .bInterval = 0,
    }, {
        .bLength = USB_DT_ENDPOINT_SIZE,
        .bDescriptorType = USB_DT_ENDPOINT,
        .bEndpointAddress = 0x82,
        .bmAttributes = USB_ENDPOINT_ATTR_BULK,
        .wMaxPacketSize = 64,
        .bInterval = 0,
    }}

};

static usbd_msc *ms;
    
static void msc_ep_set_config(usbd_device *usbd_dev,
        const struct usb_config_descriptor *cfg)
{
    (void)cfg;

    usbd_ep_prepare(usbd_dev, 0x01, USBD_EP_BULK, 64, USBD_INTERVAL_NA, USBD_EP_DOUBLE_BUFFER);
    usbd_ep_prepare(usbd_dev, 0x82, USBD_EP_BULK, 64, USBD_INTERVAL_NA, USBD_EP_DOUBLE_BUFFER);
    usbd_msc_start(ms);

}

static const struct usbd_info_string msc_string = {
    .lang_list = &string_lang_list,
    .count = 3,
    .data = string_data
};

static const struct usbd_info msc_info = {
    .device = {
        .desc = &msc_dev_desc,
        .string = &msc_string
    },

    .config = {{
        .desc = (const struct usb_config_descriptor *) &msc_config,
        .string = &msc_string
    }}

};

static void setup_callback(usbd_device *usbd_dev, uint8_t ep_addr,
			const struct usb_setup_data *setup_data)
{
	(void) ep_addr; /* assuming ep_addr == 0 */

	if (!usbd_msc_setup_ep0(ms, setup_data)) {
		usbd_ep0_setup(usbd_dev, setup_data);
	}
}

static int usb_enabled = 0;

void usb_init(void)
{
    if (usb_enabled == 0) {
        rcc_periph_clock_enable(RCC_GPIOA);
        rcc_periph_clock_enable(RCC_OTGFS);
        gpio_mode_setup(GPIOA, GPIO_MODE_AF, GPIO_PUPD_NONE,
                GPIO9 | GPIO11 | GPIO12);
        gpio_set_af(GPIOA, GPIO_AF10, GPIO9 | GPIO11 | GPIO12);
        nvic_enable_irq(NVIC_OTG_FS_IRQ);
        flash_backend.block_count = sdcard_getsize(NULL);
        usbd_dev = usbd_init(msc_target_usb_driver(), NULL, &msc_info);
        ms = usbd_msc_init(usbd_dev, 0x82, 64, 0x01, 64, &flash_backend);
        usbd_register_set_config_callback(usbd_dev, msc_ep_set_config);
        usbd_register_setup_callback(usbd_dev, setup_callback);
        usb_enabled = 1;
    }
}

void usb_deinit(void)
{
    if (usb_enabled) {
        rcc_periph_clock_disable(RCC_OTGFS);
        nvic_disable_irq(NVIC_OTG_FS_IRQ);
        gpio_mode_setup(GPIOA, GPIO_MODE_INPUT, GPIO_PUPD_NONE,
                GPIO9 | GPIO11 | GPIO12);
        usb_enabled = 0;
    }
}

static uint32_t last_call = 0;
void usb_poll(void)
{
    bridge_lock();
    if (usb_enabled && usbd_dev) {
        if ((jiffies - last_call > 4)) {
            usbd_poll(usbd_dev, 0);
            last_call = jiffies;
        }
    }
    bridge_unlock();
}

int bridge_init(void)
{

    int ret = 0;
    void *sdcard = sdcard_open();
    if (!sdcard)
        return -1;
    usb_init();
    return 0;
}

void otg_fs_isr(void)
{
    last_call = jiffies;
    usbd_poll(usbd_dev, 0);
}