/* rfb.c : Remote Frame Buffer Server */
/****************************************************************************/
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <signal.h>
#include <stdint.h>
/**/
#include <unistd.h>
#include <pthread.h>
#include <sys/wait.h>
#include <sys/uio.h>
/**/
#include "net.h"
/****************************************************************************/

#define DEFAULT_RFB_PORT            5900
#define DEFAULT_LISTEN_BACKLOG      10
#define DEFAULT_RFB_VERSION         "RFB 003.003\n"

typedef enum  {
    RFB_AUTH_CONNECTION_FAILED=            0,
    RFB_AUTH_PUBLIC=            1,
    RFB_AUTH_VNC=               2
} rfb_auth_t;

typedef enum {
    RFB_MSG_CLIENT_SETPIXELFORMAT=0,
    RFB_MSG_CLIENT_FIXCOLORMAPENTRIES=1,
    RFB_MSG_CLIENT_SETENCODINGS=2,
    RFB_MSG_CLIENT_FRAMEBUFFERUPDATEREQUEST=3,
    RFB_MSG_CLIENT_KEYEVENT=4,
    RFB_MSG_CLIENT_POINTEREVENT=5,
    RFB_MSG_CLIENT_CUTTEXT=6,
} rfb_msg_client_t;

typedef enum {
    RFB_MSG_SERVER_FRAMEBUFFERUPDATE=0,
    RFB_MSG_SERVER_SETCOLORMAPENTRIES=1,
    RFB_MSG_SERVER_BELL=2,
    RFB_MSG_SERVER_CUTTEXT=3,
} rfb_msg_server_t;

/****************************************************************************/
void shit_fill(uint8_t *fb, int w, int h)
{
    int x,y;
    for(y=0;y<h;y++) {
        for(x=0;x<w;x++) {
            fb[x+y*w]=x;
        }
    }
}

/****************************************************************************/
#ifndef __BYTE_ORDER
#error __BYTE_ORDER must be defined to be __LITTLE_ENDIAN or __BIG_ENDIAN
#endif

#if __BYTE_ORDER == __LITTLE_ENDIAN
#define LE32(x) (x)
#define LE64(x) (x)
#define LE16(x) (x)
#define BE16(x) (((uint16_t)((x)<<8))|((uint16_t)(x)>>8))
#define BE32(x) (((uint32_t)(BE16((x)))<<16)|(BE16((uint32_t)(x)>>16)))
#define BE64(x) (((uint64_t)(BE32((x)))<<32)|(BE32((uint64_t)(x)>>32)))
#else
#define LE16(x) (((uint16_t)((x)<<8))|((uint16_t)(x)>>8))
#define LE32(x) (((uint32_t)(LE16((x)))<<16)|(LE16((uint32_t)(x)>>16)))
#define LE64(x) (((uint64_t)(LE32((x)))<<32)|(LE32((uint64_t)(x)>>32)))
#define BE32(x) (x)
#define BE64(x) (x)
#define BE16(x) (x)
#endif
/****************************************************************************/
struct rfb_pixel_format {
    uint8_t                     bits_per_pixel;
    uint8_t                     depth;
    uint8_t                     bigendian_flag;
    uint8_t                     truecolor_flag;
    uint16_t                    redmax;
    uint16_t                    greenmax;
    uint16_t                    bluemax;
    uint8_t                     redshift;
    uint8_t                     greenshift;
    uint8_t                     blueshift;
    uint8_t                     pad[3];
};

struct rfb_client_info
{
    int clientfd;
    int width;
    int height;
    int bytes_per_line;
    uint8_t *fb;
    char *desktop_name;
};

/****************************************************************************/
/* RFB Protocol */

/* send the RFB version we support and recieve the version supported by the
 * client. */
static int exchange_version(int clientfd)
{
    unsigned char version_buf[12];
    int version_buf_len;
    write(clientfd, DEFAULT_RFB_VERSION, 12); 
    version_buf_len=read(clientfd, version_buf, 12);
    
    printf("RFB version: %.*s\n",version_buf_len-1,version_buf);
    return 1;
}

/* sends the authentication type that the client failed for some reason. The
 * reason is passed as a null terminated string in plain 7-bit ascii. If reason
 * is left NULL then the default message "Failed" is used. 
 * Return: 0=failure, 1=success */
