#include "socket.h"

#include <library/cpp/coroutine/engine/network.h>

#include <balancer/kernel/memory/split.h>

#include <util/generic/scope.h>

#include <utility>
#include <cstdlib>

namespace NSrvKernel {

    struct TSocketIn::TEdgeTriggeredEvent {
        TSock& Sock;
        bool NeedWait = true;
        int Err = 0;
        TPollerEvent Event;
    public:
        TEdgeTriggeredEvent(TSock& sock, int what)
            : Sock(sock)
            , Event(*sock.Executor(), *sock.Socket(), what | CONT_POLL_EDGE_TRIGGERED)
        {
            Event.SetUnwaitedEventCallback([this](int status) mutable {
                Err = status;
                NeedWait = false;
            });
        }

        TError MayWaitForEvent(TInstant deadline) {
            if (NeedWait) {
                if (Err) {
                    return Y_MAKE_ERROR(TSystemError(Err) << "io error");
                }

                NeedWait = false;
                const int waitResult = Event.Wait(Sock.Executor()->Running(), deadline);

                if (waitResult == ETIMEDOUT || waitResult == ECANCELED || waitResult == EINPROGRESS) {
                    NeedWait = true; //allow to wait more next time
                    return Y_MAKE_ERROR(TSystemError{waitResult} << "io error");
                }

                Err = waitResult;
            }

            return {};
        }
    };

    bool UseEdgeTriggered(TSock& sock, EPollMode pollMode) {
        const bool turnOff = false;
        return !turnOff && pollMode == PM_EDGE_TRIGGERED && sock.Executor()->Poller()->PollEngine() == EContPoller::Epoll;
    }

    TSocketIn::TSocketIn(TSock& sock, EPollMode pollMode) noexcept
        : Sock_(sock)
        , EdgeTriggered_(UseEdgeTriggered(sock, pollMode))
        , ReadyRead_(EdgeTriggered_ ? MakeHolder<TEdgeTriggeredEvent>(sock, CONT_POLL_READ) : nullptr)
        , ReadyEvent_(*sock.Executor(), *sock.Socket(), CONT_POLL_READ | (EdgeTriggered_ ? CONT_POLL_EDGE_TRIGGERED : 0))
    {
        ReadyEvent_.SetUnwaitedEventCallback([this](int status){
            if (ReadyCallback_) {
                ReadyCallback_(status);
            }
        });
    }

    TSocketIn::~TSocketIn() {
    }

    void TSocketIn::SetOnReadyCallback(std::function<void(int)> callback) noexcept {
        ReadyCallback_ = callback;
    }

    TError TSocketIn::DoRecv(TChunkList& lst, TInstant deadline) noexcept {
        TChunkPtr chunk = NextChunk();
        size_t readed = 0;

        if (!ReadyRead_) {
            const auto status = NCoro::ReadD(Sock_.Executor()->Running(),
                                             *Sock_.Socket(), chunk->Data(), chunk->Length(), deadline);

            Y_REQUIRE(!status.Status(), TSystemError{status.Status()} << "io error");

            readed = status.Processed();
        } else {
            while (true) {
                Y_PROPAGATE_ERROR(ReadyRead_->MayWaitForEvent(deadline));

                const ssize_t readResult = read(*Sock_.Socket(), chunk->Data(), chunk->Length());

                if (readResult >= 0) {
                    readed = readResult;
                    break;
                } else {
                    const int err = LastSystemError();

                    if (err == EAGAIN || err == EWOULDBLOCK) {
                        ReadyRead_->NeedWait = true;
                        continue;
                    } else if (err == EINTR) {
                        ReadyRead_->NeedWait = false;
                        continue;
                    } else if (ReadyRead_->Err) {
                        return Y_MAKE_ERROR(TSystemError(ReadyRead_->Err) << "io error");
                    } else {
                        return Y_MAKE_ERROR(TSystemError(err) << "io error");
                    }
                }
            }
        }

        if (readed) {
            const auto chunkLen = chunk->Length();
            if (readed == chunkLen) {
                //fill full chunk
                lst.Push(std::move(chunk));
            } else {
                const size_t left = chunkLen - readed;

                if (left < readed) {
                    //too small data chunk left, better do not reuse it
                    chunk->Shrink(readed);
                    lst.Push(std::move(chunk));
                } else {
                    //something left for future use
                    lst.Push(chunk->SubChunk(readed));
                    chunk->Skip(readed);
                    L_ = std::move(chunk);
                }
            }
        }
        return {};
    }

