//
// Created by beerpsi on 12/30/2023.
//

#include "chuniio.h"

#include <stdio.h>
#include <inttypes.h>
#include <signal.h>
#include <time.h>
#include <process.h>
#include <assert.h>

#include "config.h"
#include "struct.h"
#include "util/dprintf.h"

//region Brokenithm
struct IPCMemoryInfo* chuni_io_file_mapping;

const char *memFileName = "Local\\BROKENITHM_SHARED_BUFFER";

char remote_address[BUFSIZ];
uint16_t remote_port = 52468;
uint16_t server_port = 52468;
bool tcp_mode = false;

volatile sig_atomic_t EXIT_FLAG = false, CONNECTED = false;

uint32_t last_input_packet_id = 0;

enum {
    CARD_AIME,
    CARD_FELICA,
};

enum {
    FUNCTION_COIN = 1,
    FUNCTION_CARD
};

typedef struct {
    SOCKET sHost;
    struct IPCMemoryInfo* memory;
} thread_args;

void socket_set_timeout(SOCKET sHost, int timeout) {
    setsockopt(sHost, SOL_SOCKET, SO_SNDTIMEO, (char*)&timeout, sizeof(int));
    setsockopt(sHost, SOL_SOCKET, SO_RCVTIMEO, (char*)&timeout, sizeof(int));
}

int socket_bind(SOCKET sHost, unsigned long addr, uint16_t port) {
    struct sockaddr_in srcaddr = {};
    memset(&srcaddr, 0, sizeof(srcaddr));
    srcaddr.sin_family = AF_INET;
    srcaddr.sin_addr.s_addr = addr;
    srcaddr.sin_port = htons(port);
    return bind(sHost, (struct sockaddr*)&srcaddr, sizeof(srcaddr));
}

int socket_send_to(SOCKET sHost, const struct sockaddr_in* addr, const char* buf, int len) {
    return sendto(sHost, buf, len, 0, (struct sockaddr*)&addr, sizeof(&addr));
}

void print_err(const char* fmt, ...) {
    time_t lt = time(NULL);
    struct tm *local = localtime(&lt);
    char tmpbuf[32];

    strftime(tmpbuf, 32, "%Y-%m-%d %H:%M:%S", local);

    dprintf("brokenithm_server: [%s] ", tmpbuf);

    va_list ap;
    va_start(ap, fmt);
    dprintfv(fmt, ap);
    va_end(ap);
}

void get_socks_address(const struct PacketConnect* pkt, char* address, int address_len, uint16_t *port) {
    if (!pkt || !address || !port) {
        return;
    }

    *port = ntohs(pkt->port);

    switch (pkt->addrType) {
        case 1:
            inet_ntop(AF_INET, pkt->addr.addr4.addr, address, address_len);
            break;
        case 2:
            inet_ntop(AF_INET6, pkt->addr.addr6, address, address_len);
            break;
        default:
            return;
    }
}

void update_packet_id(uint32_t new_packet_id)
{
    if (last_input_packet_id > new_packet_id) {
        print_err("[WARN] Packet #%" PRIu32 " came too late\n", new_packet_id);
    } else if (new_packet_id > last_input_packet_id + 1) {
        print_err("[WARN] Packets between #%" PRIu32 " and #%" PRIu32 " total %" PRIu32 " packet(s) are missing, probably too late or dropped\n",
                  last_input_packet_id, new_packet_id, new_packet_id - last_input_packet_id - 1);
    } else if (new_packet_id == last_input_packet_id) {
        print_err("[WARN] Packet #%" PRIu32 " duplicated\n", new_packet_id);
    }
    last_input_packet_id = new_packet_id;
}

void dump_bytes(const void *ptr, size_t nbytes, bool hex_string)
{
    const uint8_t *bytes;
    uint8_t c;
    size_t i;
    size_t j;

    if (nbytes == 0) {
        dprintf("\t--- Empty ---\n");
    }

    bytes = (const unsigned char*)ptr;

    if (hex_string) {
        for (i = 0 ; i < nbytes ; i++) {
            dprintf("%02x", bytes[i]);
        }
        dprintf("\n");
        return;
    }

    for (i = 0 ; i < nbytes ; i += 16) {
        dprintf("    %08x:", (int) i);

        for (j = 0 ; i + j < nbytes && j < 16 ; j++) {
            dprintf(" %02x", bytes[i + j]);
        }

        while (j < 16) {
            dprintf("   ");
            j++;
        }

        dprintf(" ");

        for (j = 0 ; i + j < nbytes && j < 16 ; j++) {
            c = bytes[i + j];

            if (c < 0x20 || c >= 0x7F) {
                c = '.';
            }

            dprintf("%c", c);
        }

        dprintf("\n");
    }

    dprintf("\n");
}

