/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#include <fstream>
#include <iostream>
#include <stdio.h>
#include <stdlib.h>
#include <boost/algorithm/string.hpp>
#include <boost/asio.hpp>
#include <boost/assign.hpp>
#include <boost/bind.hpp>
#include "drill/drillc.hpp"
#include "drill/drillError.hpp"
#include "clientlib/channel.hpp"
#include "clientlib/drillClientImpl.hpp"
#include "clientlib/errmsgs.hpp"
#include "clientlib/logger.hpp"
#include "clientlib/rpcMessage.hpp"
#include "clientlib/utils.hpp"
#include "protobuf/GeneralRPC.pb.h"
#include "protobuf/UserBitShared.pb.h"

namespace Drill {

class DrillTestClient {

    public:

    DrillTestClient(Channel* pChannel):
    m_handshakeStatus(exec::user::SUCCESS),
    m_wbuf(MAX_SOCK_RD_BUFSIZE),
    m_rbuf(0){
            m_pChannel=pChannel;
            m_pError=NULL;
            m_coordinationId=Utils::s_randomNumber()%1729+1;
        }

        connectionStatus_t recvHandshake(){
            if(m_rbuf==NULL){
                m_rbuf = Utils::allocateBuffer(MAX_SOCK_RD_BUFSIZE);
            }

            m_pChannel->getIOService().reset();

            m_pChannel->getSocketStream().asyncRead(
                    boost::asio::buffer(m_rbuf, LEN_PREFIX_BUFLEN),
                    boost::bind(
                        &DrillTestClient::handleHandshake,
                        this,
                        m_rbuf,
                        boost::asio::placeholders::error,
                        boost::asio::placeholders::bytes_transferred)
                    );
            DRILL_MT_LOG(DRILL_LOG(LOG_DEBUG) << "DrillClientImpl::recvHandshake: async read waiting for server handshake response.\n";)
            m_pChannel->getIOService().run();
            if(m_rbuf!=NULL){
                Utils::freeBuffer(m_rbuf, MAX_SOCK_RD_BUFSIZE); m_rbuf=NULL;
            }

            if (m_pError != NULL) {
                DRILL_MT_LOG(DRILL_LOG(LOG_ERROR) << "DrillClientImpl::recvHandshake: failed to complete handshake with server."
                        << m_pError->msg << "\n";)
                    return static_cast<connectionStatus_t>(m_pError->status);
            }

            return CONN_SUCCESS;
        }

        void doReadFromSocket(ByteBuf_t inBuf, size_t bytesToRead, boost::system::error_code& errorCode) {
            // Check if bytesToRead is zero
            if(0 == bytesToRead) {
                return;
            }

            // Read all the bytes. In case when all the bytes were not read the proper
            // errorCode will be set.
            while(1){
                size_t dataBytesRead = m_pChannel->getSocketStream().readSome(boost::asio::buffer(inBuf, bytesToRead), errorCode);
                // Update the state
                bytesToRead -= dataBytesRead;
                inBuf += dataBytesRead;

                // Check if errorCode is EINTR then just retry otherwise break from loop
                if(EINTR != errorCode.value()) break;

                // Check if all the data is read then break from loop
                if(0 == bytesToRead) break;
            }
        }

