#include "driver.hpp"
#include <map>
#include <iostream>
#include <cstdio>
#include <cstring>
#include <cassert>
#include <ws2tcpip.h>
#include <winsock2.h>
#include <synchapi.h>




class Driver {
public:
    static const int MAX_CHAN = DrivenEngine::MAX_CHAN;
    static const int CONSOLE_MAX = 512;
    DrivenEngine *driven_;
    SOCKET socket_[MAX_CHAN];
    bool connected_[MAX_CHAN];
    bool engine_wakeup_;
    char console_line_[CONSOLE_MAX + 1];
    int console_len_;
    std::unique_ptr<char> chbuf;

    static PADDRINFOA find_good_addr(PADDRINFOA addrinfo) {
        for (PADDRINFOA addr = addrinfo; addr != nullptr; addr = addr->ai_next) {
            if ((addr->ai_family == AF_INET) || (addr->ai_family == AF_INET6)) {
                return addr;
            }
        }
        return nullptr;
    }

    SOCKET open_connection(const std::string &target, std::string &err) {
        PADDRINFOA addrs = nullptr;
        PADDRINFOA goodaddr = nullptr;
        SOCKET sock = INVALID_SOCKET;
        std::string host, port;
        u_long mode = 1;  // 1 to enable non-blocking socket
        
        err = "";
        util::split_host_port(target, host, port);
        int status = getaddrinfo(host.c_str(), port.c_str(), nullptr, &addrs);
        while (status == WSATRY_AGAIN) {
            status = getaddrinfo(host.c_str(), port.c_str(), nullptr, &addrs);
            goto error;
        }
        if (status == WSAHOST_NOT_FOUND) {
            err = "host not found";
            goto error;
        }
        if (status != 0) {
            err = "DNS resolution hard failure";
            goto error;
        }
        goodaddr = find_good_addr(addrs);
        if (goodaddr == nullptr) {
            err = "host not an internet host";
            goto error;
        }
        sock = socket(goodaddr->ai_family, SOCK_STREAM, IPPROTO_TCP);
        assert(sock != INVALID_SOCKET);
        status = ioctlsocket(sock, FIONBIO, &mode);
        assert(status == 0);
        status = connect(sock, goodaddr->ai_addr, goodaddr->ai_addrlen);
        if (status != 0) {
            int errcode = WSAGetLastError();
            if (errcode != WSAEWOULDBLOCK) {
                std::ostringstream oss;
                oss << "Error " << errcode;
                err = oss.str();
                goto error;
            }
        }
        freeaddrinfo(addrs);
        return sock;
        
    error:
        if (sock != INVALID_SOCKET) closesocket(sock);
        if (addrs != nullptr) freeaddrinfo(addrs);
        return SOCKET_ERROR;
    }

    void init(DrivenEngine *de) {
        driven_ = de;
        for (int i = 0; i < MAX_CHAN; i++) {
            socket_[i] = INVALID_SOCKET;
            connected_[i] = false;
        }
        console_len_ = 0;
        engine_wakeup_ = false;
        chbuf.reset(new char[65536]);
    }

    void handle_lua_source() {
        if (driven_->drv_get_rescan_lua_source()) {
            driven_->drv_set_lua_source(util::read_lua_source("lua"));
            engine_wakeup_ = true;
        }
    }
        
    void handle_new_closed_sockets() { 
        std::set<int> chans;
        driven_->drv_get_new_closed(chans);
        for (int chid : chans) {
            if (socket_[chid] != INVALID_SOCKET) {
                assert(closesocket(socket_[chid] == 0));
                socket_[chid] = INVALID_SOCKET;
                connected_[chid] = false;
            }
        }
    }

    void handle_new_outgoing_sockets() {
        std::set<int> chans;
        driven_->drv_get_new_outgoing(chans);
        for (int chid : chans) {
            assert(socket_[chid] == INVALID_SOCKET);
            std::string err;
            SOCKET sock = open_connection(driven_->drv_get_target(chid), err);
            if (sock == INVALID_SOCKET) {
                driven_->drv_notify_close(chid);
                engine_wakeup_ = true;
            } else {
                socket_[chid] = sock;
                connected_[chid] = false;
            }
        }
    }

    void handle_console_output() {
        int nbytes; const char *bytes;
        driven_->drv_peek_outgoing(0, &nbytes, &bytes);
        if (nbytes > 0) {
            HANDLE hstdout = GetStdHandle(STD_OUTPUT_HANDLE);
            assert(hstdout != INVALID_HANDLE_VALUE);
            DWORD nwrote;
            if (nbytes > 10000) nbytes = 10000;
            assert(WriteConsoleA(hstdout, bytes, nbytes, &nwrote, nullptr));
            assert(nwrote > 0);
            driven_->drv_sent_outgoing(0, nwrote);
        }
    }

