#include <app/zedmon/usb.h>

#include <assert.h>
#include <app/zedmon/ina.h>
#include <app/zedmon/usb_proto.h>
#include <lk/debug.h>
#include <dev/udc.h>
#include <dev/usb.h>
#include <dev/usbc.h>
#include <lk/err.h>
#include <kernel/mutex.h>
#include <kernel/port.h>
#include <kernel/thread.h>
#include <lib/cbuf.h>
#include <lib/console.h>
#include <platform.h>
#include <stdint.h>
#include <string.h>
#include <target/parameters.h>
#include <lk/trace.h>

#define DATA_IN_EP_ADDR_OFFSET (0x0B)
#define DATA_OUT_EP_ADDR_OFFSET (0x12)

#define W(w) (w & 0xff), (w >> 8)

static uint8_t zedmon_if_descriptor[] = {
    0x09,           /* length */
    INTERFACE,      /* type */
    0x00,           /* interface num */
    0x00,           /* alternates */
    0x02,           /* endpoint count */
    0xff,           /* interface class */
    0xff,           /* interface subclass */
    0x00,           /* interface protocol */
    0x00,           /* string index */

    /* endpoint 1 IN */
    0x07,           /* length */
    ENDPOINT,       /* type */
    0x83,           /* address: 1 IN */
    0x02,           /* type: bulk */
    W(64),          /* max packet size: 64 */
    00,             /* interval */

    /* endpoint 1 OUT */
    0x07,           /* length */
    ENDPOINT,       /* type */
    0x03,           /* address: 1 OUT */
    0x02,           /* type: bulk */
    W(64),          /* max packet size: 64 */
    00,             /* interval */
};

static zedmon_usb_report_format_t zedmon_report_formats[] = {
    {
        .packet_type = ZEDMON_USB_PACKET_REPORT_FORMAT,
        .index = 0,
        .type = ZEDMON_USB_TYPE_INT16,
        .unit = ZEDMON_USB_UNIT_VOLTS,
        .scale = 2.5e-6f,
        .name = "v_shunt",
    },
    {
        .packet_type = ZEDMON_USB_PACKET_REPORT_FORMAT,
        .index = 1,
        .type = ZEDMON_USB_TYPE_INT16,
        .unit = ZEDMON_USB_UNIT_VOLTS,
        .scale = 1.25e-3f,
        .name = "v_bus",
    },
    {
        .packet_type = ZEDMON_USB_PACKET_REPORT_FORMAT,
        .index = 0xff,
        .type = 0,
        .unit = 0,
        .scale =0.0f,
        .name = "",
    },
};

static zedmon_usb_parameter_value_t zedmon_parameter_values[] = {
    {
        .packet_type = ZEDMON_USB_PACKET_PARAMETER_VALUE,
        .name = "shunt_resistance",
        .data_type = ZEDMON_USB_TYPE_FLOAT32,
        .data.float_value = SHUNT_RESISTANCE,
    },
    {
        .packet_type = ZEDMON_USB_PACKET_PARAMETER_VALUE,
        .name = "",
        .data_type = ZEDMON_USB_TYPE_UINT8,
        .data.bytes = {0},
    }
};

typedef struct __attribute__((packed)) {
    lk_bigtime_t timestamp;
    uint16_t     v_shunt;
    uint16_t     v_bus;
} zedmon_sample_t;

static_assert(sizeof(zedmon_sample_t) == 12);
static_assert(offsetof(zedmon_sample_t, timestamp) == 0);
static_assert(offsetof(zedmon_sample_t, v_shunt) == 8);
static_assert(offsetof(zedmon_sample_t, v_bus) == 10);

typedef struct __attribute__((packed)) {
    uint8_t         packet_type;
    zedmon_sample_t samples[5];
} zedmon_report_packet_t;
static_assert(sizeof(zedmon_report_packet_t) <= 64);

typedef struct __attribute__((packed)) {
    uint8_t packet_type;
    uint64_t timestamp;
} zedmon_timestamp_packet_t;

static cbuf_t zedmon_sample_cbuf;

static int zedmon_ep;

typedef enum {
    ZEDMON_PORT_CTX_USB_RX     = 0x0,
    ZEDMON_PORT_CTX_USB_TX     = 0x1,
    ZEDMON_PORT_CTX_USB_SAMPLE = 0x2,
} zedmon_port_context_t;

