/////////////////////////////////////////////////////////////////////////////////
//
//  Copyright (C) 2019-     Equinor ASA
//
//  ResInsight is free software: you can redistribute it and/or modify
//  it under the terms of the GNU General Public License as published by
//  the Free Software Foundation, either version 3 of the License, or
//  (at your option) any later version.
//
//  ResInsight 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 General Public License at <http://www.gnu.org/licenses/gpl.html>
//  for more details.
//
//////////////////////////////////////////////////////////////////////////////////

#include "RiaGrpcServer.h"

#include "RiaDefines.h"
#include "RiaLogging.h"

#include "RiaGrpcCallbacks.h"
#include "RiaGrpcCaseService.h"
#include "RiaGrpcServiceInterface.h"

#include "RigCaseCellResultsData.h"
#include "RigMainGrid.h"
#include "RimEclipseCase.h"
#include "RimProject.h"

#include "cafAssert.h"
#include "cafProgressInfo.h"

#include <grpc/support/log.h>
#include <grpcpp/grpcpp.h>

#include <QTcpServer>

#include <atomic>
#include <mutex>
#include <thread>

using grpc::CompletionQueue;
using grpc::Server;
using grpc::ServerAsyncResponseWriter;
using grpc::ServerBuilder;
using grpc::ServerCompletionQueue;
using grpc::ServerContext;
using grpc::Status;

//==================================================================================================
//
// The GRPC server implementation
//
//==================================================================================================
class RiaGrpcServerImpl
{
public:
    RiaGrpcServerImpl( int portNumber );
    ~RiaGrpcServerImpl();
    int    portNumber() const;
    bool   isRunning() const;
    void   run();
    void   runInThread();
    void   initialize();
    size_t processAllQueuedRequests();
    void   quit();

private:
    void waitForNextRequest();
    void process( RiaGrpcCallbackInterface* method );

private:
    int                                                 m_portNumber;
    std::unique_ptr<grpc::ServerCompletionQueue>        m_completionQueue;
    std::unique_ptr<grpc::Server>                       m_server;
    std::list<std::shared_ptr<RiaGrpcServiceInterface>> m_services;
    std::list<RiaGrpcCallbackInterface*>                m_unprocessedRequests;
    std::list<RiaGrpcCallbackInterface*>                m_allocatedCallbakcs;
    std::mutex                                          m_requestMutex;
    // Guards m_allocatedCallbakcs, which is touched both from the main thread and from the
    // completion queue thread when serving callbacks flagged as runsOnServerThread().
    std::mutex m_allocatedCallbackMutex;
    // Serializes shutdown against an in-flight dispatch on the completion queue thread, so that
    // no request handler can be registered on a queue that is about to be shut down.
    std::mutex        m_shutdownMutex;
    std::atomic<bool> m_isShuttingDown{ false };
    std::thread       m_thread;
};

//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
RiaGrpcServerImpl::RiaGrpcServerImpl( int portNumber )
    : m_portNumber( portNumber )
{
}

//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
RiaGrpcServerImpl::~RiaGrpcServerImpl()
{
    quit();
}

//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
int RiaGrpcServerImpl::portNumber() const
{
    return m_portNumber;
}

//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
bool RiaGrpcServerImpl::isRunning() const
{
    return m_server != nullptr;
}

//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
void RiaGrpcServerImpl::run()
{
    initialize();
    while ( true )
    {
        waitForNextRequest();
    }
}

//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
void RiaGrpcServerImpl::runInThread()
{
    initialize();
    if ( m_server )
    {
        m_thread = std::thread( &RiaGrpcServerImpl::waitForNextRequest, this );
    }
}