        void handleHandshake(ByteBuf_t inBuf,
                const boost::system::error_code& err,
                size_t bytes_transferred) {
            boost::system::error_code error=err;
            DRILL_MT_LOG(DRILL_LOG(LOG_TRACE) << "Deadline timer cancelled." << std::endl;)
                if(!error){
                    rpc::InBoundRpcMessage msg;
                    uint32_t length = 0;
                    std::size_t bytes_read = rpc::lengthDecode(m_rbuf, length);
                    if(length>0){
                        const size_t leftover = LEN_PREFIX_BUFLEN - bytes_read;
                        const ByteBuf_t b = m_rbuf + LEN_PREFIX_BUFLEN;
                        const size_t bytesToRead=length - leftover;
                        doReadFromSocket(b, bytesToRead, error);

                        // Check if any error happen while reading the message bytes. If yes then return before decoding the Msg
                        if(error) {
                            DRILL_MT_LOG(DRILL_LOG(LOG_TRACE) << "DrillClientImpl::handleHandshake: ERR_CONN_RDFAIL. "
                                    << " Failed to read entire handshake message. with error: "
                                    << error.message().c_str() << "\n";)
                                handleConnError(CONN_FAILURE, getMessage(ERR_CONN_RDFAIL, "Failed to read entire handshake message"));
                            return;
                        }

                        // Decode the bytes into a valid RPC Message
                        if (!decode(m_rbuf+bytes_read, length, msg)) {
                            DRILL_MT_LOG(DRILL_LOG(LOG_TRACE) << "DrillClientImpl::handleHandshake: ERR_CONN_RDFAIL. Cannot decode handshake.\n";)
                                handleConnError(CONN_FAILURE, getMessage(ERR_CONN_RDFAIL, "Cannot decode handshake"));
                            return;
                        }
                    }else{
                        DRILL_MT_LOG(DRILL_LOG(LOG_TRACE) << "DrillClientImpl::handleHandshake: ERR_CONN_RDFAIL. No handshake.\n";)
                            handleConnError(CONN_FAILURE, getMessage(ERR_CONN_RDFAIL, "No handshake"));
                        return;
                    }
                    exec::user::BitToUserHandshake b2u;
                    b2u.ParseFromArray(msg.m_pbody.data(), msg.m_pbody.size());
                    this->m_handshakeErrorId=b2u.errorid();
                    this->m_handshakeErrorMsg=b2u.errormessage();
                }else{
                    // boost error
                    if(error==boost::asio::error::eof){ // Server broke off the connection
                        handleConnError(CONN_HANDSHAKE_FAILED, getMessage(ERR_CONN_NOHSHAKE, DRILL_RPC_VERSION));
                    }else{
                        handleConnError(CONN_FAILURE, getMessage(ERR_CONN_RDFAIL, error.message().c_str()));
                    }
                    return;
                }
            return;
        }

        connectionStatus_t handleConnError(connectionStatus_t status, const std::string& msg){
            DrillClientError* pErr = new DrillClientError(status, DrillClientError::CONN_ERROR_START+status, msg);
            if(m_pError!=NULL){ delete m_pError; m_pError=NULL;}
            m_pError=pErr;
            return status;
        }

        connectionStatus_t sendSyncCommon(rpc::OutBoundRpcMessage& msg) {
            encode(m_wbuf, msg);
            boost::system::error_code ec;
            doWriteToSocket(reinterpret_cast<char*>(m_wbuf.data()), m_wbuf.size(), ec);

            if(!ec) {
                return CONN_SUCCESS;
            } else {
                return handleConnError(CONN_FAILURE, getMessage(ERR_CONN_WFAIL, ec.message().c_str()));
            }
        }

        void doWriteToSocket(const char* dataPtr, size_t bytesToWrite,
                boost::system::error_code& errorCode) {
            if(0 == bytesToWrite) {
                return;
            }

            // Write all the bytes to socket. In case of error when all bytes are not successfully written
            // proper errorCode will be set.
            while(1) {
                size_t bytesWritten = m_pChannel->getSocketStream().writeSome(boost::asio::buffer(dataPtr, bytesToWrite), errorCode);
                // Update the state
                bytesToWrite -= bytesWritten;
                dataPtr += bytesWritten;

                if(EINTR != errorCode.value()) break;

                // Check if all the data is written then break from loop
                if(0 == bytesToWrite) break;
            }
        }