void print_card_info(uint8_t card_type, uint8_t *card_id) {
    switch (card_type) {
        case CARD_AIME:
            print_err("[INFO] Card type: AiMe, ID: ");
            dump_bytes(card_id, 10, true);
            break;
        case CARD_FELICA:
            print_err("[INFO] Card type: FeliCa, ID: ");
            dump_bytes(card_id, 8, true);
            break;
        default:
            break;
    }
}

int make_ipv4_address(struct sockaddr_in* addr, char* host, uint16_t port) {
    addr->sin_family = AF_INET;
    addr->sin_port = htons(port);
    return inet_pton(AF_INET, host, (struct in_addr *)&addr->sin_addr.s_addr);
}

uint8_t previous_led_status[3 * 32];
bool has_previous_led_status = false;
int skip_count = 0;

unsigned int __stdcall thread_led_broadcast(void *ctx) {
    thread_args *args = (thread_args*)ctx;

    SOCKET sHost = args->sHost;
    struct IPCMemoryInfo* memory = args->memory;

    struct sockaddr_in addr = {};
    make_ipv4_address(&addr, remote_address, remote_port);

    char send_buffer[4 + 3 * 32];
    send_buffer[0] = 100;
    send_buffer[1] = 'L';
    send_buffer[2] = 'E';
    send_buffer[3] = 'D';

    uint8_t current_led_status[3 * 32];

    while (!EXIT_FLAG) {
        if (!CONNECTED) {
            Sleep(50);
            continue;
        }

        memcpy(current_led_status, memory->ledRgbData, 3 * 32);

        bool same;

        if (!has_previous_led_status) {
            same = (memcmp(previous_led_status, current_led_status, 3 * 32) == 0);
        } else {
            same = false;
        }

        memcpy(previous_led_status, current_led_status, 3 * 32);
        has_previous_led_status = true;

        if (!same || ++skip_count > 50) {
            memcpy(send_buffer + 4, current_led_status, 3 * 32);

            if (socket_send_to(sHost, &addr, send_buffer, 100) < 0) {
                print_err("[ERROR] Cannot send packet: error %lu\n", GetLastError());

                if (tcp_mode) {
                    if (errno == EINTR || errno == EWOULDBLOCK || errno == EAGAIN) {
                        continue;
                    } else {
                        print_err("[INFO] Device disconnected!\n");
                        CONNECTED = false;
                        EXIT_FLAG = true;
                        break;
                    }
                }
            }

            skip_count = 0;
        }

        Sleep(10);
    }

    return 0;
}

uint8_t last_card_id[10];