//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
void RiaGrpcServerImpl::initialize()
{
    CAF_ASSERT( m_portNumber >= 0 && m_portNumber <= (int)std::numeric_limits<quint16>::max() );

    ServerBuilder builder;

    // When setting port number to 0, grpc will find and use a valid port number
    // The port number is assigned to the m_portNumber variable after calling builder.BuildAndStart()
    QString requestedServerAddress = QString( "localhost:%1" ).arg( m_portNumber );
    builder.AddListeningPort( requestedServerAddress.toStdString(), grpc::InsecureServerCredentials(), &m_portNumber );

    // A single request or reply can be large, e.g. the schedule text generated from a full
    // production history or a raw text block imported into the event timeline. The gRPC default
    // limit of 4 MB is easily exceeded by such payloads, so raise it well above any realistic deck
    // size. The client uses a matching limit (see rips/instance.py).
    const int maxMessageSize = 128 * 1024 * 1024;
    builder.SetMaxReceiveMessageSize( maxMessageSize );
    builder.SetMaxSendMessageSize( maxMessageSize );

    for ( auto key : RiaGrpcServiceFactory::instance()->allKeys() )
    {
        std::shared_ptr<RiaGrpcServiceInterface> service( RiaGrpcServiceFactory::instance()->create( key ) );
        builder.RegisterService( dynamic_cast<grpc::Service*>( service.get() ) );
        m_services.push_back( service );
    }

    m_completionQueue = builder.AddCompletionQueue();
    m_server          = builder.BuildAndStart();

    QString serverAddress = QString( "localhost:%1" ).arg( m_portNumber );

    if ( !m_server )
    {
        RiaLogging::error( QString( "Failed to start server on %1" ).arg( serverAddress ).toStdString() );
        return;
    }

    RiaLogging::debug( QString( "Server listening on %1" ).arg( serverAddress ).toStdString() );

    // Spawn new CallData instances to serve new clients.
    for ( auto service : m_services )
    {
        for ( auto callback : service->createCallbacks() )
        {
            process( callback );
        }
    }
}

//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
size_t RiaGrpcServerImpl::processAllQueuedRequests()
{
    std::list<RiaGrpcCallbackInterface*> waitingRequests;
    {
        // Block only while transferring the unprocessed requests to a local function list
        std::lock_guard<std::mutex> requestLock( m_requestMutex );
        waitingRequests.swap( m_unprocessedRequests );
    }
    size_t count = waitingRequests.size();

    // Now free to receive new requests from client while processing the current ones.
    while ( !waitingRequests.empty() )
    {
        RiaGrpcCallbackInterface* method = waitingRequests.front();
        waitingRequests.pop_front();
        process( method );
    }
    return count;
}

//--------------------------------------------------------------------------------------------------
/// Gracefully shut down the GRPC server.
/// BE VERY CAREFUL ABOUT CHANGING THE ORDER IN THIS METHOD. IT IS IMPORTANT!
//--------------------------------------------------------------------------------------------------
void RiaGrpcServerImpl::quit()
{
    if ( m_server )
    {
        RiaLogging::info( "Shutting down gRPC server" );

        // Stop serving callbacks on the completion queue thread before shutting the queue down.
        // Once the queue is shut down, no new request handler may be registered on it. Taking the
        // lock also waits for any dispatch already in progress on that thread.
        {
            std::lock_guard<std::mutex> shutdownLock( m_shutdownMutex );
            m_isShuttingDown = true;
        }

        // See the following link for details on how to shut down a GRPC server
        // https://github.com/grpc/grpc/blob/master/include/grpcpp/server_builder.h#L147

        // Shutdown server and queue
        m_server->Shutdown();
        m_completionQueue->Shutdown();

        // Wait for thread to join. The thread running waitForNextRequest() drains the completion
        // queue before exiting (loops until Next() returns false). Do NOT call Next() again
        // after the thread has joined - calling Next() after it returned false on a shut-down
        // queue is undefined behavior and causes a crash.
        if ( m_thread.joinable() )
        {
            m_thread.join();
        }

        {
            // Create a set of callbacks to be deleted. The same object may be present in both unprocessed and
            // allocated. Delete the callbacks before deleting the server and completion queue to avoid crash.

            std::set<RiaGrpcCallbackInterface*> toBeDeleted;
            for ( auto r : m_unprocessedRequests )
            {
                toBeDeleted.insert( r );
            }

            for ( auto r : m_allocatedCallbakcs )
            {
                toBeDeleted.insert( r );
            }

            m_unprocessedRequests.clear();
            m_allocatedCallbakcs.clear();

            for ( auto r : toBeDeleted )
            {
                delete r;
            }
        }

        // Must destroy server before services
        m_server.reset();
        m_completionQueue.reset();

        // Finally clear services
        m_services.clear();
    }
}

