/*
 * This file is part of the SSH Library
 *
 * Copyright (c) 2015 by Aris Adamantiadis <aris@badcode.be>
 *
 * The SSH Library is free software; you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published by
 * the Free Software Foundation; either version 2.1 of the License, or (at your
 * option) any later version.
 *
 * The SSH Library is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
 * or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public
 * License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with the SSH Library; see the file COPYING.  If not, write to
 * the Free Software Foundation, Inc., 59 Temple Place - Suite 330, Boston,
 * MA 02111-1307, USA.
 */

#include "libssh/priv.h"
#include "libssh/poll.h"
#include "libssh/callbacks.h"
#include "libssh/session.h"
#include <stdlib.h>
#define CHUNKSIZE 4096

#ifdef _WIN32
# ifdef HAVE_IO_H
#  include <io.h>
#  undef open
#  define open _open
#  undef close
#  define close _close
#  undef read
#  define read _read
#  undef unlink
#  define unlink _unlink
# endif /* HAVE_IO_H */
#endif

struct ssh_connector_struct {
    ssh_session session;

    ssh_channel in_channel;
    ssh_channel out_channel;

    socket_t in_fd;
    socket_t out_fd;

    ssh_poll_handle in_poll;
    ssh_poll_handle out_poll;

    ssh_event event;

    int in_available;
    int out_wontblock;

    struct ssh_channel_callbacks_struct in_channel_cb;
    struct ssh_channel_callbacks_struct out_channel_cb;

    enum ssh_connector_flags_e in_flags;
    enum ssh_connector_flags_e out_flags;
};

static int ssh_connector_channel_data_cb(ssh_session session,
                                         ssh_channel channel,
                                         void *data,
                                         uint32_t len,
                                         int is_stderr,
                                         void *userdata);
static int ssh_connector_channel_write_wontblock_cb(ssh_session session,
                                                    ssh_channel channel,
                                                    size_t bytes,
                                                    void *userdata);

ssh_connector ssh_connector_new(ssh_session session)
{
    ssh_connector connector;

    connector = calloc(1, sizeof(struct ssh_connector_struct));
    if (connector == NULL){
        ssh_set_error_oom(session);
        return NULL;
    }

    connector->session = session;
    connector->in_fd = SSH_INVALID_SOCKET;
    connector->out_fd = SSH_INVALID_SOCKET;

    ssh_callbacks_init(&connector->in_channel_cb);
    ssh_callbacks_init(&connector->out_channel_cb);

    connector->in_channel_cb.userdata = connector;
    connector->in_channel_cb.channel_data_function = ssh_connector_channel_data_cb;

    connector->out_channel_cb.userdata = connector;
    connector->out_channel_cb.channel_write_wontblock_function =
            ssh_connector_channel_write_wontblock_cb;

    return connector;
}

void ssh_connector_free (ssh_connector connector)
{
    if (connector->event != NULL){
        ssh_connector_remove_event(connector);
    }

    if (connector->in_poll != NULL) {
        ssh_poll_free(connector->in_poll);
        connector->in_poll = NULL;
    }

    if (connector->out_poll != NULL) {
        ssh_poll_free(connector->out_poll);
        connector->out_poll = NULL;
    }

    if (connector->in_channel != NULL) {
        ssh_remove_channel_callbacks(connector->in_channel,
                                     &connector->in_channel_cb);
    }
    if (connector->out_channel != NULL) {
        ssh_remove_channel_callbacks(connector->out_channel,
                                     &connector->out_channel_cb);
    }

    free(connector);
}

int ssh_connector_set_in_channel(ssh_connector connector,
                                  ssh_channel channel,
                                  enum ssh_connector_flags_e flags)
{
    connector->in_channel = channel;
    connector->in_fd = SSH_INVALID_SOCKET;
    connector->in_flags = flags;

    /* Fallback to default value for invalid flags */
    if (!(flags & SSH_CONNECTOR_STDOUT) && !(flags & SSH_CONNECTOR_STDERR)) {
        connector->in_flags = SSH_CONNECTOR_STDOUT;
    }

    return ssh_add_channel_callbacks(channel, &connector->in_channel_cb);
}

int ssh_connector_set_out_channel(ssh_connector connector,
                                  ssh_channel channel,
                                  enum ssh_connector_flags_e flags)
{
    connector->out_channel = channel;
    connector->out_fd = SSH_INVALID_SOCKET;
    connector->out_flags = flags;

    /* Fallback to default value for invalid flags */
    if (!(flags & SSH_CONNECTOR_STDOUT) && !(flags & SSH_CONNECTOR_STDERR)) {
        connector->in_flags = SSH_CONNECTOR_STDOUT;
    }

