#include "master_main.h"
#include "connection.h"
#include "connection_callback.h"
#include "master_process.h"
#include "initializer_message.h"

#include <balancer/kernel/http/parser/httpdecoder.h>
#include <balancer/kernel/http/parser/output.h>
#include <balancer/kernel/http/parser/response_builder.h>
#include <balancer/kernel/net/sockops.h>
#include <balancer/kernel/process/children/worker_process.h>
#include <kernel/p0f/p0f.h>

namespace {
    NP0f::TP0fOrError GetP0f(const NSrvKernel::TAddrHolder& remoteAddr, const NSrvKernel::TAddrHolder& localAddr) {
#ifdef _linux_
        const sockaddr* saddr = remoteAddr.Addr()->Addr();
        switch (saddr->sa_family) {
            case AF_INET:
                return NP0f::ExtractFingerprint(((const sockaddr_in*)saddr)->sin_addr, remoteAddr.Port(), localAddr.Port());
            case AF_INET6:
                return NP0f::ExtractFingerprint(((const sockaddr_in6*)saddr)->sin6_addr, remoteAddr.Port(), localAddr.Port());
        }
#else
        Y_UNUSED(remoteAddr);
        Y_UNUSED(localAddr);
#endif

        return {};
    }
}

namespace NSrvKernel::NProcessCore {
    TBaseConnection::TBaseConnection(const TAcceptFull& a)
        : S_(a.S->Release())
        , RemoteAddrOwner_(a.Remote)
        , LocalAddrOwner_(a.Local)
        , RemoteAddress_(&RemoteAddrOwner_)
        , LocalAddress_(&LocalAddrOwner_)
    {
    }

    TCont* TBaseConnection::GetCont() noexcept {
        return C_;
    }

    void TBaseConnection::Run(TCont* cont) {
        try {
            Y_UNUSED(DoRun(cont));
        } catch (...) {
        }
        delete this;
    }

    class TConnection::TPriorityIo final: public IIoInput, public IHttpOutput {
    public:
        TPriorityIo(TConnection& parent, TSocketIo& slave) noexcept
            : Parent_(parent)
            , Slave_(slave)
        {
        }

    private:
        TError DoSendHead(TResponse&&, const bool, TInstant) override {
            Y_FAIL("TPriorityIo doesn't know what to do in SendHead.");
        }

        TError DoSend(TChunkList lst, TInstant deadline) noexcept override {
            Y_PROPAGATE_ERROR(Slave_.Out().Send(std::move(lst), deadline));
            Parent_.OnConnIo();
            return {};
        }

        TError DoSendTrailers(THeaders&&, TInstant) override {
            Y_FAIL("TPriorityIo doesn't know what to do in SendTrailers.");
        }

        TError DoRecv(TChunkList& lst, TInstant deadline) noexcept override {
            Y_PROPAGATE_ERROR(Slave_.In().Recv(lst, deadline));
            Parent_.OnConnIo();
            return {};
        }

    private:
        TConnection& Parent_;
        TSocketIo& Slave_;
    };

    TConnection::TConnection(IConnectionCallback* parent, const IModule& module, const TAcceptFull& a, TCpuLimiter* cpuLimiter)
        : TBaseConnection(a)
        , Parent_(parent)
        , Module_(module)
        , CpuLimiter_(cpuLimiter)
    {
        C_ = Parent_->WorkerProcess_->WorkerExecutor_->Executor()
                 .Create<TBaseConnection, &TBaseConnection::Run>(this, "connection");

        Acquire(Parent_);
    }

    TConnection::~TConnection() {
        Release(Parent_);
    }

    void TConnection::Acquire(IConnectionCallback* task) noexcept {
        ++task->InProgr_;
        if (task->TotalConnectionInProgress_) {
            AtomicIncrement(*task->TotalConnectionInProgress_);
        }
        ++task->WorkerProcess_->MainStats_->TcpConns;
        ++task->WorkerProcess_->MainStats_->TcpConnsOpen;
    }