unsigned int __stdcall thread_input_recv(void *ctx) {
    thread_args *args = (thread_args*)ctx;

    SOCKET sHost = args->sHost;
    struct IPCMemoryInfo* memory = args->memory;

    char buffer[BUFSIZ];

    struct sockaddr_in addr = {};
    make_ipv4_address(&addr, remote_address, remote_port);

    int recv_len, packet_len;
    uint8_t real_len;

    while (!EXIT_FLAG) {
        if (!tcp_mode) {
            /**
                on UDP mode data is sent as packets, so just receive into a buffer big enough for 1 packet
                each recvfrom call will only get 1 packet of data, the remaining data is discarded
            **/

            if ((recv_len = recvfrom(sHost, buffer, BUFSIZ - 1, 0, NULL, NULL)) == -1) {
                continue;
            }

            real_len = (unsigned char)buffer[0];

            if (real_len > recv_len) {
                continue;
            }

            packet_len = real_len + 1;
        } else {
            /**
                on TCP mode packets are length-prefixed, so we read in the first 4 bytes to
                figure out how much we need to read, then read in the full data.
            **/
            recv_len = 0;

            while (recv_len < 4) {
                int read = recv(sHost, buffer + recv_len, 4 - recv_len, 0);

                if (read == -1) {
                    if (errno == EINTR || errno == EWOULDBLOCK || errno == EAGAIN) {
                        continue;
                    } else {
                        print_err("[INFO] Device disconnected!\n");
                        CONNECTED = false;
                        EXIT_FLAG = true;
                        break;
                    }
                }

                recv_len = recv_len + read;
            }

            real_len = buffer[0];
            packet_len = real_len + 1;

            while (recv_len < packet_len) {
                int read = recv(sHost, buffer + recv_len, packet_len - recv_len, 0);

                if (read == -1) {
                    if (errno == EINTR || errno == EWOULDBLOCK || errno == EAGAIN) {
                        continue;
                    } else {
                        print_err("[INFO] Device disconnected!\n");
                        CONNECTED = false;
                        EXIT_FLAG = true;
                        break;
                    }
                }

                recv_len = recv_len + read;
            }
        }

        if (packet_len >= sizeof(struct PacketInput) && buffer[1] == 'I' && buffer[2] == 'N' && buffer[3] == 'P') {
            struct PacketInput* pkt = (struct PacketInput*)buffer;

            memcpy(memory->airIoStatus, pkt->airIoStatus, sizeof(pkt->airIoStatus));
            memcpy(memory->sliderIoStatus, pkt->sliderIoStatus, sizeof(pkt->sliderIoStatus));
            memory->testBtn = pkt->testBtn;
            memory->serviceBtn = pkt->serviceBtn;

            update_packet_id(ntohl(pkt->packetId));
        } else if (packet_len >= sizeof(struct PacketInputNoAir) && buffer[1] == 'I' && buffer[2] == 'P' && buffer[3] == 'T') {  // without air
            struct PacketInputNoAir* pkt = (struct PacketInputNoAir*)buffer;

            memcpy(memory->sliderIoStatus, pkt->sliderIoStatus, sizeof(pkt->sliderIoStatus));
            memory->testBtn = pkt->testBtn;
            memory->serviceBtn = pkt->serviceBtn;

            update_packet_id(ntohl(pkt->packetId));
        } else if (packet_len >= sizeof(struct PacketFunction) && buffer[1] == 'F' && buffer[2] == 'N' && buffer[3] == 'C') {
            struct PacketFunction* pkt = (struct PacketFunction*)buffer;

            switch (pkt->funcBtn) {
                case FUNCTION_COIN:
                    memory->coinInsertion = 1;
                    break;
                case FUNCTION_CARD:
                    memory->cardRead = 1;
                    break;
            }
        } else if (packet_len >= sizeof(struct PacketConnect) && buffer[1] == 'C' && buffer[2] == 'O' && buffer[3] == 'N') {
            struct PacketConnect* pkt = (struct PacketConnect*)buffer;

            get_socks_address(pkt, remote_address, BUFSIZ - 1, &remote_port);
            print_err("[INFO] Device %s:%d connected.\n", remote_address, remote_port);

            last_input_packet_id = 0;
            CONNECTED = true;
        } else if (packet_len >= 4 && buffer[1] == 'D' && buffer[2] == 'I' && buffer[3] == 'S') {
            CONNECTED = false;

            if (tcp_mode) {
                EXIT_FLAG = 1;
                print_err("[INFO] Device disconnected!\n");
                break;
            }

            if (strlen(remote_address)) {
                print_err("[INFO] Device %s:%d disconnected.\n", remote_address, remote_port);
                memset(remote_address, 0, BUFSIZ);
            }
        } else if (packet_len >= sizeof(struct PacketPing) && buffer[1] == 'P' && buffer[2] == 'I' && buffer[3] == 'N') {
            if (!CONNECTED) {
                continue;
            }

            char response[13];
            memcpy(response, buffer, 12);
            response[2] = 'O';

            socket_send_to(sHost, &addr, response, 13);
        } else if (packet_len >= sizeof(struct PacketCard) && buffer[1] == 'C' && buffer[2] == 'R' && buffer[3] == 'D') {
            struct PacketCard* pkt = (struct PacketCard*)buffer;

            if (pkt->remoteCardRead) {
                if (memcmp(last_card_id, pkt->remoteCardId, 10) != 0) {
                    print_err("[INFO] Got remote card.\n");
                    print_card_info(pkt->remoteCardType, pkt->remoteCardId);
                    memcpy(last_card_id, pkt->remoteCardId, 10);
                }
            } else if (memory->remoteCardRead) {
                print_err("[INFO] Remote card removed.\n");
                memset(last_card_id, 0, 10);
            }

            memory->remoteCardRead = pkt->remoteCardRead;
            memory->remoteCardType = pkt->remoteCardType;
            memcpy(memory->remoteCardId, pkt->remoteCardId, 10);
        }
    }

    return 0;
}