    return ssh_add_channel_callbacks(channel, &connector->out_channel_cb);
}

void ssh_connector_set_in_fd(ssh_connector connector, socket_t fd)
{
    connector->in_fd = fd;
    connector->in_channel = NULL;
}

void ssh_connector_set_out_fd(ssh_connector connector, socket_t fd)
{
    connector->out_fd = fd;
    connector->out_channel = NULL;
}

/* TODO */
static void ssh_connector_except(ssh_connector connector, socket_t fd)
{
    (void) connector;
    (void) fd;
}

/* TODO */
static void ssh_connector_except_channel(ssh_connector connector,
                                         ssh_channel channel)
{
    (void) connector;
    (void) channel;
}

/**
 * @internal
 *
 * @brief Reset the poll events to be followed for each file descriptors.
 */
static void ssh_connector_reset_pollevents(ssh_connector connector)
{
    if (connector->in_fd != SSH_INVALID_SOCKET) {
        if (connector->in_available) {
            ssh_poll_remove_events(connector->in_poll, POLLIN);
        } else {
            ssh_poll_add_events(connector->in_poll, POLLIN);
        }
    }

    if (connector->out_fd != SSH_INVALID_SOCKET) {
        if (connector->out_wontblock) {
            ssh_poll_remove_events(connector->out_poll, POLLOUT);
        } else {
            ssh_poll_add_events(connector->out_poll, POLLOUT);
        }
    }
}

/**
 * @internal
 *
 * @brief Callback called when a poll event is received on an input fd.
 */
static void ssh_connector_fd_in_cb(ssh_connector connector)
{
    unsigned char buffer[CHUNKSIZE];
    int r;
    int toread = CHUNKSIZE;
    int w;
    int total = 0;
    int rc;

    SSH_LOG(SSH_LOG_TRACE, "connector POLLIN event for fd %d", connector->in_fd);

    if (connector->out_wontblock) {
        if (connector->out_channel != NULL) {
            size_t size = ssh_channel_window_size(connector->out_channel);

            /* Don't attempt reading more than the window */
            toread = MIN(size, CHUNKSIZE);
        }

        r = read(connector->in_fd, buffer, toread);
        if (r < 0) {
            ssh_connector_except(connector, connector->in_fd);
            return;
        }

        if (connector->out_channel != NULL) {
            if (r == 0) {
                rc = ssh_channel_send_eof(connector->out_channel);
                (void)rc; /* TODO Handle rc? */
            } else if (r> 0) {
                /* loop around ssh_channel_write in case our window reduced due to a race */
                while (total != r){
                    if (connector->out_flags & SSH_CONNECTOR_STDOUT) {
                        w = ssh_channel_write(connector->out_channel,
                                              buffer + total,
                                              r - total);
                    } else {
                        w = ssh_channel_write_stderr(connector->out_channel,
                                                     buffer + total,
                                                     r - total);
                    }
                    if (w == SSH_ERROR) {
                        return;
                    }
                    total += w;
                }
            }
        } else if (connector->out_fd != SSH_INVALID_SOCKET) {
            if (r == 0){
                close (connector->out_fd);
                connector->out_fd = SSH_INVALID_SOCKET;
            } else {
                /*
                 * Loop around write in case the write blocks even for CHUNKSIZE
                 * bytes
                 */
                while (total != r) {
                    w = write(connector->out_fd, buffer + total, r - total);
                    if (w < 0){
                        ssh_connector_except(connector, connector->out_fd);
                        return;
                    }
                    total += w;
                }
            }
        } else {
            ssh_set_error(connector->session, SSH_FATAL, "output socket or channel closed");
            return;
        }
        connector->out_wontblock = 0;
        connector->in_available = 0;
    } else {
        connector->in_available = 1;
    }
}

/** @internal
 * @brief Callback called when a poll event is received on an output fd
 */