static port_t zedmon_usb_rx_port;
static port_t zedmon_usb_tx_port;
static port_t zedmon_usb_sample_port;
static port_t zedmon_usb_port_group;

static thread_t *zedmon_usb_thread;
static bool zedmon_tx_busy;

static bool zedmon_usb_reporting_enabled;
static zedmon_report_packet_t zedmon_usb_report_packet;
static zedmon_timestamp_packet_t zedmon_usb_timestamp_packet;


static status_t ep_cb_rx(ep_t endpoint, usbc_transfer_t *t);
static status_t ep_cb_tx(ep_t endpoint, usbc_transfer_t *t);

static void queue_rx(void)
{
    static usbc_transfer_t transfer;
    static uint8_t buf[512];

    transfer.callback = &ep_cb_rx;
    transfer.result = 0;
    transfer.buf = &buf;
    transfer.buflen = sizeof(buf);
    transfer.bufpos = 0;
    transfer.extra = 0;

    usbc_queue_rx(zedmon_ep, &transfer);
}

static status_t queue_tx(void *data, size_t len)
{
    if (zedmon_tx_busy) {
        return ERR_BUSY;
    }

    static usbc_transfer_t transfer;
    transfer.callback = &ep_cb_tx;
    transfer.result = 0;
    transfer.buf = data;
    transfer.buflen = len;
    transfer.bufpos = 0;
    transfer.extra = 0;

    zedmon_tx_busy = true;
    usbc_queue_tx(zedmon_ep, &transfer);
    return NO_ERROR;
}

static status_t ep_cb_rx(ep_t endpoint, usbc_transfer_t *t)
{
#if LOCAL_TRACE
    LTRACEF("ep %u transfer %p\n", endpoint, t);
    usbc_dump_transfer(t);

    if (t->result >= 0) {
        hexdump8(t->buf, t->bufpos);
    }
#endif

    //printf("ep %u transfer %p\n", endpoint, t);

    if (t->result >= 0) {
        if (t->bufpos >= 1) {
            port_packet_t packet;
            int command_length =
                t->bufpos < sizeof(packet.value) ? t->bufpos : sizeof(packet.value);
            memcpy(&packet.value, t->buf, command_length);
            port_write(zedmon_usb_rx_port, &packet, 1);
        }
        queue_rx();
    } else {
        printf("bad rx\n");
    }


    return NO_ERROR;
}

static status_t ep_cb_tx(ep_t endpoint, usbc_transfer_t *t)
{
    //printf("ep %u transfer %p\n", endpoint, t);
#if LOCAL_TRACE
    LTRACEF("ep %u transfer %p\n", endpoint, t);
    usbc_dump_transfer(t);
#endif

    if (t->result >= 0) {
        port_packet_t packet;
        packet.value[0] = 0x0;
        port_write(zedmon_usb_tx_port, &packet, 1);
    } else {
        printf("bad tx\n");
    }
    return NO_ERROR;
}


static status_t zedmon_usb_callback(void *cookie, usb_callback_op_t op,
                                    const union usb_callback_args *args)
{
#if LOCAL_TRACE
    LTRACEF("cookie %p, op %u, args %p\n", cookie, op, args);
#endif

    if (op == USB_CB_ONLINE) {
        usbc_setup_endpoint(zedmon_ep, USB_IN, 0x40, USB_BULK);
        usbc_setup_endpoint(zedmon_ep, USB_OUT, 0x40, USB_BULK);

        queue_rx();
    } else if (op == USB_CB_SETUP_MSG) {
        zedmon_usb_reporting_enabled = false;
    }
    return NO_ERROR;
}


static void zedmon_usb_handle_query_report_format_packet(const uint8_t values[7]) {
    uint8_t index = values[0];

    if (index >= countof(zedmon_report_formats)) {
        index = countof(zedmon_report_formats) - 1;
    }

    queue_tx(&zedmon_report_formats[index], sizeof(zedmon_report_formats[index]));
}

static void zedmon_usb_handle_query_parameter_value_packet(const uint8_t values[7]) {
    uint8_t index = values[0];

    if (index >= countof(zedmon_parameter_values)) {
        index = countof(zedmon_parameter_values) - 1;
    }

    queue_tx(&zedmon_parameter_values[index], sizeof(zedmon_parameter_values[index]));
}