    void TConnection::Release(IConnectionCallback* task) noexcept {
        --task->InProgr_;
        if (task->TotalConnectionInProgress_) {
            AtomicDecrement(*task->TotalConnectionInProgress_);
        }
        --task->WorkerProcess_->MainStats_->TcpConns;
        ++task->WorkerProcess_->MainStats_->TcpConnsClose;
    }

    void TConnection::Cancel() noexcept {
        Parent_->AddCanceled(LocalAddrOwner_);
        C_->Cancel();

        if (C_) {
            C_->Executor()->Running()->Join(C_, TInstant());
        }
    }

    void TConnection::OnConnIo() noexcept {
        Parent_->UpPriority(this);
    }

    TError TConnection::DoRun(TCont* cont) {
        ++Parent_->WorkerProcess_->MainStats_->ConnectionsCounter();

        const auto& config = Parent_->WorkerProcess_->MainConfig();

        Y_PROPAGATE_ERROR(EnableNoDelay(S_));
        Y_PROPAGATE_ERROR(SetSockBufSize(S_, config.SockBufSize));
        Y_PROPAGATE_ERROR(EnableKeepalive(S_, config.TcpKeepalive));
        if (config.TcpCongestionControlName) {
            Y_PROPAGATE_ERROR(SetCongestionControl(S_, config.TcpCongestionControlName));
        }
        if (config.TcpNotsentLowat.Defined()) {
            Y_PROPAGATE_ERROR(SetNotsentLowat(S_, *config.TcpNotsentLowat.Get()));
        }

        TSocketIo io{&S_, cont->Executor(), config.Buffer, PM_EDGE_TRIGGERED};
        TPriorityIo prioIo{*this, io};

        TTcpConnProps tcpProps(*Parent_->WorkerProcess_, RemoteAddress_, LocalAddress_, &S_);
        tcpProps.CpuLimiter = CpuLimiter_;
        tcpProps.TcpRstOnError = config.TcpRstOnError;

        if (config.P0fEnabled) {
            auto p0f = GetP0f(RemoteAddress_, LocalAddress_);
            if (p0f.Value) {
                tcpProps.P0f = std::move(p0f.Value);
            } else if (p0f.Error) {
                Parent_->WorkerProcess_->Log_ << "failed to get p0f: " << *p0f.Error << Endl;
            }
        }

        TConnProps props(tcpProps, TInstant::Now(), RandomNumber<ui64>(), &io);
        TConnDescr connDescr(prioIo, prioIo, props);

        TError moduleRunError{};

        Y_TRY(TError, error) {
            return Module_.Run(connDescr);
        }
        Y_CATCH {
            moduleRunError = std::move(error);
        }

        LOG_ERROR(TLOG_INFO, connDescr, "connection finished: Connection lasted for " << Now() - connDescr.Properties->Start);
        Y_PROPAGATE_ERROR(io.Out().Flush(TSocketOut::FlushAll, TInstant::Max()));

        if (moduleRunError && tcpProps.TcpRstOnError) {
            Y_UNUSED(EnableRstOnClose(S_));
        }
        return moduleRunError;
    }

    TStatsConnection::TStatsConnection(TStatsCallback* parent, const TAcceptFull& a, TSharedStatsManager& statsManager, TMasterProcess& masterProcess)
        : TBaseConnection(a)
        , StatsManager_(statsManager)
        , MasterProcess_(masterProcess)
        , Parent_(parent)
        , JsonRegexp_("/json/(.*?)(/.*)?")
    {
        C_ = Parent_->Executor.Create<TBaseConnection, &TBaseConnection::Run>(this, "stats connection");

        Acquire(Parent_);
    }

    TStatsConnection::~TStatsConnection() {
        Release(Parent_);
    }

