#include "string.h"
#include "lwip/netif.h"
#include "wm_include.h"
#include "FreeRTOS.h"
#include "FreeRTOSConfig.h"
#include "nano_shell_server_task.h"
#include "common.h"

extern int shell_printf(const char *format, ...);
extern void network_rx_callback(u16 len, char *data);

tls_os_mutex_t *socket_mutex = NULL;
int nano_shell_srv_sock = 0, client_sock = 0;
const char greeting_buffer[] = "\r\n"
             " _   _                     ____  _          _ _\r\n"
             "| \\ | | __ _ _ __   ___   / ___|| |__   ___| | |\r\n"
             "|  \\| |/ _` | '_ \\ / _ \\  \\___ \\| '_ \\ / _ \\ | |\r\n"
             "| |\\  | (_| | | | | (_) |  ___) | | | |  __/ | |\r\n"
             "|_| \\_|\\__,_|_| |_|\\___/  |____/|_| |_|\\___|_|_|\r\n"
             "\r\n"
             "Welcome to Nano-Shell remote access\r\n"
             "\r\n"
             " * Source:    https://github.com/lebinlv/nano-shell\r\n"
             " * Copyright: (c) Liber 2020\r\n"
             "\r\n";
char identity_buffer[250] = "";

bool network_write_char(const char c)
{
    bool toReturn = false;

    tls_os_mutex_acquire(socket_mutex, 0);
    if(client_sock > 0)
    {
        
        if(send(client_sock, &c, 1, 0) < 0)
        {
            // Failed to send data to client because he probably disconnected
            // or the connection broke
            client_sock = 0;
        }
        else
        {
            toReturn = true;
        }
    }
    tls_os_mutex_release(socket_mutex);
    
    return toReturn;
}

bool network_write_string(const char *str, size_t size)
{
    bool toReturn = false;

    if(client_sock > 0)
    {
        tls_os_mutex_acquire(socket_mutex, 0);
        if(send(client_sock, str, size, 0) < 0)
        {
            // Failed to send data to client because he probably disconnected
            // or the connection broke
            client_sock = 0;
        }
        else
        {
            toReturn = true;
        }
        tls_os_mutex_release(socket_mutex);
    }
    
    return toReturn;
}

bool disconnect_client(void)
{
    bool toReturn = false;

    if(client_sock > 0)
    {
        if(shutdown(client_sock, SHUT_WR) < 0)
        {

        }
        else
        {
            toReturn = true;
        }
        client_sock = 0;
    }

    return toReturn;
}

void nano_shell_server_task(void* param)
{
    (void)param;

    bool setup_error = false;
    char recv_buffer[256] = "";
    
    struct sockaddr_in nano_shell_srv_addr = { .sin_family = AF_INET, .sin_addr.s_addr = INADDR_ANY, .sin_port = htons(NANO_SHELL_SERVER_PORT)}, client_addr;
    socklen_t sockaddr_in_len = sizeof(struct sockaddr_in);
    
    //We initialize the mutex
    if(tls_os_mutex_create(0, &socket_mutex) != TLS_OS_SUCCESS)
    {
        shell_printf("Failed to create the mutex."NEW_LINE);
        setup_error = true;
    }

    //We setup the listening socket :
    if((nano_shell_srv_sock = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP)) < 0)
    {
        shell_printf("Failed to create nano_shell listening socket."NEW_LINE);
        setup_error = true;
    }

    if(bind(nano_shell_srv_sock, (struct sockaddr *) &nano_shell_srv_addr, sockaddr_in_len) < 0)
    {
        shell_printf("Failed to bind nano_shell listening socket to addr."NEW_LINE);
        setup_error = true;
    }

    //We only wait for one connection at a time because the nano_shell is not multi user anyway
    if(listen(nano_shell_srv_sock, 0) < 0)
    {
        shell_printf("Failed to mark nano_shell_sock as a listening socket."NEW_LINE);
        setup_error = true;
    }

    if(setup_error)
    {
        for(;;)
            tls_os_time_delay(portMAX_DELAY);
    }
    
    for(;;)
    {
        client_sock = accept(nano_shell_srv_sock, (struct sockaddr *)&client_addr, &sockaddr_in_len);

        if(client_sock < 0)
        {
            shell_printf("Failed to accept incoming connection."NEW_LINE);
        }

        if(!network_write_string(greeting_buffer, sizeof greeting_buffer))
        {
            shell_printf("Failed to send greetings to client - errno(%d)."NEW_LINE, errno);
        }

        sprintf(identity_buffer, "Connected from : %u.%u.%u.%u:%u"NEW_LINE NEW_LINE,((u8 *)&client_addr.sin_addr)[0], 
        ((u8 *)&client_addr.sin_addr)[1], 
        ((u8 *)&client_addr.sin_addr)[2], 
        ((u8 *)&client_addr.sin_addr)[3],
        ntohs(client_addr.sin_port));

        if(!network_write_string(identity_buffer, strlen(identity_buffer)))
        {
            shell_printf("Failed to send greetings to client - errno(%d)."NEW_LINE, errno);
        }

        for(;client_sock > 0;)
        {
            int result = recv(client_sock, recv_buffer, 255, 0);

            if(result < 0)
            {
                shell_printf("Failed to receive data from client - errno(%d)."NEW_LINE"Closing connection."NEW_LINE, errno);
                if(close(client_sock) < 0)
                {
                    shell_printf("Failed to close socket - errno(%d)."NEW_LINE, errno);
                }
                client_sock = 0;
            }
            else if(result == 0)
            {
                shell_printf("Client disconnected."NEW_LINE);
                if(close(client_sock) < 0)
                {
                    shell_printf("Failed to close socket - errno(%d)."NEW_LINE, errno);
                }
                client_sock = 0;
            }
            else //We pass the received data to the nano shell process
            {
                //Need to remove the \n at the end
                char *pos = strchr(recv_buffer, '\r');
                if(pos)
                {
                    *pos = '\n';
                    result = pos + 1 - recv_buffer;
                }
                network_rx_callback(result, recv_buffer);
            }
        }  
    }
}