static int send_authentication_failed(int clientfd, const char *reason)
{
    uint32_t authvalue_be;
    uint32_t reason_len_be;
    int reason_len;
    struct iovec vector[] = {
        {&authvalue_be, sizeof authvalue_be},
        {&reason_len_be, sizeof reason_len_be},
        {&reason, 0}
    };

    if(!reason) {
        reason="Failed";
    }
    authvalue_be=BE32(RFB_AUTH_CONNECTION_FAILED);
    vector[2].iov_len=reason_len=strlen(reason)-1;
    reason_len_be=BE32(reason_len);
    writev(clientfd, vector, sizeof vector/sizeof *vector);
    return 1;
}

/* sends authentication type indication that no authentication is required
 * Return: 0=failure, 1=success */
static int send_authentication_public(int clientfd)
{
    uint32_t value;
    value=BE32(RFB_AUTH_PUBLIC);
    write(clientfd, &value, sizeof value);
    return 1;
}

/* sets shared flag if the client is allowing shared access. clear flag if
 * client requests exclusive access. shared_flag is only set if not NULL.
 * Return: 0=failure, 1=success */
static int get_shared_flag(int clientfd, int *shared_flag)
{
    uint8_t shared;
    read(clientfd, &shared, sizeof shared);
    if(shared_flag) *shared_flag=shared;
    return 1;
};

static int rfb_server_initialize(int clientfd, int width, int height, const
char *desktop_name)
{
    uint16_t                    width_be;
    uint16_t                    height_be;
    struct rfb_pixel_format     format;
    uint32_t                    name_len;
    int desktop_len;
    struct iovec vector[] = {
        {&width_be, sizeof width_be},
        {&height_be, sizeof height_be},
        {&format, sizeof format},
        {&name_len, sizeof name_len},
        {(void*)desktop_name, 0},
    };

    width_be=BE16(width);
    height_be=BE16(height);
    /* TODO: figure out what these format values should be */
    format.bits_per_pixel=8;
    format.depth=8;
    format.bigendian_flag=0;
    format.truecolor_flag=1;
    format.redmax=BE16(0x7);
    format.bluemax=BE16(0x3);
    format.greenmax=BE16(0x7);
    format.redshift=0;
    format.blueshift=6;
    format.greenshift=3;
    memset(format.pad,0,sizeof format.pad);

    /* desktop string */
    desktop_len=strlen(desktop_name)-1;
    name_len=BE32(desktop_len);
    vector[4].iov_len=desktop_len;

    writev(clientfd, vector, sizeof vector / sizeof *vector);
    return 1;
}

static int get_client_msgtype(int clientfd, int *msg_type)
{
    uint8_t type;
    int ret;
    ret=read(clientfd, &type, 1)>0;
    if(ret && msg_type) *msg_type=type;
    return ret;
};

static int get_client_setpixelformat(int clientfd, struct rfb_pixel_format *format)
{
    uint8_t pad[3];
    struct iovec vector[] = {
        {&pad, sizeof pad},
        {format, sizeof *format},
    };
    int ret;

    ret=readv(clientfd, vector, sizeof vector / sizeof *vector)>0;
    printf( "Client Format:\n"
            "  bits per pixel  = %d\n"
            "  depth           = %d\n"
            "  big endian      = %d\n"
            "  true-color      = %d\n"
            "  Red max         = %u\n"
            "  Green max       = %u\n"
            "  Blue max        = %u\n"
            "  Red Shift       = %u\n"
            "  Green Shift     = %u\n"
            "  Blue Shift      = %u\n"
            ,
            format->bits_per_pixel,
            format->depth,
            format->bigendian_flag,
            format->truecolor_flag,
            BE16(format->redmax),
            BE16(format->greenmax),
            BE16(format->bluemax),
            format->redshift,
            format->greenshift,
            format->blueshift);
    return ret;
};

static int get_client_fixcolormapentries(int clientfd)
{
    int nr_color;
    int start_color;
    uint16_t start_color_be; /* first color */
    uint16_t nr_color_be; /* number of colors */
    uint8_t pad[3];
    struct iovec vector[] = {
        {pad, sizeof pad},
        {&start_color_be, sizeof start_color_be},
        {&nr_color_be, sizeof nr_color_be},
    };

    readv(clientfd, vector, sizeof vector / sizeof *vector);
    nr_color=BE32(nr_color_be);
    start_color=BE32(start_color_be);
    {
        uint16_t color_data[nr_color*3];
        read(clientfd, color_data, sizeof color_data);
    }
    return 1;
}