static void zedmon_usb_handle_set_output_packet(const uint8_t values[7]) {
    zedmon_usb_set_target_out(values[0], values[1]);
}

static void zedmon_usb_handle_usb_rx_port(const port_packet_t *packet) {
    zedmon_usb_packet_type_t type = (zedmon_usb_packet_type_t)packet->value[0];

    switch (type) {
        case ZEDMON_USB_PACKET_QUERY_REPORT_FORMAT:
            zedmon_usb_handle_query_report_format_packet((const uint8_t *)packet->value + 1);
            break;
        case ZEDMON_USB_PACKET_QUERY_TIME:
            zedmon_usb_timestamp_packet.packet_type = ZEDMON_USB_PACKET_TIMESTAMP;
            zedmon_usb_timestamp_packet.timestamp = ina_get_current_time();
            queue_tx(&zedmon_usb_timestamp_packet, sizeof(zedmon_usb_timestamp_packet));
            break;
        case ZEDMON_USB_PACKET_QUERY_PARAMETER_VALUE:
            zedmon_usb_handle_query_parameter_value_packet((const uint8_t *)packet->value + 1);
            break;
        case ZEDMON_USB_PACKET_ENABLE_REPORTING:
            zedmon_usb_reporting_enabled = true;
            break;
        case ZEDMON_USB_PACKET_DISABLE_REPORTING:
            zedmon_usb_reporting_enabled = false;
            break;
        case ZEDMON_USB_PACKET_SET_OUTPUT:
            zedmon_usb_handle_set_output_packet((const uint8_t *)packet->value + 1);
            break;

            // These are Device -> Host packets which we should never see.
        case ZEDMON_USB_PACKET_REPORT_FORMAT:
            break;
        case ZEDMON_USB_PACKET_REPORT:
            break;
        case ZEDMON_USB_PACKET_TIMESTAMP:
            break;
        case ZEDMON_USB_PACKET_PARAMETER_VALUE:
            break;
    }
}

static void zedmon_usb_check_for_repot_tx(void) {
    if (zedmon_tx_busy) {
        return;
    }
    zedmon_report_packet_t *packet = &zedmon_usb_report_packet;
    if (cbuf_space_used(&zedmon_sample_cbuf) >= sizeof(packet->samples)) {
        size_t read_len =
            cbuf_read(&zedmon_sample_cbuf, &packet->samples, sizeof(packet->samples), false);
        assert(read_len == sizeof(packet->samples));

        if (zedmon_usb_reporting_enabled) {
            queue_tx(packet, sizeof(*packet));
        }
    }
}

static int zedmon_usb_thread_entry(void *arg) {
    while (true) {
        port_result_t result;
        port_read(zedmon_usb_port_group, INFINITE_TIME, &result);
        zedmon_port_context_t context = (zedmon_port_context_t)result.ctx;
        switch(context) {
            case ZEDMON_PORT_CTX_USB_RX:
                zedmon_usb_handle_usb_rx_port(&result.packet);
                break;

            case ZEDMON_PORT_CTX_USB_TX:
                zedmon_tx_busy = false;
                zedmon_usb_check_for_repot_tx();
                break;

            case ZEDMON_PORT_CTX_USB_SAMPLE:
                zedmon_usb_check_for_repot_tx();
                break;

        }
    }
    return 0;

}

void zedmon_usb_add_sample(lk_bigtime_t timestamp, uint16_t v_shunt, uint16_t v_bus) {
    zedmon_sample_t sample = {
        .timestamp = timestamp,
        .v_shunt   = v_shunt,
        .v_bus     = v_bus,
    };

    // If there's no space, drop the sample.
    if (cbuf_space_avail(&zedmon_sample_cbuf) < sizeof(sample)) {
        return;
    }
    size_t write_len = cbuf_write(&zedmon_sample_cbuf, &sample, sizeof(sample), false);
    assert(write_len == sizeof(sample));

    port_packet_t packet;
    packet.value[0] = 0x0;
    port_write(zedmon_usb_sample_port, &packet, 1);
}

void __WEAK zedmon_usb_set_target_out(int index, bool value) {
}