unsigned int __stdcall server_thread_proc(void* ctx) {
    struct IPCMemoryInfo* memory = (struct IPCMemoryInfo*)ctx;

    if (!tcp_mode) {
        print_err("[INFO] Mode: UDP\n");

        SOCKET sHost = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
        socket_set_timeout(sHost, 2000);
        socket_bind(sHost, htonl(INADDR_ANY), server_port);

        print_err("[INFO] Waiting for device on port %d...\n", server_port);

        thread_args args = { .memory = memory, .sHost = sHost };

        HANDLE led_thread = (HANDLE)_beginthreadex(NULL, 0, thread_led_broadcast, &args, 0, NULL);
        HANDLE input_thread = (HANDLE)_beginthreadex(NULL, 0, thread_input_recv, &args, 0, NULL);

        WaitForSingleObject(led_thread, INFINITE);
        WaitForSingleObject(input_thread, INFINITE);

        CloseHandle(led_thread);
        CloseHandle(input_thread);
    } else {
        print_err("[INFO] Mode: TCP\n");
        SOCKET sHost = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);

        socket_set_timeout(sHost, 50);
        socket_bind(sHost, htonl(INADDR_ANY), server_port);

        listen(sHost, 10);

#pragma clang diagnostic push
#pragma ide diagnostic ignored "EndlessLoop"
        for (;;) {
            print_err("[INFO] Waiting for device on port %d...\n", server_port);
            struct sockaddr_in user_socket = {};
            socklen_t sock_size = sizeof(struct sockaddr_in);
            SOCKET acc_socket = accept(sHost, (struct sockaddr *)&user_socket, &sock_size);

            char buffer[20] = {};
            const char* user_address = inet_ntop(AF_INET, &user_socket.sin_addr, buffer, 20);
            if (user_address != NULL) {
                print_err("[INFO] Device %s:%d connected.\n", user_address, user_socket.sin_port);
            }

            CONNECTED = true;
            EXIT_FLAG = false;

            thread_args args = { .memory = memory, .sHost = acc_socket };

            HANDLE led_thread = (HANDLE)_beginthreadex(NULL, 0, thread_led_broadcast, &args, 0, NULL);
            HANDLE input_thread = (HANDLE)_beginthreadex(NULL, 0, thread_input_recv, &args, 0, NULL);

            WaitForSingleObject(led_thread, INFINITE);
            WaitForSingleObject(input_thread, INFINITE);

            CloseHandle(led_thread);
            CloseHandle(input_thread);

            print_err("[INFO] Exiting gracefully...\n");
            last_input_packet_id = 0;
            EXIT_FLAG = true;
            CONNECTED = false;
        }
#pragma clang diagnostic pop
    }

    return 0;
}

HRESULT server_start() {
    tcp_mode = GetPrivateProfileIntW(L"brokenithm", L"tcp", 0, L".\\segatools.ini") == 1;
    server_port = GetPrivateProfileIntW(L"brokenithm", L"port", 52468, L".\\segatools.ini");

    struct WSAData wsaData = {};
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
        print_err("[ERROR] WSA startup failed!\n");
        return E_FAIL;
    }

    HANDLE hMapFile = OpenFileMappingA(FILE_MAP_ALL_ACCESS, false, memFileName);

    if (hMapFile == NULL) {
        hMapFile = CreateFileMappingA(INVALID_HANDLE_VALUE, NULL, PAGE_READWRITE, 0, 1024, memFileName);
        if (hMapFile == NULL) {
            print_err("[ERROR] CreateFileMapping failed! error: %lu\n", GetLastError());
            return E_FAIL;
        }
    }

    struct IPCMemoryInfo* memory = (struct IPCMemoryInfo*)MapViewOfFileEx(hMapFile, FILE_MAP_ALL_ACCESS, 0, 0, 1024, NULL);
    chuni_io_file_mapping = memory;

    if (memory == NULL) {
        print_err("[ERROR] Cannot get view of memory map! error: %lu\n", GetLastError());
        return E_FAIL;
    }

    _beginthreadex(NULL, 0, server_thread_proc, memory, 0, NULL);

    return S_OK;
}
//endregion

//region ChuniIO stuff
static unsigned int __stdcall chuni_io_slider_thread_proc(void *ctx);

static bool chuni_io_coin;
static uint16_t chuni_io_coins;
static uint8_t chuni_io_hand_pos;
static HANDLE chuni_io_slider_thread;
static bool chuni_io_slider_stop_flag;
static struct chuni_io_config chuni_io_cfg;