        connectionStatus_t validateHandshake(DrillUserProperties* properties){

            DRILL_MT_LOG(DRILL_LOG(LOG_TRACE) << "validateHandShake\n";)

                exec::user::UserToBitHandshake u2b;
            u2b.set_channel(exec::shared::USER);
            u2b.set_rpc_version(DRILL_RPC_VERSION);
            u2b.set_support_listening(true);
            u2b.set_support_timeout(DrillClientConfig::getHeartbeatFrequency() > 0);
            u2b.set_sasl_support(exec::user::SASL_PRIVACY);

            // Adding version info
            exec::user::RpcEndpointInfos* infos = u2b.mutable_client_infos();
            infos->set_name(DrillClientConfig::getClientName());
            infos->set_application(DrillClientConfig::getApplicationName());
            infos->set_version(DRILL_VERSION_STRING);
            infos->set_majorversion(DRILL_VERSION_MAJOR);
            infos->set_minorversion(DRILL_VERSION_MINOR);
            infos->set_patchversion(DRILL_VERSION_PATCH);

            if(properties != NULL && properties->size()>0){
                std::string username;
                std::string err;
                if(!properties->validate(err)){
                    DRILL_MT_LOG(DRILL_LOG(LOG_INFO) << "Invalid user input:" << err << std::endl;)
                }
                exec::user::UserProperties* userProperties = u2b.mutable_properties();

                std::map<char,int>::iterator it;
                for (std::map<std::string,std::string>::const_iterator propIter=properties->begin(); propIter!=properties->end(); ++propIter){
                    std::string currKey=propIter->first;
                    std::string currVal=propIter->second;
                    std::map<std::string,uint32_t>::const_iterator it=DrillUserProperties::USER_PROPERTIES.find(currKey);
                    if(it==DrillUserProperties::USER_PROPERTIES.end()){
                        DRILL_MT_LOG(DRILL_LOG(LOG_INFO) << "Connection property ("<< currKey
                                << ") is unknown" << std::endl;)
                            exec::user::Property* connProp = userProperties->add_properties();
                        connProp->set_key(currKey);
                        connProp->set_value(currVal);
                        continue;
                    }
                    if(IS_BITSET((*it).second,USERPROP_FLAGS_SERVERPROP)){
                        exec::user::Property* connProp = userProperties->add_properties();
                        connProp->set_key(currKey);
                        connProp->set_value(currVal);
                        //Username(but not the password) also needs to be set in UserCredentials
                        if(IS_BITSET((*it).second,USERPROP_FLAGS_USERNAME)){
                            exec::shared::UserCredentials* creds = u2b.mutable_credentials();
                            username=currVal;
                            creds->set_user_name(username);
                            //u2b.set_credentials(&creds);
                        }
                        if(IS_BITSET((*it).second,USERPROP_FLAGS_PASSWORD)){
                            DRILL_MT_LOG(DRILL_LOG(LOG_INFO) <<  currKey << ": ********** " << std::endl;)
                        }else{
                            DRILL_MT_LOG(DRILL_LOG(LOG_INFO) << currKey << ":" << currVal << std::endl;)
                        }
                    }// Server properties
                }
            }

            {
                boost::lock_guard<boost::mutex> lock(this->m_dcMutex);
                uint64_t coordId = ++m_coordinationId;

                rpc::OutBoundRpcMessage out_msg(exec::rpc::REQUEST, exec::user::HANDSHAKE, coordId, &u2b);
                sendSyncCommon(out_msg);
                DRILL_MT_LOG(DRILL_LOG(LOG_TRACE) << "Sent handshake request message. Coordination id: " << coordId << "\n";)
            }

            connectionStatus_t ret = recvHandshake();
            if(ret!=CONN_SUCCESS){
                return ret;
            }

            switch(this->m_handshakeStatus) {
                case exec::user::SUCCESS:
                    // reset io_service after handshake is validated before running queries
                    m_pChannel->getIOService().reset();
                    return CONN_SUCCESS;
                case exec::user::RPC_VERSION_MISMATCH:
                    DRILL_MT_LOG(DRILL_LOG(LOG_TRACE) << "Invalid rpc version.  Expected "
                            << DRILL_RPC_VERSION << ", actual "<< 0 << "." << std::endl;)
                        return handleConnError(CONN_BAD_RPC_VER, getMessage(ERR_CONN_BAD_RPC_VER, DRILL_RPC_VERSION,
                                    0,
                                    this->m_handshakeErrorId.c_str(),
                                    this->m_handshakeErrorMsg.c_str()));
                case exec::user::AUTH_FAILED:
                    DRILL_MT_LOG(DRILL_LOG(LOG_TRACE) << "Authentication failed." << std::endl;)
                        return handleConnError(CONN_AUTH_FAILED, getMessage(ERR_CONN_AUTHFAIL,
                                    this->m_handshakeErrorId.c_str(),
                                    this->m_handshakeErrorMsg.c_str()));
                case exec::user::UNKNOWN_FAILURE:
                    DRILL_MT_LOG(DRILL_LOG(LOG_TRACE) << "Unknown error during handshake." << std::endl;)
                        return handleConnError(CONN_HANDSHAKE_FAILED, getMessage(ERR_CONN_UNKNOWN_ERR,
                                    this->m_handshakeErrorId.c_str(),
                                    this->m_handshakeErrorMsg.c_str()));
                case exec::user::AUTH_REQUIRED:
                    DRILL_MT_LOG(DRILL_LOG(LOG_TRACE) << "Server requires SASL authentication." << std::endl;)
                    return handleConnError(CONN_HANDSHAKE_FAILED, getMessage(ERR_CONN_UNKNOWN_ERR,
                                                                             this->m_handshakeErrorId.c_str(),
                                                                             this->m_handshakeErrorMsg.c_str()));
                default:
                    DRILL_MT_LOG(DRILL_LOG(LOG_TRACE) << "Unknown return status." << std::endl;)
                        return handleConnError(CONN_HANDSHAKE_FAILED, getMessage(ERR_CONN_UNKNOWN_ERR,
                                    this->m_handshakeErrorId.c_str(),
                                    this->m_handshakeErrorMsg.c_str()));
            }
        }