static int get_client_setencodings(int clientfd)
{
    int nr_encodings;
    uint8_t pad;
    uint16_t nr_encodings_be;
    struct iovec vector[] = {
        {&pad, sizeof pad},
        {&nr_encodings_be, sizeof nr_encodings_be},
    };

    readv(clientfd, vector, sizeof vector / sizeof *vector);
    nr_encodings=BE16(nr_encodings_be);
    {
        uint32_t encodings_data[nr_encodings];
        read(clientfd, encodings_data, sizeof encodings_data);
    }
    /* TODO: put encodings_data somewhere */
    return 1;
}

static int get_client_framebufferupdaterequest(int clientfd, int *x, int *y,
int *w, int *h, int *inc)
{
    uint8_t incremental;
    uint16_t x_be, y_be, w_be, h_be;
    struct iovec vector[] = {
        {&incremental, sizeof incremental},
        {&x_be, sizeof x_be},
        {&y_be, sizeof y_be},
        {&w_be, sizeof w_be},
        {&h_be, sizeof h_be},
    };
    readv(clientfd, vector, sizeof vector / sizeof *vector);
    if(x) *x=BE16(x_be);
    if(y) *y=BE16(y_be);
    if(w) *w=BE16(w_be);
    if(h) *h=BE16(h_be);
    if(inc) *inc=incremental;
    return 1;
};

static int get_client_keyevent(int clientfd, int *key, int *down_flag)
{
    uint32_t key_be;
    uint8_t  down_flag_be;
    uint8_t  pad[2];
    struct iovec vector[] = {
        {&down_flag_be, sizeof down_flag_be},
        {pad, sizeof pad},
        {&key_be, sizeof key_be},
    };
    int ret;

    ret=readv(clientfd, vector, sizeof vector / sizeof *vector)>0;
    if(ret) {
        if(key) *key=BE32(key_be);
        if(down_flag) *down_flag=BE32(down_flag_be);
    }
    return ret;
}

static int get_client_pointerevent(int clientfd, int *x, int *y, int *button)
{
    uint8_t button_be;
    uint16_t x_be;
    uint16_t y_be;
    struct iovec vector[] = {
        {&button_be, sizeof button_be},
        {&x_be, sizeof x_be},
        {&y_be, sizeof y_be},
    };
    int ret;

    ret=readv(clientfd, vector, sizeof vector / sizeof *vector)>0;
    if(ret) {
        if(x) *x=BE16(x_be);
        if(y) *y=BE16(y_be);
        if(button) *button=button_be;
    }
    return ret;
};

static int get_client_cuttext(int clientfd, int len, char *str)
{
    uint8_t pad[3];
    uint32_t len_be;
    struct iovec vector[] = {
        {pad, sizeof pad},
        {&len_be, sizeof len_be},
    };
    int ret;
    int clen;

    ret=readv(clientfd, vector, sizeof vector / sizeof *vector)>0;
    if(ret) {
        clen=BE32(len_be);
        if(clen>=len) {
            /* TODO: truncate the buffer instead of failing */
            return 0; /* failure */
        }
        read(clientfd,str,clen);
        str[clen]=0;
    }
    return ret;
}
static int put_server_bell(int clientfd)
{
    uint8_t bell=RFB_MSG_SERVER_BELL;
    printf("Bell.\n");
    return write(clientfd, &bell, sizeof bell)>0;
}

static int put_server_frameupdate(struct rfb_client_info *client, int x, int
y, int w, int h)
{
    struct {
        uint8_t type;
        uint8_t pad;
        uint16_t nr_rect;
    } msg_header = {RFB_MSG_SERVER_FRAMEBUFFERUPDATE,0,BE16(1)};
    struct { 
        uint16_t    x;
        uint16_t    y;
        uint16_t    w;
        uint16_t    h;
        uint32_t    type;
    } rect_body;
    struct iovec vector[] = {
        {&msg_header, sizeof msg_header},
        {&rect_body, sizeof rect_body},
        {0, 0},
    };
    /* TODO: this is just a test. replace with real stuff */
    rect_body.x=BE16(x);
    rect_body.y=BE16(y);
    rect_body.w=BE16(w);
    rect_body.h=BE16(h);
    rect_body.type=BE32(0);
    vector[2].iov_base=client->fb + x + y*client->bytes_per_line;
    vector[2].iov_len=w*h;
    return writev(client->clientfd, vector, sizeof vector / sizeof *vector)>0;
}