uint16_t chuni_io_get_api_version() {
    return 0x0102;
}

HRESULT chuni_io_jvs_init() {
    chuni_io_config_load(&chuni_io_cfg, L".\\segatools.ini");

    HRESULT result = server_start();
    if (result != S_OK) {
        return result;
    }

    return S_OK;
}

void chuni_io_jvs_read_coin_counter(uint16_t *out) {
    if (out == NULL) {
        return;
    }

    if (chuni_io_file_mapping && chuni_io_file_mapping->coinInsertion) {
        chuni_io_coins++;
        chuni_io_file_mapping->coinInsertion = 0;
    } else {
        if (GetAsyncKeyState(chuni_io_cfg.vk_coin)) {
            if (!chuni_io_coin) {
                chuni_io_coin = true;
                chuni_io_coins++;
            }
        } else {
            chuni_io_coin = false;
        }
    }

    *out = chuni_io_coins;
}

void chuni_io_jvs_poll(uint8_t *opbtn, uint8_t *beams) {
    size_t i;

    if ((chuni_io_file_mapping && chuni_io_file_mapping->testBtn) || GetAsyncKeyState(chuni_io_cfg.vk_test)) {
        *opbtn |= CHUNI_IO_OPBTN_TEST; /* Test */
    }

    if ((chuni_io_file_mapping && chuni_io_file_mapping->serviceBtn) || GetAsyncKeyState(chuni_io_cfg.vk_service)) {
        *opbtn |= CHUNI_IO_OPBTN_SERVICE; /* Service */
    }

    if (GetAsyncKeyState(chuni_io_cfg.vk_ir_emu)) {
        if (chuni_io_hand_pos < 6) {
            chuni_io_hand_pos++;
        }
    } else {
        if (chuni_io_hand_pos > 0) {
            chuni_io_hand_pos--;
        }
    }

    for (i = 0 ; i < 6 ; i++) {
        if (chuni_io_hand_pos > i) {
            *beams |= (1 << i);
        }
    }

    // IR format is beams[5:0] = {b5,b6,b3,b4,b1,b2};
    for (i = 0 ; i < 3 ; i++) {
        if (chuni_io_file_mapping && chuni_io_file_mapping->airIoStatus[i*2])
            *beams |= (1 << (i*2+1));
        if (chuni_io_file_mapping && chuni_io_file_mapping->airIoStatus[i*2+1])
            *beams |= (1 << (i*2));
    }
}

HRESULT chuni_io_slider_init() {
    return S_OK;
}

void chuni_io_slider_start(void* callback) {
    if (chuni_io_slider_thread != NULL) {
        return;
    }

    chuni_io_slider_thread = (HANDLE) _beginthreadex(
            NULL,
            0,
            chuni_io_slider_thread_proc,
            callback,
            0,
            NULL);
}

void chuni_io_slider_stop(void) {
    if (chuni_io_slider_thread == NULL) {
        return;
    }

    chuni_io_slider_stop_flag = true;

    WaitForSingleObject(chuni_io_slider_thread, INFINITE);
    CloseHandle(chuni_io_slider_thread);
    chuni_io_slider_thread = NULL;
    chuni_io_slider_stop_flag = false;
}

void chuni_io_slider_set_leds(const uint8_t *rgb) {
    if (chuni_io_file_mapping) {
        memcpy(chuni_io_file_mapping->ledRgbData, rgb, 32 * 3);
    }
}

HRESULT chuni_io_led_init(void) {
    return S_OK;
}

void chuni_io_led_set_colors(uint8_t board, uint8_t *rgb) {}

static unsigned int __stdcall chuni_io_slider_thread_proc(void *ctx)
{
    chuni_io_slider_callback_t callback;
    uint8_t pressure[32];

    callback = (chuni_io_slider_callback_t)ctx;

#pragma clang diagnostic push
#pragma ide diagnostic ignored "LoopDoesntUseConditionVariableInspection"
    while (!chuni_io_slider_stop_flag) {
        if (chuni_io_file_mapping) {
            memcpy(pressure, chuni_io_file_mapping->sliderIoStatus, 32);
        }

        callback(pressure);
        Sleep(1);
    }
#pragma clang diagnostic pop

    return 0;
}
//endregion

BOOL APIENTRY DllMain(HMODULE hModule, DWORD ul_reason_for_call, LPVOID lpReserved) {
    return TRUE;
}