    TError TSocketIn::WaitRdhup() noexcept {
        int ret = 0;
        if (EdgeTriggered_) {
            TPollerEvent rdHup(*Sock_.Executor(), *Sock_.Socket(), CONT_POLL_RDHUP | CONT_POLL_EDGE_TRIGGERED);
            ret = rdHup.Wait(Sock_.Executor()->Running(), TInstant::Max());
        } else {
            ret = NCoro::PollI(Sock_.Executor()->Running(), *Sock_.Socket(), CONT_POLL_RDHUP);
        }

        if (ret != 0) {
            return Y_MAKE_ERROR(TSystemError{ret} << "wait rdhup poll error");
        } else {
            return {};
        }
    }

    TSocketOut::TSocketOut(TSock& sock, size_t maxInQueue, EPollMode pollMode) noexcept
        : IPollEvent(*sock.Socket(), CONT_POLL_WRITE | (UseEdgeTriggered(sock, pollMode) ? CONT_POLL_EDGE_TRIGGERED : 0))
        , EdgeTriggered_(UseEdgeTriggered(sock, pollMode))
        , Sock_(sock)
        , MaxInQueue_(maxInQueue)
    {
    }

    TSocketOut::~TSocketOut() {
        Sock_.Executor()->Poller()->Remove(this);
        InPoller_ = false;
    }

    TError TSocketOut::Flush(size_t keep, TInstant deadline) noexcept {
        while (W_.size() > keep && !(Sock_.Executor()->Running()->Cancelled())) {
            Y_REQUIRE(!Err_, TSystemError{Err_} << "send failed");

            Wake_ = Sock_.Executor()->Running();
            int res = Wake_->SleepD(deadline);
            Wake_ = nullptr;

            Y_REQUIRE(res != ETIMEDOUT, TSystemError{ETIMEDOUT} << "send failed");
        }

        W_.MakeOwned();

        return {};
    }

    TError TSocketOut::DoSend(TChunkList in, TInstant deadline) noexcept {
        Y_REQUIRE(!Err_, TSystemError{Err_} << "send failed");

        if (in.Empty()) {
            return {};
        }

        W_.Append(std::move(in));
        if (EdgeTriggered_) {
            Sock_.Executor()->ScheduleUserEvent(this);
        } else {
            Sock_.Executor()->Poller()->Schedule(this);
            InPoller_ = true;
        }
        return Flush(MaxInQueue_, deadline);
    }

    void TSocketOut::OnPollEvent(int status) noexcept {
        using TPart = IOutputStream::TPart;

        if (Wake_) {
            Wake_->ReSchedule();
        }

        if (status) {
            Err_ = status;
            Sock_.Executor()->Poller()->Remove(this);
            InPoller_ = false;

            return;
        }

        while (true) {
            THolder<void, TFree> parts(malloc(W_.ChunksCount() * sizeof(TPart)));

            size_t bytesToWrite = 0;
            TPart* pb = reinterpret_cast<TPart*>(parts.Get());
            TPart* pe = std::transform(W_.ChunksBegin(), W_.ChunksEnd(), pb, [&](const auto& chunk) {
                const size_t chunkLen = chunk.Length();
                bytesToWrite += chunkLen;
                return TPart(chunk.Data(), chunkLen);
            });

            TContIOVector vec(pb, pe - pb);

            const ssize_t res = writev(*Sock_.Socket(), (const iovec*) vec.Parts(), Min(IOV_MAX, (int) vec.Count()));

            if (res < 0) {
                const int err = LastSystemError();

                if (err != EAGAIN && err != EWOULDBLOCK) {
                    Err_ = err;
                    Sock_.Executor()->Poller()->Remove(this);
                    InPoller_ = false;
                } else if (!EdgeTriggered_) {
                    Y_FAIL("fd isn't ready for write");
                } else if (!InPoller_) {
                    Sock_.Executor()->Poller()->Schedule(this);
                    InPoller_ = true;
                }

                return;
            } else {
                if ((size_t) res == bytesToWrite) {
                    //full write complete, nothing to wait
                    W_.Clear();
                    if (!EdgeTriggered_) {
                        Sock_.Executor()->Poller()->Remove(this);
                        InPoller_ = false;
                    }
                    return;
                } else {
                    //some data left
                    W_.Skip((size_t) res);
                    if (EdgeTriggered_) {
                        continue;
                    } else {
                        return;
                    }
                }
            }
        }
    }
}