static void ssh_connector_fd_out_cb(ssh_connector connector){
    unsigned char buffer[CHUNKSIZE];
    int r;
    int w;
    int total = 0;
    SSH_LOG(SSH_LOG_TRACE, "connector POLLOUT event for fd %d", connector->out_fd);

    if(connector->in_available){
        if (connector->in_channel != NULL){
            r = ssh_channel_read_nonblocking(connector->in_channel, buffer, CHUNKSIZE, 0);
            if(r == SSH_ERROR){
                ssh_connector_except_channel(connector, connector->in_channel);
                return;
            } else if(r == 0 && ssh_channel_is_eof(connector->in_channel)){
                close(connector->out_fd);
                connector->out_fd = SSH_INVALID_SOCKET;
                return;
            } else if(r>0) {
                /* loop around write in case the write blocks even for CHUNKSIZE bytes */
                while (total != r){
                    w = write(connector->out_fd, buffer + total, r - total);
                    if (w < 0){
                        ssh_connector_except(connector, connector->out_fd);
                        return;
                    }
                    total += w;
                }
            }
        } else if (connector->in_fd != SSH_INVALID_SOCKET){
            /* fallback on the socket input callback */
            connector->out_wontblock = 1;
            ssh_connector_fd_in_cb(connector);
        } else {
            ssh_set_error(connector->session,
                          SSH_FATAL,
                          "Output socket or channel closed");
            return;
        }
        connector->in_available = 0;
        connector->out_wontblock = 0;
    } else {
        connector->out_wontblock = 1;
    }
}

/**
 * @internal
 *
 * @brief Callback called when a poll event is received on a file descriptor.
 *
 * This is for (input or output.
 *
 * @param[in] fd file descriptor receiving the event
 *
 * @param[in] revents received Poll(2) events
 *
 * @param[in] userdata connector
 *
 * @returns 0
 */
static int ssh_connector_fd_cb(ssh_poll_handle p,
                               socket_t fd,
                               int revents,
                               void *userdata)
{
    ssh_connector connector = userdata;

    (void)p;

    if (revents & POLLERR) {
        ssh_connector_except(connector, fd);
    } else if((revents & POLLIN) && fd == connector->in_fd) {
        ssh_connector_fd_in_cb(connector);
    } else if((revents & POLLOUT) && fd == connector->out_fd) {
        ssh_connector_fd_out_cb(connector);
    }
    ssh_connector_reset_pollevents(connector);

    return 0;
}

/**
 * @internal
 *
 * @brief Callback called when data is received on channel.
 *
 * @param[in] data Pointer to the data
 *
 * @param[in] len Length of data
 *
 * @param[in] is_stderr Set to 1 if the data are out of band
 *
 * @param[in] userdata The ssh connector
 *
 * @returns Amount of data bytes consumed
 */
static int ssh_connector_channel_data_cb(ssh_session session,
                                         ssh_channel channel,
                                         void *data,
                                         uint32_t len,
                                         int is_stderr,
                                         void *userdata)
{
    ssh_connector connector = userdata;
    int w;
    size_t window;

    (void) session;
    (void) channel;
    (void) is_stderr;

    SSH_LOG(SSH_LOG_TRACE,"connector data on channel");

    if (is_stderr && !(connector->in_flags & SSH_CONNECTOR_STDERR)) {
        /* ignore stderr */
        return 0;
    } else if (!is_stderr && !(connector->in_flags & SSH_CONNECTOR_STDOUT)) {
        /* ignore stdout */
        return 0;
    }

    if (connector->out_wontblock) {
        if (connector->out_channel != NULL) {
            int window_len;

            window = ssh_channel_window_size(connector->out_channel);
            window_len = MIN(window, len);

            /* Route the data to the right exception channel */
            if (is_stderr && (connector->out_flags & SSH_CONNECTOR_STDERR)) {
                w = ssh_channel_write_stderr(connector->out_channel,
                                             data,
                                             window_len);
            } else if (!is_stderr &&
                       (connector->out_flags & SSH_CONNECTOR_STDOUT)) {
                w = ssh_channel_write(connector->out_channel,
                                      data,
                                      window_len);
            } else if (connector->out_flags & SSH_CONNECTOR_STDOUT) {
                w = ssh_channel_write(connector->out_channel,
                                      data,
                                      window_len);
            } else {
                w = ssh_channel_write_stderr(connector->out_channel,
                                             data,
                                             window_len);
            }
            if (w == SSH_ERROR) {
                ssh_connector_except_channel(connector, connector->out_channel);
            }
        } else if (connector->out_fd != SSH_INVALID_SOCKET) {
            w = write(connector->out_fd, data, len);
            if (w < 0)
                ssh_connector_except(connector, connector->out_fd);
        } else {
            ssh_set_error(session, SSH_FATAL, "output socket or channel closed");
            return SSH_ERROR;
        }

        connector->out_wontblock = 0;
        connector->in_available = 0;
        if ((unsigned int)w < len) {
            connector->in_available = 1;
        }
        ssh_connector_reset_pollevents(connector);

        return w;
    } else {
        connector->in_available = 1;

        return 0;
    }
}

/**
 * @internal
 *
 * @brief Callback called when the channel is free to write.
 *
 * @param[in] bytes Amount of bytes that can be written without blocking
 *
 * @param[in] userdata The ssh connector
 *
 * @returns Amount of data bytes consumed
 */
