#pragma once

#include <balancer/kernel/custom_io/iterator.h>
#include <balancer/kernel/helpers/errors.h>
#include <balancer/kernel/io/iobase.h>
#include <balancer/kernel/memory/chunks.h>

#include <util/generic/strbuf.h>
#include <util/stream/format.h>
#include <util/string/hex.h>
#include <util/string/builder.h>


namespace NSrvKernel {
    // FIXME: decoder should be fixed to comply with section 3.6.1 of standard
    class TChunkedTransferDecoder : public TStreamIterator, public IIoInput {
    public:
        TChunkedTransferDecoder(IIoInput* slave) noexcept
            : TStreamIterator(slave)
        {}

        TError DoRecv(TChunkList& lst, TInstant deadline = TInstant::Max()) noexcept override {
            if (Finished_) {
                return {};
            }

            Y_TRY(TError, error) {
                if (!L_) {
                    Y_PROPAGATE_ERROR(ParseNextChunkLength(deadline).AssignTo(L_));
                    Finished_ = !L_;
                }

                if (Finished_) {
                    return SkipCRLF(deadline);
                }

                Y_PROPAGATE_ERROR(Next(L_, lst, deadline));
                L_ -= lst.size();

                if (!L_) {
                    Y_PROPAGATE_ERROR(SkipCRLF(deadline)); // Skip trailing \r\n
                }

                Y_REQUIRE(!lst.Empty(),
                    (THttpError{400, "chunked transfer cropped data"}) << "not all data received");
                return {};
            } Y_CATCH {
                Y_REQUIRE(!error.GetAs<TStreamError>(),
                    (THttpError{400, "chunked transfer cropped data"}) << "not all data received");
                return error;
            }
            return {};
        }

    private:
        TErrorOr<size_t> ParseNextChunkLength(TInstant deadline) noexcept {
            size_t ret = 0;
            size_t logLen = 0;

            while (true) {
                char ch;
                Y_PROPAGATE_ERROR(Next(deadline).AssignTo(ch));

                switch (ch) {
                    case '\r':
                        Y_PROPAGATE_ERROR(Next(deadline).AssignTo(ch));

                        Y_REQUIRE(ch == '\n' && logLen,
                            (THttpError{400, "invalid symbol in chunk length"}) << "invalid symbol in chunk length");

                        return ret;

                    case '\n':
                        Y_REQUIRE(logLen,
                            (THttpError{400, "invalid symbol in chunk length"}) << "invalid symbol in chunk length");

                        return ret;

                    default:
                        Y_REQUIRE(logLen < sizeof(ret) * 2,
                            (THttpError{400, "too long chunk length"}) << "too long chunk length");

                        Y_REQUIRE(!!((ch >= '0') ^ (ch > '9') ^ (ch >= 'a') ^ (ch > 'f') ^ (ch >= 'A') ^ (ch > 'F')),
                            (THttpError{400, "invalid symbol in chunk length"}) << "invalid symbol in chunk length");

                        try {
                            ret = ret * 16 + Char2Digit(ch);
                        } Y_TRY_STORE(yexception);
                        logLen += 1;

                        break;
                }
            }
        }

        TError SkipCRLF(TInstant deadline) noexcept {
            char ch;
            Y_PROPAGATE_ERROR(Next(deadline).AssignTo(ch));

            if (ch == '\r') {
                Y_PROPAGATE_ERROR(Next(deadline).AssignTo(ch));

                if (ch == '\n') {
                    return {};
                }
            }

            return Y_MAKE_ERROR(THttpError{(THttpError{400, "invalid symbol"}) << "invalid symbol"});
        }

    private:
        size_t L_ = 0;
        bool Finished_ = false;
    };


    class TChunkedTransferEncoder : public IIoOutput {
    public:
        TChunkedTransferEncoder(IIoOutput* slave) noexcept
            : S_(slave)
        {}

        TError DoSend(TChunkList lst, TInstant deadline) noexcept override {
            if (!lst.Empty()) {
                TStringBuilder sb;
                sb.reserve(16 + 2);
                sb << Hex(lst.size(), 0) << "\r\n";
                lst.PushFront(NewChunk(std::move(sb)));
                lst.Push("\r\n");

                return S_->Send(std::move(lst), deadline);
            } else {
                lst.Push("0\r\n\r\n");
                Y_PROPAGATE_ERROR(S_->Send(std::move(lst), deadline));

                return S_->SendEof(deadline);
            }
        }

    private:
        IIoOutput* const S_ = nullptr;
    };
};