/****************************************************************************/
static void *do_rfb_client(void *p) __attribute__((noreturn));
static void *do_rfb_client(void *p)
{
    struct rfb_client_info *client=p;
    int shared_flag;
    int msg_type;
    printf("Client connection!\n");
    /* you can do stuff with the client here */
    exchange_version(client->clientfd);
    if(0) send_authentication_failed(client->clientfd,"Monkeys got it.");
    send_authentication_public(client->clientfd);
    get_shared_flag(client->clientfd, &shared_flag);
    printf("Shared Flag: %d\n",shared_flag);
    if(!client->fb) {
        client->fb=calloc(client->bytes_per_line, client->height);
        shit_fill(client->fb,client->bytes_per_line, client->height);
    }
    rfb_server_initialize(client->clientfd, 
        client->width, client->height,      /* width, height */
        client->desktop_name
        );

    put_server_frameupdate(client,0,0,client->width,client->height);

    while(1) {
        if(!get_client_msgtype(client->clientfd, &msg_type)) {
            close(client->clientfd);
            pthread_exit(NULL);
        }
        printf("Message Type: %d\n",msg_type);
        switch((rfb_msg_client_t)msg_type) {
            case RFB_MSG_CLIENT_SETPIXELFORMAT: {
                struct rfb_pixel_format     pixel_format;
                get_client_setpixelformat(client->clientfd, &pixel_format);
            }
            break;
            case RFB_MSG_CLIENT_FIXCOLORMAPENTRIES: {
                get_client_fixcolormapentries(client->clientfd);
            }
            break;
            case RFB_MSG_CLIENT_SETENCODINGS: {
                get_client_setencodings(client->clientfd);
            }
            break;
            case RFB_MSG_CLIENT_FRAMEBUFFERUPDATEREQUEST: {
                int x,y,w,h;
                int incremental;
                get_client_framebufferupdaterequest(client->clientfd,&x,&y,&w,&h,&incremental);
                printf("Update: (%d,%d) - (%d,%d) Inc:%d\n",x,y,w,h,incremental);
                if(incremental) {
                } else {
                    put_server_frameupdate(client,x,y,w,h);
                }
            }
            break;
            case RFB_MSG_CLIENT_KEYEVENT: {
                int key;
                int down_flag;
                get_client_keyevent(client->clientfd, &key, &down_flag);
                printf("KeyEvent: key=%04X\n",key);
                put_server_bell(client->clientfd);
            }
            break;
            case RFB_MSG_CLIENT_POINTEREVENT: {
                int x, y, button;
                get_client_pointerevent(client->clientfd, &x, &y, &button);
                printf("Pointer: x=%d y=%d button=%02X\n",x,y,button);
            }
            break;
            case RFB_MSG_CLIENT_CUTTEXT: {
                char str[256];
                if(!get_client_cuttext(client->clientfd, sizeof str, str)) {
                    close(client->clientfd);
                    pthread_exit(NULL);
                }
                printf("Cuttext: %s\n",str);
            }
            break;
            default:
                /* TODO: handle failures better */
                printf("Unknown message type!\n");
                close(client->clientfd);
                pthread_exit(NULL);
        }
    }
    
    pthread_exit(NULL);
}

/****************************************************************************/
/* Main */
int main(/*int argc, char **argv*/)
{
    int clientfd;
    int listenfd;
    struct sockaddr_in clientaddr;
    pthread_t thr[16];
    struct rfb_client_info clientinfo[16];
    void *thr_ret;

    listenfd=net_listen_port(DEFAULT_RFB_PORT, DEFAULT_LISTEN_BACKLOG);
    if(listenfd<0) { return EXIT_FAILURE; }
    printf("Listening!\n");

    clientfd=net_client_accept(listenfd, &clientaddr);
    if(clientfd<0) { close(listenfd); return EXIT_FAILURE; }
    /* TODO: the creation of clients needs to be done with a function */
    clientinfo[0].clientfd=clientfd;
    clientinfo[0].width=320;
    clientinfo[0].height=240;
    clientinfo[0].bytes_per_line=320*1; /* must be a multiple of something? */
    clientinfo[0].fb=NULL;
    clientinfo[0].desktop_name=strdup("Jon's RFB Server!");

    pthread_create(&thr[0], NULL, do_rfb_client, &clientinfo[0]);

    pthread_join(thr[0], &thr_ret);
    close(clientfd);
    close(listenfd);

    return EXIT_SUCCESS;
}