static int ssh_connector_channel_write_wontblock_cb(ssh_session session,
                                                    ssh_channel channel,
                                                    size_t bytes,
                                                    void *userdata)
{
    ssh_connector connector = userdata;
    uint8_t buffer[CHUNKSIZE];
    int r, w;

    (void) channel;

    SSH_LOG(SSH_LOG_TRACE, "Channel write won't block");
    if (connector->in_available) {
        if (connector->in_channel != NULL) {
            size_t len = MIN(CHUNKSIZE, bytes);

            r = ssh_channel_read_nonblocking(connector->in_channel,
                                             buffer,
                                             len,
                                             0);
            if (r == SSH_ERROR) {
                ssh_connector_except_channel(connector, connector->in_channel);
            } else if(r == 0 && ssh_channel_is_eof(connector->in_channel)){
                ssh_channel_send_eof(connector->out_channel);
            } else if (r > 0) {
                w = ssh_channel_write(connector->out_channel, buffer, r);
                if (w == SSH_ERROR) {
                    ssh_connector_except_channel(connector,
                                                 connector->out_channel);
                }
            }
        } else if (connector->in_fd != SSH_INVALID_SOCKET) {
            /* fallback on on the socket input callback */
            connector->out_wontblock = 1;
            ssh_connector_fd_in_cb(connector);
            ssh_connector_reset_pollevents(connector);
        } else {
            ssh_set_error(session,
                          SSH_FATAL,
                          "Output socket or channel closed");

            return 0;
        }
        connector->in_available = 0;
        connector->out_wontblock = 0;
    } else {
        connector->out_wontblock = 1;
    }

    return 0;
}

int ssh_connector_set_event(ssh_connector connector, ssh_event event)
{
    int rc = SSH_OK;

    if ((connector->in_fd == SSH_INVALID_SOCKET &&
         connector->in_channel == NULL)
        || (connector->out_fd == SSH_INVALID_SOCKET &&
            connector->out_channel == NULL)) {
        rc = SSH_ERROR;
        ssh_set_error(connector->session,SSH_FATAL,"Connector not complete");
        goto error;
    }

    connector->event = event;
    if (connector->in_fd != SSH_INVALID_SOCKET) {
        if (connector->in_poll == NULL) {
            connector->in_poll = ssh_poll_new(connector->in_fd,
                                              POLLIN|POLLERR,
                                              ssh_connector_fd_cb,
                                              connector);
        }
        rc = ssh_event_add_poll(event, connector->in_poll);
        if (rc != SSH_OK) {
            goto error;
        }
    }

    if (connector->out_fd != SSH_INVALID_SOCKET) {
        if (connector->out_poll == NULL) {
            connector->out_poll = ssh_poll_new(connector->out_fd,
                                               POLLOUT|POLLERR,
                                               ssh_connector_fd_cb,
                                               connector);
        }

        rc = ssh_event_add_poll(event, connector->out_poll);
        if (rc != SSH_OK) {
            goto error;
        }
    }
    if (connector->in_channel != NULL) {
        rc = ssh_event_add_session(event,
                ssh_channel_get_session(connector->in_channel));
        if (rc != SSH_OK)
            goto error;
        if (ssh_channel_poll_timeout(connector->in_channel, 0, 0) > 0){
            connector->in_available = 1;
        }
    }
    if(connector->out_channel != NULL) {
        ssh_session session = ssh_channel_get_session(connector->out_channel);

        rc =  ssh_event_add_session(event, session);
        if (rc != SSH_OK) {
            goto error;
        }
        if (ssh_channel_window_size(connector->out_channel) > 0) {
            connector->out_wontblock = 1;
        }
    }

error:
    return rc;
}

int ssh_connector_remove_event(ssh_connector connector) {
    ssh_session session;

    if (connector->in_poll != NULL) {
        ssh_event_remove_poll(connector->event, connector->in_poll);
        ssh_poll_free(connector->in_poll);
        connector->in_poll = NULL;
    }

    if (connector->out_poll != NULL) {
        ssh_event_remove_poll(connector->event, connector->out_poll);
        ssh_poll_free(connector->out_poll);
        connector->out_poll = NULL;
    }

    if (connector->in_channel != NULL) {
        session = ssh_channel_get_session(connector->in_channel);

        ssh_event_remove_session(connector->event, session);
        connector->in_channel = NULL;
    }

    if (connector->out_channel != NULL) {
        session = ssh_channel_get_session(connector->out_channel);

        ssh_event_remove_session(connector->event, session);
        connector->out_channel = NULL;
    }
    connector->event = NULL;

    return SSH_OK;
}