    TError TStatsConnection::WriteResponse(TCont* cont, IIoInput* input, IIoOutput* output) noexcept {
        TFromClientDecoder decoder{input};
        TRequest request;
        Y_PROPAGATE_ERROR(decoder.ReadRequest(request, TInstant::Max()));

        TStringBuf path = request.RequestLine().Path.AsStringBuf();

        int status = HTTP_BAD_REQUEST;
        TChunksOutputStream out;

        if (request.RequestLine().Method == EMethod::GET && path == "/unistat") {
            StatsManager_.WriteResponse(cont, TSharedStatsManager::TResponseType::Unistat, &out);
            status = HTTP_OK;
        } else if (request.RequestLine().Method == EMethod::GET && path == "/solomon") {
            StatsManager_.WriteResponse(cont, TSharedStatsManager::TResponseType::Solomon, &out);
            status = HTTP_OK;
        } else if (request.RequestLine().Method == EMethod::GET && path.StartsWith("/json/")) {
            TVector<TStringBuf> groups;
            if (JsonRegexp_.Extract(path, &groups, false) && groups.size() == 3) {
                TEventData eventData{groups[1], out};
                MasterProcess_.ProcessJsonEvent(eventData);
                status = HTTP_OK;
            }
        }

        THttpOutput httpOutput{output, 1, false, false, false};

        Parent_->Log << "[" << NAddr::PrintHostAndPort(RemoteAddrOwner_) << "] stats request: " << path << " " << status << Endl;
        auto content = std::move(out.Chunks());
        TResponse response = BuildResponse().Version11().Code(status).ContentLength(content.size());
        Y_PROPAGATE_ERROR(httpOutput.SendHead(std::move(response), true, TInstant::Max()));
        if (!content.Empty()) {
            Y_PROPAGATE_ERROR(httpOutput.Send(std::move(content), TInstant::Max()));
        }
        return httpOutput.SendEof(TInstant::Max());
    }

    TError TStatsConnection::DoRun(TCont* cont) {
        ++Parent_->ConnectionsCounter;
        Y_PROPAGATE_ERROR(EnableNoDelay(S_));
        TSocketIo io{&S_, cont->Executor(), Parent_->BufferSize, PM_EDGE_TRIGGERED};

        Y_TRY(TError, error) {
            Y_PROPAGATE_ERROR(WriteResponse(cont, &io.In(), &io.Out()));
            return io.Out().Flush(TSocketOut::FlushAll, TInstant::Max());
        }
        Y_CATCH {
            Parent_->Log << "[" << NAddr::PrintHostAndPort(RemoteAddrOwner_) << "] stats error: " << error->what() << Endl;
        }
        return {};
    }

    void TStatsConnection::Acquire(TStatsCallback* cb) noexcept {
        ++cb->InProgr;
    }

    void TStatsConnection::Release(TStatsCallback* cb) noexcept {
        --cb->InProgr;
    }

    void TStatsConnection::Cancel() noexcept {
        C_->Cancel();

        if (C_) {
            C_->Executor()->Running()->Join(C_, TInstant());
        }
    }

    TAdminConnection::TAdminConnection(TAdminCallback* parent, const TAcceptFull& a, TAdminManager& adminManager)
        : TBaseConnection(a)
        , AdminManager_(adminManager)
        , Parent_(parent)
    {
        C_ = Parent_->Executor.Create<TBaseConnection, &TBaseConnection::Run>(this, "connection");

        Acquire(Parent_);
    }

    TAdminConnection::~TAdminConnection() {
        Release(Parent_);
    }

    TError TAdminConnection::DoRun(TCont* cont) {
        ++Parent_->ConnectionsCounter;
        Y_PROPAGATE_ERROR(EnableNoDelay(S_));
        TSocketIo io(&S_, cont->Executor(), Parent_->BufferSize, PM_EDGE_TRIGGERED);

        Y_TRY(TError, error) {
            Y_PROPAGATE_ERROR(AdminManager_.MakeAdminAnswer(&io.In(), &io.Out()))
            return io.Out().Flush(TSocketOut::FlushAll, TInstant::Max());
        }
        Y_CATCH {
            Parent_->Log << "[" << NAddr::PrintHostAndPort(RemoteAddrOwner_) << "] admin error: " << error->what() << Endl;
        }

        return {};
    }

    void TAdminConnection::Acquire(TAdminCallback* cb) noexcept {
        ++cb->InProgr;
    }

    void TAdminConnection::Release(TAdminCallback* cb) noexcept {
        --cb->InProgr;
    }

    void TAdminConnection::Cancel() noexcept {
        C_->Cancel();

        if (C_) {
            C_->Executor()->Running()->Join(C_, TInstant());
        }
    }
}
