#include <stdint.h>
#include <stdio.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/mman.h>
#include <arpa/inet.h>
#include <string.h>
#include <fcntl.h>

/*
 * TPM2.net handling
 */
#define TPM2_NET_PORT 65506
#define TPM2_NET_START_BYTE 0x9c
#define TPM2_START_BYTE 0xc9
#define TPM2_END_BYTE 0x36

enum tpm2_packet_type {
    DATA_FRAME = 0xda,
    COMMAND = 0xc0,
    RESPONSE = 0xaa,
};

struct tpm2_packet {
    uint8_t start_byte;
    uint8_t packet_type;
    uint16_t len;
    uint8_t packet_id;
    uint8_t packet_count;
    uint8_t data[];
};

#define BUF_SIZE 0xffff
uint8_t buf[BUF_SIZE];

/*
 * Q3KLED handling
 */
#define Q3KLED_ADDRESS 0x7aa00000
#define Q3KLED_SIZE 0x10000
#define Q3KLED_WIDTH 128
#define Q3KLED_HEIGHT 128

#define FRAMEBUFFER_SIZE (Q3KLED_WIDTH*Q3KLED_HEIGHT*3)
uint8_t framebuffer[FRAMEBUFFER_SIZE];
uint8_t framebuffer_cnt;

#ifdef TEST // Use this flag to disable /dev/mem mmapping
uint32_t q3kled_fb_test[Q3KLED_SIZE];
volatile uint32_t* q3kled_fb = q3kled_fb_test;
#else
volatile uint32_t* q3kled_fb;
#endif

int main(int argc, char *argv[])
{
    int sock;
    struct sockaddr_in addr;

#ifndef TEST
    int memfd;

    if((memfd = open("/dev/mem", O_RDWR | O_SYNC)) < 0) {
        perror("/dev/mem open failed");
        return 1;
    }

    q3kled_fb = (uint32_t *)(mmap(0, Q3KLED_SIZE, PROT_READ|PROT_WRITE, MAP_SHARED, memfd, Q3KLED_ADDRESS));
    if(q3kled_fb == NULL) {
        perror("framebuffer mmap failed");
        return 1;
    }
#endif

    if((sock = socket(AF_INET, SOCK_DGRAM, 0)) < 0) {
        perror("socket failed");
        return 1;
    }

    memset((uint8_t*) &addr, 0, sizeof(struct sockaddr_in));
    addr.sin_family = AF_INET;
    addr.sin_port = htons(TPM2_NET_PORT);
    addr.sin_addr.s_addr = htonl(INADDR_ANY);

    if(bind(sock, (struct sockaddr*) &addr, sizeof(addr)) < 0) {
        perror("bind failed");
        return 1;
    }

    printf("Listening on port %d...\n", TPM2_NET_PORT);

    while(1) {
        int recvlen = recvfrom(sock, buf, BUF_SIZE, 0, 0, 0);
        printf("received %d bytes\n", recvlen);

        if(recvlen < sizeof(struct tpm2_packet)) {
            printf("frame too short\n");
            continue;
        }

        struct tpm2_packet* header = (struct tpm2_packet*) buf;
        header->len = htons(header->len);
        if(header->start_byte == TPM2_NET_START_BYTE) {
            printf("TPM2.net packet\n");
            if(header->len + sizeof(struct tpm2_packet) + 1 > recvlen) {
                printf("invalid size\n");
                continue;
            }

            if(header->len + sizeof(struct tpm2_packet) + 1 < recvlen) {
                printf("too much data? ignoring the rest\n");
            }

            if(buf[header->len + sizeof(struct tpm2_packet)] != TPM2_END_BYTE) {
                printf("invalid end byte\n");
                continue;
            }

            if(header->packet_type == DATA_FRAME) {
                printf("Data frame, %d / %d\n", header->packet_id, header->packet_count);

                if(header->packet_id == 0) {
                    memset(framebuffer, 0, sizeof(framebuffer));
                    framebuffer_cnt = 0;
                }

                uint32_t bytes_left = header->len > FRAMEBUFFER_SIZE-framebuffer_cnt ? FRAMEBUFFER_SIZE-framebuffer_cnt : header->len;
                memcpy(framebuffer + framebuffer_cnt, header->data, bytes_left);
                framebuffer_cnt += bytes_left;

                if(header->packet_id == header->packet_count - 1) {
                    printf("end of transmission, blitting\n");
                    for(uint32_t b = 0; b < FRAMEBUFFER_SIZE/3; b++) {
                        q3kled_fb[b] = framebuffer[3*b] << 16 | framebuffer[3*b+1] << 8 | framebuffer[3*b+2];
                    }
                }
            }
        } else if(header->start_byte == TPM2_START_BYTE) {
            printf("TPM2 packet (unsupported)\n");
        }
    }

    return 0;
}