#pragma once

#include <balancer/serval/core/buffer.h>
#include <balancer/serval/core/config.h>
#include <balancer/serval/mod/apphost/apphost.ev.pb.h>
#include <balancer/serval/mod/apphost/stream.pb.h>

#include <library/cpp/json/json_reader.h>
#include <library/cpp/json/json_writer.h>
#include <library/cpp/json/yson/json2yson.h>
#include <library/cpp/string_utils/base64/base64.h>

#include <apphost/lib/compression/compression.h>
#include <apphost/lib/compression/compression_codecs.h>

namespace NSv::NAppHost {
    static inline TStringBuf EncodeChunk(const NProtoBuf::MessageLite& msg, bool grpc, TString& out) {
        auto delta = msg.ByteSizeLong();
        const ui8 prefix[] = {0, (ui8)(delta >> 24), (ui8)(delta >> 16), (ui8)(delta >> 8), (ui8)delta};
        if (grpc) {
            out.append((const char*)prefix, (const char*)(prefix + 5));
        }
        Y_PROTOBUF_SUPPRESS_NODISCARD msg.AppendPartialToString(&out);
        return out;
    }

    template <typename T>
    static inline TMaybe<T> DecodeChunk(NSv::IStream& stream, bool grpc, bool complete = false) {
        auto ret = MakeMaybe<T>();
        while (!stream.AtEnd() && (complete || !ret->ByteSizeLong())) {
            size_t size = -1;
            if (grpc) {
                auto head = NSv::ReadFrom(stream, 5);
                if (!head MUN_RETHROW) {
                    return {};
                }
                if (head->size() < 5) {
                    break; // actually a protocol error, but whatever
                }
                size = (ui32)(ui8)(*head)[1] << 24 | (ui32)(ui8)(*head)[2] << 16
                     | (ui32)(ui8)(*head)[3] << 8  | (ui32)(ui8)(*head)[4];
            }
            auto payload = NSv::ReadFrom(stream, size);
            if (!payload MUN_RETHROW) {
                return {};
            }
            google::protobuf::io::CodedInputStream input(reinterpret_cast<const ui8*>(payload->data()), payload->size());
            Y_PROTOBUF_SUPPRESS_NODISCARD ret->MergePartialFromCodedStream(&input);
        }
        return ret;
    }

    class TStream : public NSv::TConstRequestStream {
    public:
        TStream(THead head, bool grpc, NSv::TLogFrame log)
            : TConstRequestStream(std::move(head), {}, std::move(log))
            , Grpc(grpc)
        {
        }

        TMaybe<TStringBuf> Read() noexcept override {
            while (!Done && NSv::TConstRequestStream::AtEnd() && !Merged.ByteSizeLong()) {
                if (!HaveData.wait() MUN_RETHROW) {
                    return {};
                }
            }
            if (NSv::TConstRequestStream::AtEnd() && Merged.ByteSizeLong()) {
                Log().Push<NSv::NEv::TAppHostSendChunk>(Merged.ByteSizeLong());
                NSv::TConstRequestStream::UpdatePayload(EncodeChunk(Merged, Grpc, Send));
                Merged.Clear();
            }
            return NSv::TConstRequestStream::Read();
        }

        TMaybe<TRequest> ReadDirect(bool all = false) noexcept {
            while (!Done && (all || !Merged.ByteSizeLong())) {
                if (!HaveData.wait() MUN_RETHROW) {
                    return {};
                }
            }
            return std::move(Merged);
        }

        bool AtEnd() const noexcept override {
            return Done && NSv::TConstRequestStream::AtEnd() && !Merged.ByteSizeLong();
        }

        TRequest& GetBuffer() noexcept {
            HaveData.wake();
            return Merged;
        }

        void CloseBuffer() noexcept {
            Done = true;
            HaveData.wake();
        }

        bool WriteHeadEx(THead&) noexcept override {
            return true; // TODO handle non-200 codes?
        }

        bool WriteEx(TStringBuf data) noexcept override {
            Recv += data;
            while (Grpc && Recv.size() >= 5) {
                ui32 s = (ui32)(ui8)Recv[1] << 24 | (ui32)(ui8)Recv[2] << 16
                       | (ui32)(ui8)Recv[3] << 8  | (ui32)(ui8)Recv[4];
                if (s > Recv.size() - 5) {
                    break;
                }
                TResponse chunk;
                Y_PROTOBUF_SUPPRESS_NODISCARD chunk.ParsePartialFromArray(Recv.data() + 5, s);
                if (!OnChunk(std::move(chunk))) {
                    return false;
                }
                Recv.erase(0, s + 5);
            }
            return true;
        }

