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

#include "android.h"
#include "arch.h"

#ifdef ENV32BIT
#include <inttypes.h>
#include <process.h>
#include <stdatomic.h>
#include <stdbool.h>
#include <stdint.h>

#include "servers/common.h"
#include "servers/socket.h"

bool tcp_mode = true;
uint16_t server_port = 52468;

enum CardType {
    CARD_AIME,
    CARD_FELICA,
};

typedef struct {
    SOCKET sock;
    char remote_address[40];
    uint16_t remote_port;

    atomic_bool exit_flag;
    atomic_bool connected;

    uint32_t last_input_packet_id;
    uint8_t last_card_id[10];

    bool has_previous_led_status;
    uint8_t previous_led_status[3 * 32];
    uint8_t led_skip_count;

    struct IPCMemoryInfo *memory;
} android_thread_ctx;

void socket_set_timeout(const 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(const SOCKET sHost, const unsigned long addr, const 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(const SOCKET sHost, const struct sockaddr_in *addr, const char *buf,
                   const int len) {
    return sendto(sHost, buf, len, 0, (struct sockaddr *)&addr, sizeof(&addr));
}

int make_ipv4_address(struct sockaddr_in *addr, const char *host, const 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);
}

void get_socks_address(const struct PacketConnect *pkt, char *address,
                       const 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 print_card_info(const uint8_t card_type, const uint8_t *card_id) {
    switch (card_type) {
    case CARD_AIME:
        print_err("[Android: INFO] Card type: AiMe, ID: ");
        dump_bytes(card_id, 10, true);
        break;
    case CARD_FELICA:
        print_err("[Android: INFO] Card type: FeliCa, ID: ");
        dump_bytes(card_id, 8, true);
        break;
    default:
        break;
    }
}

void update_packet_id(android_thread_ctx *ctx, const uint32_t new_packet_id) {
    if (ctx->last_input_packet_id > new_packet_id) {
        print_err("[WARN] Packet #%" PRIu32 " came too late\n", new_packet_id);
    } else if (new_packet_id > ctx->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",
                  ctx->last_input_packet_id, new_packet_id,
                  new_packet_id - ctx->last_input_packet_id - 1);
    } else if (new_packet_id == ctx->last_input_packet_id) {
        print_err("[WARN] Packet #%" PRIu32 " duplicated\n", new_packet_id);
    }
    ctx->last_input_packet_id = new_packet_id;
}

unsigned int __stdcall led_broadcast_thread_proc(void *v) {
    android_thread_ctx *ctx = v;

    const SOCKET sock = ctx->sock;
    const struct IPCMemoryInfo *memory = ctx->memory;

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

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

    while (!atomic_load(&ctx->exit_flag)) {
        uint8_t current_led_status[3 * 32];

        if (!atomic_load(&ctx->connected)) {
            Sleep(50);
            continue;
        }

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

        bool same;

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

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

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

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

                if (tcp_mode) {
                    if (errno == EINTR || errno == EWOULDBLOCK || errno == EAGAIN) {
                        continue;
                    }

                    print_err("[Android: INFO] Device disconnected!\n");
                    ctx->connected = false;
                    ctx->exit_flag = true;
                    break;
                }
            }

            ctx->led_skip_count = 0;
        }

        Sleep(10);
    }

    return 0;
}

