// Copyright (c) 2017, Xiaomi, Inc.  All rights reserved.
// This source code is licensed under the Apache License Version 2.0, which
// can be found in the LICENSE file in the root directory of this source tree.

#include <dsn/tool-api/task_spec.h>
#include <dsn/tool-api/uri_address.h>

#include <rrdb/rrdb.code.definition.h>
#include "proxy_layer.h"

namespace pegasus {
namespace proxy {

proxy_stub::proxy_stub(const proxy_session::factory &factory, const char *uri)
    : serverlet<proxy_stub>("proxy_stub"), _factory(factory)
{
    dsn::task_spec::get(RPC_CALL_RAW_MESSAGE)->allow_inline = true;
    dsn::task_spec::get(RPC_CALL_RAW_SESSION_DISCONNECT)->allow_inline = true;

    dsn::task_spec::get(dsn::apps::RPC_RRDB_RRDB_PUT_ACK)->allow_inline = true;
    dsn::task_spec::get(dsn::apps::RPC_RRDB_RRDB_MULTI_PUT_ACK)->allow_inline = true;
    dsn::task_spec::get(dsn::apps::RPC_RRDB_RRDB_REMOVE_ACK)->allow_inline = true;
    dsn::task_spec::get(dsn::apps::RPC_RRDB_RRDB_MULTI_REMOVE_ACK)->allow_inline = true;
    dsn::task_spec::get(dsn::apps::RPC_RRDB_RRDB_GET_ACK)->allow_inline = true;
    dsn::task_spec::get(dsn::apps::RPC_RRDB_RRDB_MULTI_GET_ACK)->allow_inline = true;
    dsn::task_spec::get(dsn::apps::RPC_RRDB_RRDB_SORTKEY_COUNT_ACK)->allow_inline = true;
    dsn::task_spec::get(dsn::apps::RPC_RRDB_RRDB_TTL_ACK)->allow_inline = true;
    dsn::task_spec::get(dsn::apps::RPC_RRDB_RRDB_GET_SCANNER_ACK)->allow_inline = true;
    dsn::task_spec::get(dsn::apps::RPC_RRDB_RRDB_SCAN_ACK)->allow_inline = true;
    dsn::task_spec::get(dsn::apps::RPC_RRDB_RRDB_CLEAR_SCANNER_ACK)->allow_inline = true;

    _uri_address.assign_uri(uri);
    open_service();
}

proxy_stub::~proxy_stub() {}

void proxy_stub::on_rpc_request(dsn_message_t request)
{
    ::dsn::rpc_address source = dsn_msg_from_address(request);
    std::shared_ptr<proxy_session> ps;
    {
        ::dsn::service::zauto_read_lock l(_lock);
        auto it = _sessions.find(source);
        if (it != _sessions.end()) {
            ps = it->second;
        }
    }
    if (nullptr == ps) {
        ::dsn::service::zauto_write_lock l(_lock);
        auto it = _sessions.find(source);
        if (it != _sessions.end()) {
            ps = it->second;
        } else {
            ps = _factory(this, request);
            _sessions.emplace(source, ps);
        }
    }

    ps->on_recv_request(request);
}

void proxy_stub::on_recv_remove_session_request(dsn_message_t request)
{
    ::dsn::rpc_address source = dsn_msg_from_address(request);
    std::shared_ptr<proxy_session> ps = remove_session(source);
    if (ps != nullptr) {
        ps->on_remove_session();
    }
}

std::shared_ptr<proxy_session> proxy_stub::remove_session(dsn::rpc_address remote_address)
{
    ::dsn::service::zauto_write_lock l(_lock);
    auto iter = _sessions.find(remote_address);
    if (iter == _sessions.end()) {
        dwarn("%s has been removed from proxy stub", remote_address.to_string());
        return nullptr;
    }
    ddebug("remove %s from proxy stub", remote_address.to_string());
    std::shared_ptr<proxy_session> result = std::move(iter->second);
    _sessions.erase(iter);
    return result;
}

proxy_session::proxy_session(proxy_stub *op, dsn_message_t first_msg)
    : stub(op), is_session_reset(false), backup_one_request(first_msg)
{
    dassert(first_msg != nullptr, "null msg when create session");
    dsn_msg_add_ref(backup_one_request);

    remote_address = dsn_msg_from_address(backup_one_request);
    dassert(remote_address.type() == HOST_TYPE_IPV4,
            "invalid rpc_address type, type = %d",
            (int)remote_address.type());
}

proxy_session::~proxy_session()
{
    dsn_msg_release_ref(backup_one_request);
    ddebug("proxy session %s destroyed", remote_address.to_string());
}

void proxy_session::on_recv_request(dsn_message_t msg)
{
    // NOTICE:
    // 1. in the implementation of "parse", the msg may add_ref & release_ref.
    //    so if the ref_count of msg is 0 before call "parse", the msg may be released already
    //    after "parse" returns. so please take care when you want to
    //    use "msg" after call "parse"
    //
    // 2. as "on_recv_request" won't be called concurrently, it's not necessary to call
    //    "parse" with a lock. a subclass may implement a lock inside parse if necessary
    if (!parse(msg)) {
        derror("%s: got invalid message, try to remove proxy session from proxy stub",
               remote_address.to_string());
        stub->remove_session(remote_address);

        derror("close the rpc session %s", remote_address.to_string());
        ((dsn::message_ex *)backup_one_request)->io_session->close();
    }
}

void proxy_session::on_remove_session() { is_session_reset.store(true); }

dsn_message_t proxy_session::create_response()
{
    return dsn_msg_create_response(backup_one_request);
}
}
} // namespace