void zedmon_usb_init(int data_ep_addr) {
    cbuf_initialize(&zedmon_sample_cbuf, 128);
    zedmon_usb_report_packet.packet_type = ZEDMON_USB_PACKET_REPORT;

    status_t ret;

    ret = port_create("zm_usb_rx", PORT_MODE_UNICAST, &zedmon_usb_rx_port);
    if (ret != NO_ERROR) {
        printf("error creating zm_usb_rx: %d\n", ret);
    }

    port_t usb_rx_reader;
    ret = port_open("zm_usb_rx", (void *)ZEDMON_PORT_CTX_USB_RX, &usb_rx_reader);
    if (ret != NO_ERROR) {
        printf("error opening zm_usb_rx: %d\n", ret);
    }

    ret = port_create("zm_usb_tx", PORT_MODE_UNICAST, &zedmon_usb_tx_port);
    if (ret != NO_ERROR) {
        printf("error creating zm_usb_rx: %d\n", ret);
    }

    port_t usb_tx_reader;
    ret = port_open("zm_usb_tx", (void *)ZEDMON_PORT_CTX_USB_TX, &usb_tx_reader);
    if (ret != NO_ERROR) {
        printf("error opening zm_usb_tx: %d\n", ret);
    }

    ret = port_create("zm_usb_sam", PORT_MODE_UNICAST, &zedmon_usb_sample_port);
    if (ret != NO_ERROR) {
        printf("error creating zm_usb_sample: %d\n", ret);
    }

    port_t usb_sample_reader;
    ret = port_open("zm_usb_sam", (void *)ZEDMON_PORT_CTX_USB_SAMPLE, &usb_sample_reader);
    if (ret != NO_ERROR) {
        printf("error opening zm_usb_sample: %d\n", ret);
    }

    port_t ports[] = {usb_rx_reader, usb_tx_reader, usb_sample_reader};

    ret = port_group(ports, countof(ports), &zedmon_usb_port_group);

    zedmon_ep = data_ep_addr;

    zedmon_if_descriptor[DATA_IN_EP_ADDR_OFFSET]  = data_ep_addr | 0x80;
    zedmon_if_descriptor[DATA_OUT_EP_ADDR_OFFSET] = data_ep_addr;

    usb_append_interface_lowspeed(zedmon_if_descriptor, sizeof(zedmon_if_descriptor));
    usb_append_interface_highspeed(zedmon_if_descriptor, sizeof(zedmon_if_descriptor));

    usb_register_callback(&zedmon_usb_callback, NULL);

    zedmon_tx_busy = false;
    zedmon_usb_thread = thread_create("zedmon_usb", &zedmon_usb_thread_entry, 0x0,
                                      DEFAULT_PRIORITY, DEFAULT_STACK_SIZE);
    thread_detach(zedmon_usb_thread);
    thread_resume(zedmon_usb_thread);
}

#include "stm32f0xx_hal.h"
static void print_epr(int i, uint16_t epr) {
    printf("ep%dr: ea:%d stat_tx:%d dtog_tx:%d ctr_tx:%d ep_kind:%d ep_type:%d\n"
           "       setup:%d stat_rx:%d dtog_rx:%d ctr_rx:%d\n",
           i,
           epr & 0xf,
           (epr >> 4) & 0x3,
           (epr >> 6) & 0x1,
           (epr >> 7) & 0x1,
           (epr >> 8) & 0x1,
           (epr >> 9) & 0x3,
           (epr >> 11) & 0x1,
           (epr >> 12) & 0x3,
           (epr >> 14) & 0x1,
           (epr >> 15) & 0x1);
}

static int cmd_usb_dump(int argc, const cmd_args *argv) {
    print_epr(0, USB->EP0R);
    print_epr(1, USB->EP1R);
    print_epr(2, USB->EP2R);
    print_epr(3, USB->EP3R);
    print_epr(4, USB->EP4R);
    print_epr(5, USB->EP5R);
    print_epr(6, USB->EP6R);
    print_epr(7, USB->EP7R);

    return NO_ERROR;
}

STATIC_COMMAND_START
STATIC_COMMAND("usb_dump", "usb dump", &cmd_usb_dump)
STATIC_COMMAND_END(usb);