#include "robomaster.h"
#include "client.h"
#include "crc.h"
#include "connection.h"

#include <stdlib.h>
#include <sys/socket.h>
#include <sys/select.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <string.h>
#include <fcntl.h>
#include <stdio.h>

static
inline
size_t
message_length(int cmd) {
    switch(cmd) {
        case SET_SDK_CONNECTION_CMD:
            return sizeof(struct SetSdkConnectionReq);
        case SDK_HEARTBEAT_CMD:
            return sizeof(struct SdkHeartbeatReq);
        case SET_SDK_MODE_CMD:
            return sizeof(struct SetSdkModeReq);
        case SET_SYSTEM_LED_CMD:
            return sizeof(struct SetSystemLedReq);
        case SET_ROBOT_MODE_CMD:
            return sizeof(struct SetRobotModeReq);
        case SUBNODE_RESET_CMD:
            return sizeof(struct SubNodeResetReq);
        case SUBSCRIBE_ADD_NODE_CMD:
            return sizeof(struct SubscribeAddNodeReq);
        case SET_WHEEL_SPEED_CMD:
            return sizeof(struct SetWheelSpeedReq);
        case CHASSIS_SPEED_MODE_CMD:
            return sizeof(struct ChassisSpeedModeReq);
        case GIMBAL_CTRL_SPEED_CMD:
            return sizeof(struct GimbalCtrlSpeedReq);
        case BLASTER_FIRE_CMD:
            return sizeof(struct BlasterFireReq);
        case STREAM_CTRL_CMD:
            return sizeof(struct StreamCtrlReq);
        case VISION_DETECT_ENABLE_CMD:
            return sizeof(struct VisionDetectEnableReq);
        default:
            return 0;
    }
}

static
inline
uint8_t
message_module(int cmd) {
    switch(cmd) {
        case SET_SDK_CONNECTION_CMD:
        case SDK_HEARTBEAT_CMD:
        case SET_SDK_MODE_CMD:
        case SET_SYSTEM_LED_CMD:
        case SET_ROBOT_MODE_CMD:
        case SUBNODE_RESET_CMD:
        case SUBSCRIBE_ADD_NODE_CMD:
            return host2byte(SDK_HOST, SDK_INDEX);
        case SET_WHEEL_SPEED_CMD:
        case CHASSIS_SPEED_MODE_CMD:
            return host2byte(CHASSIS_HOST, CHASSIS_INDEX);
        case GIMBAL_CTRL_SPEED_CMD:
            return host2byte(GIMBAL_HOST, GIMBAL_INDEX);
        case BLASTER_FIRE_CMD:
            return host2byte(BLASTER_HOST, BLASTER_INDEX);
        case STREAM_CTRL_CMD:
            return host2byte(STREAM_HOST, STREAM_INDEX);
        case VISION_DETECT_ENABLE_CMD:
            return host2byte(VISION_HOST, VISION_INDEX);
        default:
            return 0;
    }
}

// The greated file descriptor is needed for polling the sockets.
// It needs to be global for the whole process.
int max_fd = -1;

// TODO: Close the socket and return NULL on error
struct Connection*
connection_new(unsigned int source_port, const char* source_ip, unsigned int dest_port, const char* dest_ip)
{
    struct Connection* conn = malloc(sizeof(struct Connection));
    memset(conn, 0, sizeof(struct Connection));

    // Request a UDP socket
    conn->sockfd = socket(AF_INET, SOCK_DGRAM, 0);

    // Set the source address and port if they are provided
    if(source_port && source_ip) {
       struct sockaddr_in loc_addr;
       loc_addr.sin_family = AF_INET;
       loc_addr.sin_port = htons(source_port);
       loc_addr.sin_addr.s_addr = inet_addr(source_ip);
       if(bind(conn->sockfd, (struct sockaddr*)&loc_addr, sizeof(loc_addr)) < 0)
       {
           perror("unable to bind local port");
           exit(EXIT_FAILURE);
       }
    }
 
    // Make the socket non-blocking
    int flags = fcntl(conn->sockfd, F_GETFL);
    fcntl(conn->sockfd, F_SETFL, flags | O_NONBLOCK);

    // Set the address of the drone
    memset(&conn->remote_addr, 0, sizeof(conn->remote_addr));
    conn->addrlen = sizeof(conn->remote_addr);
    conn->remote_addr.sin_family = AF_INET;
    conn->remote_addr.sin_port = htons(dest_port);
    conn->remote_addr.sin_addr.s_addr = inet_addr(dest_ip);

    // File descriptors are numbers that count up sequentially,
    // so save the last one as the greatest file descriptor.
    // This is needed for polling the sockets later.
    max_fd = conn->sockfd;

    return conn;
}

struct Connection*
connection_poll_ready(struct ClientImp* client) {

    // Return a null connection if no sockets have been opened
    if(max_fd < 0)
        return NULL;

    // Add all the connections' socket file descriptors to a watch list
    fd_set read_fds;
    FD_ZERO(&read_fds);
    for(int i = 0; i < 2; i++)
        if(client->conns[i])
            FD_SET(client->conns[i]->sockfd, &read_fds);

    struct timeval timeout = {0, 0};
    int result = select(max_fd + 1, &read_fds, NULL, NULL, &timeout);

    // Check for socket polling errors
    if(result < 0) {
        perror("message polling failed");
	exit(EXIT_FAILURE);
    }

    // Return a null connection if nothing was received on any of them
    if (result == 0)
        return NULL;

    // Return the first connection with something to read
    for(int i = 0; i < 2; i++)
        if(FD_ISSET(client->conns[i]->sockfd, &read_fds))
            return client->conns[i];

    // Return a null connection if somehow none of them have anything to read
    return NULL;

}

void
connection_read(struct Connection* conn, union Message* resp) {

    memset(resp, 0, sizeof(union Message));

    if(!conn) return;

    int recvb = recvfrom(conn->sockfd, resp, sizeof(union Message), 0, (struct sockaddr*)&conn->remote_addr, &conn->addrlen);

    // Check for socket read errors
    if(recvb < 0) {
        perror("reading socket failed");
        exit(EXIT_FAILURE);
    }

    // Check for message errors
    if(message_validate(resp) != MESSAGEERR_NONE) {
        perror("invalid message");
        exit(EXIT_FAILURE);
    }
}

void
req_send(struct Connection* conn, union Request* req) {
    if(!conn || !req) return;
    size_t length = message_length(req->header.cmd);
    sendto(conn->sockfd, req, length, 0, (struct sockaddr*)&conn->remote_addr, conn->addrlen);
}

void
req_finalize(uint16_t seq, uint16_t cmd, bool need_ack, union Request* req) {

    size_t length = message_length(cmd);
    req->header.preamble = 0x55;
    req->header.length = (length & 0x1FFF) | 0x400;
    req->header.crc = crc8(req, 3);
    req->header.seq_id = seq;
    req->header.sender = host2byte(CLIENT_HOST, CLIENT_INDEX);
    req->header.receiver = message_module(cmd);
    req->header.ack_needed = need_ack;
    req->header.cmd = cmd;

    struct Footer* footer = (void*)req + length - sizeof(struct Footer);
    uint16_t crc = crc16(req, length - sizeof(struct Footer));
    footer->crc = crc;

}