    // This is painful.  Win32 allows nonblocking read of keyboard events. But
    // it doesn't have any way to do nonblocking read of processed lines. So we
    // have to read individual events and do the line processing ourselves.
    void handle_console_input() {
        HANDLE hstdin = GetStdHandle(STD_INPUT_HANDLE);
        HANDLE hstdout = GetStdHandle(STD_OUTPUT_HANDLE);
        assert(hstdin != INVALID_HANDLE_VALUE);
        assert(hstdout != INVALID_HANDLE_VALUE);
        INPUT_RECORD inrecords[512];
        DWORD nread, nevents, nwrote;
        if (GetNumberOfConsoleInputEvents(hstdin, &nevents)) {
            if (nevents > 512) nevents = 512;
            ReadConsoleInputA(hstdin, inrecords, nevents, &nread);
            for (int i = 0; i < int(nread); i++) {
                const INPUT_RECORD &inr = inrecords[i];
                if (inr.EventType != KEY_EVENT) continue;
                const KEY_EVENT_RECORD &key = inr.Event.KeyEvent;
                if (!key.bKeyDown) continue;
                char c = key.uChar.AsciiChar;
                if ((c == '\r') || (c == '\n')) {
                    console_line_[console_len_++] = '\n';
                    assert(WriteConsoleA(hstdout, " \r\n", 3, &nwrote, nullptr));
                    assert(nwrote==3);
                    driven_->drv_recv_incoming(0, console_len_, console_line_);
                    console_len_ = 0;
                    engine_wakeup_ = true;
                } else if (c == '\b') {
                    if (console_len_ > 0) {
                        assert(WriteConsoleA(hstdout, "\b \b", 3, &nwrote, nullptr));
                        assert(nwrote==3);
                        console_len_ -= 1;
                    }
                } else if (c >= 32) {
                    if (console_len_ < CONSOLE_MAX) {
                        console_line_[console_len_++] = c;
                        assert(WriteConsoleA(hstdout, &c, 1, &nwrote, nullptr));
                        assert(nwrote==1);
                    }
                }
            }
        }
    }

    void handle_clock() {
    }

    bool calc_select_sets(fd_set &rfds, fd_set &wfds, fd_set &efds) const {
        FD_ZERO(&rfds);
        FD_ZERO(&wfds);
        FD_ZERO(&efds);
        bool any = false;
        for (int chid = 1; chid < MAX_CHAN; chid++) {
            SOCKET sock = socket_[chid];
            if (sock == INVALID_SOCKET) continue;
            any = true;
            FD_SET(sock, &rfds);
            if ((!connected_[chid]) || (!driven_->drv_outgoing_empty(chid))) {
                FD_SET(sock, &wfds);
            }
        }
        return any;
    }

    void close_socket(int chid) {
        assert(socket_[chid] != INVALID_SOCKET);
        assert(closesocket(socket_[chid]) == 0);
        driven_->drv_notify_close(chid);
        socket_[chid] = INVALID_SOCKET;
        connected_[chid] = false;
        engine_wakeup_ = true;
    }

    void handle_socket_input_output(int mstimeout) {
        fd_set rfds, wfds, efds;
        int nbytes; const char *bytes;
        bool any = calc_select_sets(rfds, wfds, efds);
        if (!any) {
            if (mstimeout>0) Sleep(mstimeout);
            return;
        }
        TIMEVAL timeout;
        timeout.tv_sec = mstimeout / 1000;
        timeout.tv_usec = (mstimeout - (timeout.tv_sec*1000)) * 1000;
        int status = select(1, &rfds, &wfds, &efds, &timeout);
        assert(status != SOCKET_ERROR);
        for (int chid = 1; chid < MAX_CHAN; chid++) {
            SOCKET sock = socket_[chid];
            if (sock == INVALID_SOCKET) continue;
            if (FD_ISSET(sock, &wfds)) {
                connected_[chid] = true;
                driven_->drv_peek_outgoing(chid, &nbytes, &bytes);
                if (nbytes > 0) {
                    int wbytes = send(sock, bytes, nbytes, 0);
                    if (wbytes == SOCKET_ERROR) {
                        close_socket(chid);
                        continue;
                    } else {
                        driven_->drv_sent_outgoing(chid, wbytes);
                    }
                }
            }
            if (FD_ISSET(sock, &rfds)) {
                int nrecv = recv(sock, chbuf.get(), 65536, 0);
                if ((nrecv == SOCKET_ERROR) || (nrecv == 0)) {
                    close_socket(chid);
                    continue;
                } else {
                    driven_->drv_recv_incoming(chid, nrecv, chbuf.get());
                    engine_wakeup_ = true;
                }
            }
        }
    }

    void drive(DrivenEngine *de) {
        WSADATA whocares;
        assert(WSAStartup(MAKEWORD(2,2), &whocares) == 0);
        HANDLE hconsole = GetStdHandle(STD_INPUT_HANDLE);
        assert(hconsole != INVALID_HANDLE_VALUE);
        init(de);
        DrivenEngine::set(de);
        driven_->drv_set_lua_source(util::read_lua_source("lua"));
        driven_->drv_invoke_event_init();
        while (!de->drv_get_stop_driver()) {
            engine_wakeup_ = false;
            handle_lua_source();
            handle_new_closed_sockets();
            handle_new_outgoing_sockets();
            handle_console_output();
            handle_console_input();
            int mstimeout = engine_wakeup_ ? 0 : 100;
            handle_socket_input_output(mstimeout);
            handle_clock();
            if (engine_wakeup_) {
                de->drv_invoke_event_update();
            }
        }
        DrivenEngine::set(nullptr);
    }
};


void driver_drive(DrivenEngine *de) {
    Driver driver;
    driver.drive(de);
}