        bool CloseEx(THeaderVector&) noexcept override {
            if (!Grpc) {
                TResponse chunk;
                Y_PROTOBUF_SUPPRESS_NODISCARD chunk.ParsePartialFromArray(Recv.data(), Recv.size());
                if (!OnChunk(std::move(chunk))) {
                    return false;
                }
            }
            return OnClose();
        }

        virtual bool OnChunk(TResponse&&) noexcept = 0;
        virtual bool OnClose() noexcept = 0;

    private:
        TString Recv;
        TString Send;
        TRequest Merged;
        cone::event HaveData;
        bool Done = false;
        bool Grpc = false;
    };

    template <typename TChunk, typename TClose>
    class TLambdaStream : public TStream, TChunk, TClose {
    public:
        TLambdaStream(THead head, bool grpc, TChunk chunk, TClose close, NSv::TLogFrame log)
            : TStream(std::move(head), grpc, std::move(log))
            , TChunk(std::move(chunk))
            , TClose(std::move(close))
        {
        }

        bool OnChunk(TResponse&& chunk) noexcept override {
            return TChunk::operator()(std::move(chunk));
        }
        bool OnClose() noexcept override {
            return TClose::operator()();
        }
    };

    template <typename TChunk, typename TClose>
    static inline auto Stream(TStringBuf path, bool grpc, TStringBuf hint, const TVector<TString>& backends,
                              TChunk chunk = {}, TClose close = {}, NSv::TLogFrame log = {}) {
        NSv::THead head{"POST", path, {{":authority", "unknown"}, {":scheme", "http"}}};
        if (grpc) {
            head.PathWithQuery = "/NAppHostProtocol.TServant/Invoke";
            head.emplace("content-type", "application/grpc");
            head.emplace("te", "trailers");
        }
        if (hint) {
            head.emplace("x-apphost-hint", hint);
        }
        for (const auto& backend : backends) {
            head.emplace("x-apphost-srcrwr", backend);
        }
        return std::make_shared<TLambdaStream<TChunk, TClose>>(std::move(head), grpc,
            std::move(chunk), std::move(close), std::move(log));
    }

    struct TInterface {
    public:
        TInterface(NSv::IStreamPtr parent)
            : P(std::move(parent))
            , Casted(dynamic_cast<TStream*>(P.get()))
        {
            auto head = P->Head();
            auto type = head->find("content-type");
            Grpc = type != head->end() && type->second.StartsWith("application/grpc");
        }

        bool IsRoot() const {
            return !Casted;
        }

        bool StreamOut() const {
            return Casted || (Grpc && P->Head()->Path() == "/NAppHostProtocol.TServant/InvokeEx");
        }

        NSv::TLogFrame& Log() {
            return P->Log();
        }

        TMaybe<TRequest> ReadAll() {
            return Casted ? Casted->ReadDirect(true) : DecodeChunk<TRequest>(*P, Grpc, true);
        }

        TMaybe<TRequest> Read() {
            return Casted ? Casted->ReadDirect() : DecodeChunk<TRequest>(*P, Grpc);
        }

        bool Write(TResponse&& chunk) {
            if (Casted) {
                return Casted->OnChunk(std::move(chunk));
            }
            if (!SentHead) {
                if (Grpc ? !P->WriteHead({200, {{"content-type", "application/grpc"}}}) : !P->WriteHead(200) MUN_RETHROW) {
                    return false;
                }
            }
            SentHead = true;
            // FIXME can only write as separate chunks if InvokeEx was called, not Invoke;
            //       should concatenate into one object in the latter case.
            TString out;
            return P->Write(EncodeChunk(chunk, Grpc, out));
        }

        bool Close() {
            return Casted ? Casted->OnClose() && !mun_error(EREQDONE, "stream closed")
                 : Grpc ? P->Close({{"grpc-status", "0"}}) : P->Close();
        }

    private:
        NSv::IStreamPtr P;
        TStream* Casted;
        bool Grpc = false;
        bool SentHead = false;
    };

    template <typename F>
    static inline void DecodeJson(const TAnswer& answer, F&& f) {
        auto data = ::NAppHost::NCompression::Decode(answer.GetData());
        auto view = TStringBuf(data);
        NJson::TJsonValue value;
        if (view.SkipPrefix("p_")) {
            return;
        }
        if (view.SkipPrefix("y_")) {
            NJson2Yson::DeserializeYsonAsJsonValue(view, &value);
        } else {
            NJson::ReadJsonFastTree(view, &value);
        }
        if (!value.IsArray()) {
            f(std::move(value));
        } else {
            for (auto& subvalue : value.GetArraySafe()) {
                f(std::move(subvalue));
            }
        }
    }
}