//--------------------------------------------------------------------------------------------------
/// Block and wait for requests from the client from the command queue.
/// The requests are pushed onto the Unprocessed Request queue which are handled in processRequests
//--------------------------------------------------------------------------------------------------
void RiaGrpcServerImpl::waitForNextRequest()
{
    void* tag;
    bool  ok = false;

    while ( m_completionQueue->Next( &tag, &ok ) )
    {
        RiaGrpcCallbackInterface* method = static_cast<RiaGrpcCallbackInterface*>( tag );
        if ( !ok )
        {
            method->setNextCallState( RiaGrpcCallbackInterface::FINISH_REQUEST );
        }

        // Callbacks that read no application state are served here instead of being queued for
        // the main thread. ResInsight processes queued requests from the main thread, so a
        // long-running request (large export, schedule generation, heavy grid computation) or an
        // open progress dialog would otherwise leave liveness requests unanswered for as long as
        // that operation runs, making a busy ResInsight look like a dead one (issue #14683).
        //
        // During shutdown the completion queue is draining and no new request handlers may be
        // registered on it, so fall back to queueing (those callbacks are deleted in quit()).
        if ( method->runsOnServerThread() )
        {
            std::lock_guard<std::mutex> shutdownLock( m_shutdownMutex );
            if ( !m_isShuttingDown )
            {
                process( method );
                continue;
            }
        }

        std::lock_guard<std::mutex> requestLock( m_requestMutex );
        m_unprocessedRequests.push_back( method );
    }
}

//--------------------------------------------------------------------------------------------------
/// The handling of calls pushed onto the command queue. We only get one queued callback per client request.
/// The gRPC calls triggered in the callback will see each callback pushed back onto the command queue.
/// The call state will then determine what the callback should do next.
//--------------------------------------------------------------------------------------------------
void RiaGrpcServerImpl::process( RiaGrpcCallbackInterface* method )
{
    if ( method->callState() == RiaGrpcCallbackInterface::CREATE_HANDLER )
    {
        method->createRequestHandler( m_completionQueue.get() );

        std::lock_guard<std::mutex> callbackLock( m_allocatedCallbackMutex );
        m_allocatedCallbakcs.push_back( method );
    }
    else if ( method->callState() == RiaGrpcCallbackInterface::INIT_REQUEST_STARTED )
    {
        method->onInitRequestStarted();
    }
    else if ( method->callState() == RiaGrpcCallbackInterface::INIT_REQUEST_COMPLETED )
    {
        method->onInitRequestCompleted();
    }
    else if ( method->callState() == RiaGrpcCallbackInterface::PROCESS_REQUEST )
    {
        method->onProcessRequest();
    }
    else
    {
        method->onFinishRequest();

        process( method->createNewFromThis() );

        {
            std::lock_guard<std::mutex> callbackLock( m_allocatedCallbackMutex );
            m_allocatedCallbakcs.remove( method );
        }

        delete method;
    }
}

//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
RiaGrpcServer::RiaGrpcServer( int portNumber )
{
    m_serverImpl = new RiaGrpcServerImpl( portNumber );
}

//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
RiaGrpcServer::~RiaGrpcServer()
{
    delete m_serverImpl;
}

//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
int RiaGrpcServer::portNumber() const
{
    if ( m_serverImpl ) return m_serverImpl->portNumber();

    return 0;
}

//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
bool RiaGrpcServer::isRunning() const
{
    if ( m_serverImpl ) return m_serverImpl->isRunning();

    return false;
}

//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
void RiaGrpcServer::run()
{
    CVF_ASSERT( m_serverImpl );

    m_serverImpl->run();
}

//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
void RiaGrpcServer::runInThread()
{
    CVF_ASSERT( m_serverImpl );

    m_serverImpl->runInThread();
}

//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
bool RiaGrpcServer::s_receivedExitRequest = false;

//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
void RiaGrpcServer::initialize()
{
    CVF_ASSERT( m_serverImpl );

    m_serverImpl->initialize();
}

//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
size_t RiaGrpcServer::processAllQueuedRequests()
{
    CVF_ASSERT( m_serverImpl );

    return m_serverImpl->processAllQueuedRequests();
}

//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
void RiaGrpcServer::quit()
{
    if ( m_serverImpl ) m_serverImpl->quit();
}

//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
int RiaGrpcServer::findAvailablePortNumber( int defaultPortNumber )
{
    int startPort = 50051;

    if ( defaultPortNumber >= 0 && defaultPortNumber < (int)std::numeric_limits<quint16>::max() )
    {
        startPort = defaultPortNumber;
    }

    int endPort = std::min( startPort + 100, (int)std::numeric_limits<quint16>::max() );

    QTcpServer serverTest;
    for ( quint16 port = static_cast<quint16>( startPort ); port <= static_cast<quint16>( endPort ); ++port )
    {
        if ( serverTest.listen( QHostAddress::LocalHost, port ) )
        {
            return static_cast<int>( port );
        }
    }
    return -1;
}

//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
void RiaGrpcServer::setReceivedExitRequest()
{
    RiaLogging::info( "Received Exit Request" );
    s_receivedExitRequest = true;
}

//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
bool RiaGrpcServer::receivedExitRequest()
{
    return s_receivedExitRequest;
}