unsigned int __stdcall input_recv_thread_proc(void *v) {
    android_thread_ctx *ctx = v;

    const SOCKET sock = ctx->sock;
    struct IPCMemoryInfo *memory = ctx->memory;

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

    int recv_len, packet_len;
    uint8_t real_len;

    while (!atomic_load(&ctx->exit_flag)) {
        char buffer[96];

        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(sock, buffer, 96, 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) {
                const int read = recv(sock, buffer + recv_len, 4 - recv_len, 0);

                if (read == -1) {
                    int error = WSAGetLastError();

                    if (errno == EINTR || errno == EWOULDBLOCK || errno == EAGAIN ||
                        error == WSAETIMEDOUT) {
                        continue;
                    }

                    print_err("[Android: INFO] Device disconnected (could not read "
                              "data, errno "
                              "%d, os error %ld)\n",
                              errno, error);
                    atomic_store(&ctx->connected, false);
                    atomic_store(&ctx->exit_flag, true);
                    break;
                }

                recv_len = recv_len + read;
            }

            real_len = buffer[0];
            packet_len = real_len + 1; // 1 for the packet length

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

                if (read == -1) {
                    int error = WSAGetLastError();

                    if (errno == EINTR || errno == EWOULDBLOCK || errno == EAGAIN ||
                        error == WSAETIMEDOUT) {
                        continue;
                    }

                    print_err("[Android: INFO] Device disconnected (could not read "
                              "data, errno "
                              "%d, os error %ld)\n",
                              errno, error);
                    atomic_store(&ctx->connected, false);
                    atomic_store(&ctx->exit_flag, true);
                    break;
                }

                recv_len = recv_len + read;
            }
        }

        if (packet_len >= sizeof(struct PacketInput) &&
            memcmp(buffer + 1, "INP", 3) == 0) {
            const 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(ctx, ntohl(pkt->packetId));
        } else if (packet_len >= sizeof(struct PacketInputNoAir) &&
                   memcmp(buffer + 1, "IPT", 3) == 0) { // without air
            const 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(ctx, ntohl(pkt->packetId));
        } else if (packet_len >= sizeof(struct PacketFunction) &&
                   memcmp(buffer + 1, "FNC", 3) == 0) {
            const struct PacketFunction *pkt = (struct PacketFunction *)buffer;

            switch (pkt->funcBtn) {
            case FUNCTION_COIN:
                memory->coinInsertion = 1;
                break;
            case FUNCTION_CARD:
                memory->cardRead = 1;
                break;
            default:
                break;
            }
        } else if (packet_len >= sizeof(struct PacketConnect) &&
                   memcmp(buffer + 1, "CON", 3) == 0) {
            const struct PacketConnect *pkt = (struct PacketConnect *)buffer;

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

            ctx->last_input_packet_id = 0;
            atomic_store(&ctx->connected, true);
        } else if (packet_len >= 4 && memcmp(buffer + 1, "DIS", 3) == 0) {
            atomic_store(&ctx->connected, false);

            if (tcp_mode) {
                atomic_store(&ctx->exit_flag, true);
                print_err("[Android: INFO] Device disconnected (clean disconnect).\n");
                break;
            }

            if (strlen(ctx->remote_address)) {
                print_err("[Android: INFO] Device %s:%d disconnected.\n",
                          ctx->remote_address, ctx->remote_port);
                memset(ctx->remote_address, 0, 40);
            }

            if (tcp_mode) {
                break;
            }
        } else if (packet_len >= sizeof(struct PacketPing) &&
                   memcmp(buffer + 1, "PIN", 3) == 0) {
            if (!atomic_load(&ctx->connected)) {
                continue;
            }

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

            socket_send_to(sock, &addr, response, 13);
        } else if (packet_len >= sizeof(struct PacketCard) &&
                   memcmp(buffer + 1, "CRD", 3) == 0) {
            const struct PacketCard *pkt = (struct PacketCard *)buffer;

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

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

    free(ctx);

    return 0;
}

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

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

        const SOCKET sock = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);

        socket_set_timeout(sock, 2000);
        socket_bind(sock, htonl(INADDR_ANY), server_port);

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

        android_thread_ctx *ctx = malloc(sizeof(android_thread_ctx));
        ctx->sock = sock;
        ctx->exit_flag = ATOMIC_VAR_INIT(false);
        ctx->connected = ATOMIC_VAR_INIT(false);
        ctx->last_input_packet_id = 0;
        ctx->memory = memory;
        ctx->has_previous_led_status = false;
        ctx->led_skip_count = 0;

        HANDLE led_thread =
            (HANDLE)_beginthreadex(NULL, 0, led_broadcast_thread_proc, ctx, 0, NULL);
        HANDLE input_thread =
            (HANDLE)_beginthreadex(NULL, 0, input_recv_thread_proc, ctx, 0, NULL);

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

        CloseHandle(led_thread);
        CloseHandle(input_thread);

        free(ctx);
    } else {
        print_err("[Android: INFO] Mode: TCP\n");

        const SOCKET sock = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
        socket_set_timeout(sock, 50);
        socket_bind(sock, htonl(INADDR_ANY), server_port);

        listen(sock, 10);

        print_err("[Android: 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;

        while (
            (acc_socket = accept(sock, (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("[Android: INFO] Device %s:%d connected.\n", user_address,
                          user_socket.sin_port);
            }

            android_thread_ctx *ctx = malloc(sizeof(android_thread_ctx));
            ctx->sock = acc_socket;
            ctx->exit_flag = ATOMIC_VAR_INIT(false);
            ctx->connected = ATOMIC_VAR_INIT(true);
            ctx->last_input_packet_id = 0;
            ctx->memory = memory;
            ctx->has_previous_led_status = false;
            ctx->led_skip_count = 0;

            _beginthreadex(NULL, 0, led_broadcast_thread_proc, ctx, 0, NULL);
            _beginthreadex(NULL, 0, input_recv_thread_proc, ctx, 0, NULL);
        }
    }
}
#endif // defined(ENV32BIT)

HRESULT android_init_server(struct IPCMemoryInfo *memory) {
#ifdef ENV32BIT
    tcp_mode =
        GetPrivateProfileIntW(L"brokenithm", L"tcp", 1, 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("[Android:ERROR] WSA startup failed, os error %ld\n",
                  WSAGetLastError());
        return E_FAIL;
    }

    _beginthreadex(NULL, 0, server_thread_proc, memory, 0, NULL);
#endif // defined(ENV32BIT)

    return S_OK;
}