        DrillClientError* m_pError;
    private:
        Channel* m_pChannel;
        int32_t m_coordinationId;
        std::string m_handshakeErrorId;
        std::string m_handshakeErrorMsg;
        exec::user::HandshakeStatus m_handshakeStatus;
        DataBuf m_wbuf;
        ByteBuf_t m_rbuf;
        boost::mutex m_dcMutex;
    

};

} // namespace Drill

using namespace Drill;

int main(int argc, char* argv[]){
    Channel *pChannel = NULL;
    ChannelContext *pChannelContext = NULL;
    std::string connectStr = "zk=localhost:2181/drill/drillbits1";
    //std::string connectStr = "drillbit=localhost:31090";
    channelType_t type;
    boost::asio::io_service ioService;

    bool isSSL = argc==2 && !(strcmp(argv[1], "ssl"));
    type = CHANNEL_TYPE_SOCKET;
    if(isSSL){
        type = CHANNEL_TYPE_SSLSTREAM;
    }
    Drill::DrillUserProperties props;
    props.setProperty(USERPROP_USERNAME, "admin");
    props.setProperty(USERPROP_PASSWORD, "admin");
    props.setProperty(USERPROP_CERTFILEPATH, "../../../test/ssl/drillTestCert.pem");

    pChannel = ChannelFactory::getChannel(type, ioService, connectStr.c_str(), &props);
    if(pChannel != NULL){
        connectionStatus_t connStat;
        connStat = pChannel->init();
        if(connStat != CONN_SUCCESS){
            std::cout << "Init Failed." << std::endl;
            return -1;
        }
        connStat = pChannel->connect();
        if(connStat != CONN_SUCCESS){
            std::cout << "Connect Failed." << std::endl;
            std::cout << pChannel->getError()->msg << std::endl;
            return -1;
        }
    } else{
        std::cout << "Channel creation failed." << std::endl;
        return -1;
    }
    std::cout << "Connected." << std::endl;
    std::cout << "Starting Drill handshake" << std::endl;


    DrillTestClient client(pChannel);

    connectionStatus_t stat = client.validateHandshake(&props);
    if(stat == CONN_SUCCESS){
        std::cout << "Handshake validated." << std::endl;
    } else{ 
        if(client.m_pError != NULL){
            std::cout << "Handshake failed: " << client.m_pError->msg << ". " << std::endl;
        } else{
            std::cout << "Handshake failed  with unknown error" << ". " << std::endl;
        }
    }

    return 0;

}