#pragma once

#include "address.h"
#include "function.h"
#include "io.h"
#include "log.h"
#include "unistat.h"

#include <array>

#include <balancer/serval/contrib/cone/cone.hh>

#include <library/cpp/containers/sorted_vector/sorted_vector.h>

#include <util/generic/maybe.h>
#include <util/generic/strbuf.h>

enum {
    // Returned by `Close` on a server-side stream to signal that you should stop
    // processing the request and begin unwinding the handler's stack.
    EREQDONE = 30169,
};

namespace NSv {
    class IO;

    struct THeader : std::pair<TStringBuf, TStringBuf> {
    public:
        template <typename... Args>
        THeader(Args&&... args) noexcept
            : std::pair<TStringBuf, TStringBuf>(std::forward<Args>(args)...)
        {
            Y_UNUSED(Unused_);
        }

        struct TGetName {
            TStringBuf operator()(const THeader& h) const noexcept {
                return h.first;
            }
        };

    private:
        ui8 Unused_ = 0; // layout compatibility with cno_header_t
    };

    // XXX linear scan might be even faster since there are generally few (< 64) headers.
    class THeaderVector: public NSorted::TSortedVector<THeader, TStringBuf, THeader::TGetName> {
    public:
        using NSorted::TSortedVector<THeader, TStringBuf, THeader::TGetName>::TSortedVector;

        template <typename F>
        void erase_if(F&& f) noexcept(noexcept(f(*begin()))) {
            erase(std::remove_if(begin(), end(), std::forward<F>(f)), end());
        }

        // Compatibility alias. Can't actually construct in-place due to how the sorted vector
        // operates (need the constructed object to find the place to put it).
        auto emplace(TStringBuf name, TStringBuf value) {
            return insert({name, value});
        }
    };

    using TCookieVector = TVector<std::pair<TStringBuf, TMaybe<TStringBuf>>>;
    TCookieVector ParseCookie(const THeaderVector& v) noexcept;

    inline TStringBuf ParsePath(TStringBuf loc) noexcept {
        return loc.Before('#').Before('?');
    }

    inline TStringBuf ParseQuery(TStringBuf loc) noexcept {
        return loc.Before('#').SplitOff('?');
    }

    inline TStringBuf ParseFragment(TStringBuf loc) noexcept {
        return loc.SplitOff('#');
    }

    inline bool IsInformationalCode(int code) noexcept {
        // 101 is technically an informational code because the server is supposed to
        // actually respond to the request via the new protocol. Obviously, this does
        // not work if we treat the new protocol as a black box.
        return 100 <= code && code < 200 && code != 101;
    }

    class THead: public THeaderVector {
    public:
        THead() = default;
        THead(int code, THeaderVector hs = {}) noexcept // TODO reason string (set as `Method`)?
            : THeaderVector(std::move(hs))
            , Code(code)
        {}

        THead(TStringBuf method, TStringBuf path, THeaderVector hs = {}) noexcept
            : THeaderVector(std::move(hs))
            , Method(method)
            , PathWithQuery(path)
        {
        }

        TStringBuf Path() const noexcept {
            return ParsePath(PathWithQuery);
        }

        TStringBuf Query() const noexcept {
            return ParseQuery(PathWithQuery);
        }

        TStringBuf Fragment() const noexcept {
            return ParseFragment(PathWithQuery);
        }

        bool IsInformational() const noexcept {
            return IsInformationalCode(Code);
        }

    public:
        int Code = 0;
        TStringBuf Method;
        TStringBuf PathWithQuery;
    };

    // This structure holds various information about connection,
    // e.g. backend attempts or connection timeouts.
    // Information may be stored here by any action and is used by `report`
    // struct TStreamStats {
    //     ui64 BackendAttempts = 0;
    //     ui64 BackendError = 0;
    //     ui64 BackendTimeout = 0;
    //     ui64 ConnRefused = 0;
    //     ui64 ConnTimeout = 0;
    //     ui64 ConnReset = 0;
    //     std::array<ui64, 5> BackendStatusCodeCategs;
    //     TDuration ConnectTime = TDuration::Zero();
    //     TDuration FullBackendTime = TDuration::Zero();
    //     TInstant StreamStartTime = TInstant::Zero();

    //     TStreamStats() {
    //         BackendStatusCodeCategs.fill(0);
    //     }
    // };

    // TStreamStats GetStatsDelta(const TStreamStats& lhs, const TStreamStats& rhs);

    class IStream {
    public:
        virtual ~IStream() = default;
        // Return the address of the stream's other side, or AF_UNSPEC if not applicable.
        virtual NSv::IP Peer() const noexcept {
            return {};
        }
        // Return the root log frame for this stream.
        virtual NSv::TLogFrame& Log() noexcept = 0;
        // Wait until receiving the headers, then return them. For server-side streams,
        // always succeeds and returns the same pointer, since a stream is only created once
        // a request (consisting of a single head) is received. For client-side streams,
        //   * if the stream/connection is closed before receiving a response, fails with ECONNRESET;
        //   * if the returned value has code 1xx, the next call will block until the next set
        //     of headers is received. The old pointer remains valid, though.
        virtual THead* Head() noexcept = 0;
        // Return the trailers. This method does not block: if the trailers have not arrived
        // yet, it returns nullptr, in which case you should consume some payload. (It is
        // guaranteed to return non-null if `Read` previously returned an empty chunk, though.)
        // If there are no trailers, a pointer to an empty vector is returned.
        virtual THeaderVector* Tail() noexcept = 0;
        // Wait until at least one byte is available or EOF is reached, then return the contents of
        // the buffer. Fails with ECONNRESET if the connection is closed before the whole payload
        // is received (but continues to return an empty buffer if it is closed after).
        virtual TMaybe<TStringBuf> Read() noexcept = 0;
        // Discard the first n bytes from the buffer, making `Read` stop returning them.
        // This is a write operation -- you're telling the peer that you've accepted
        // the data and are ready to accept more. Can only fail on disconnect (with EPIPE).
        virtual bool Consume(size_t) noexcept = 0;
        // Whether `Read` will only return empty buffers.
        virtual bool AtEnd() const noexcept = 0;
        // Attempt to restart this stream. If unable, return -1; otherwise, undo modifications
        // to `Head()` and `Tail()`, make `Head` start from the first one again, and return
        // the offset of the next chunk that will be returned by `Read`.
        virtual size_t Replay() noexcept {
            return -1;
        }
        // Write a response. Fails with EPIPE if the connection is already closed.
        virtual bool WriteHead(THead&) noexcept = 0;
        // Rvalue reference version of the above.
        bool WriteHead(THead&& h) noexcept {
            return WriteHead(h);
        }
        // Write a chunk of the payload. Fails with EPIPE if the connection is already closed.
        virtual bool Write(TStringBuf) noexcept = 0;
        // Close the payload with trailers. Fails with EPIPE if the connection is already closed.
        // For server-side streams, fails with EREQDONE on success. Trailers can be silently
        // discarded by proxies, and also by this library if the protocol does not support them
        // (e.g. HTTP/1.1 with a `content-length` header), so don't rely on them too much.
        virtual bool Close(THeaderVector&) noexcept = 0;
        // Rvalue reference version of the above.
        bool Close(THeaderVector&& tail = {}) noexcept {
            return Close(tail);
        }

        // Keep a string alive while the request stream is. FIXME at least move this to THead!
        TStringBuf Retain(TString str) {
            Retained.emplace_back(std::move(str));
            return Retained.back();
        }

    private:
        TVector<TString> Retained;

    // public:
    //     TStreamStats Stats;
    };

    using IStreamPtr = std::shared_ptr<IStream>;

    // A wrapper class that forwards all method calls to the underlying IStream.
    // Useful as a base class when installing hooks on some of the methods.
    class TStreamProxy: public IStream {
    public:
        TStreamProxy(IStreamPtr s) noexcept
            : S_(std::move(s))
        {}
        NSv::IP Peer() const noexcept override {
            return S_->Peer();
        }
        NSv::TLogFrame& Log() noexcept override {
            return S_->Log();
        }
        THead* Head() noexcept override {
            return S_->Head();
        }
        THeaderVector* Tail() noexcept override {
            return S_->Tail();
        }
        TMaybe<TStringBuf> Read() noexcept override {
            return S_->Read();
        }
        bool Consume(size_t n) noexcept override {
            return S_->Consume(n);
        }
        bool AtEnd() const noexcept override {
            return S_->AtEnd();
        }
        size_t Replay() noexcept override {
            return S_->Replay();
        }
        bool WriteHead(THead& head) noexcept override {
            return S_->WriteHead(head);
        }
        bool Write(TStringBuf data) noexcept override {
            return S_->Write(data);
        }
        bool Close(THeaderVector& tail) noexcept override {
            return S_->Close(tail);
        }

    protected:
        IStreamPtr S_;
    };

    using THandler = NSv::TFunction<bool(IStreamPtr)>;

    struct IConnection {
        virtual ~IConnection() = default;
        // (Client only.) Send a request. Fails with EPIPE if the connection was closed.
        virtual IStreamPtr Request(const THead& head, bool payload = false, NSv::TLogFrame = {}) noexcept = 0;
        // Check whether there are no active streams on this connection.
        virtual bool IsIdle() const noexcept = 0;
        // (Server only.) Wait until closed.
        virtual bool Wait() noexcept = 0;
        // Signal that no further requests will be made/handled.
        virtual bool Shutdown() noexcept = 0;
        // Move the connection's coroutines to the current thread. The connection must be idle.
        virtual bool Migrate() noexcept = 0;
        // Please don't call this.
        virtual bool IsH2() const noexcept {
            return false;
        }
        // Check whether `Request` has a chance to succeed.
        virtual bool IsOpen() const noexcept {
            return true;
        }
    };

    struct TConnectionOptions {
        // Client-only: use h2 even if the I/O does not indicate it as the selected protocol.
        // Useful for non-TLS connections to enable prior-knowledge h2c establishment.
        bool ForceH2 = false;
        // How many bytes can be written without waiting for a flush to the I/O.
        // If 0, the buffer is not preallocated, but can grow arbitrarily large.
        size_t WriteBuffer = 32768;
        // How many bytes of received payload can be buffered per stream.
        size_t StreamWindow = 65535;
    };

    // HTTP/1.1 and HTTP 2 (over a bytestream connection).
    THolder<IConnection> H2Server(NSv::IO&, THandler, NSv::TNumber<ui64>* emptyWritesSignal, TConnectionOptions = {}, NSv::TLogFrame = {});
    THolder<IConnection> H2Client(NSv::IO&, NSv::TNumber<ui64>* emptyWritesSignal, TConnectionOptions = {});

    // Wrap a single bytestream connection into a fake HTTP stream. Inbound requests appear
    // to be `CONNECT localhost`, outbound headers are discarded.
    THolder<IConnection> RawTCPServer(NSv::IO&, THandler, NSv::TLogFrame = {});
    THolder<IConnection> RawTCPClient(NSv::IO&);

    // Sort-of HTTP messages over a datagram connection. Non-standard, don't use.
    THolder<IConnection> TestUDPServer(NSv::TFile&, THandler, NSv::TLogFrame = {});
    THolder<IConnection> TestUDPClient(const NSv::IP&);

    // Create a HTTP server connection and wait on it. On first cancellation, initiate
    // the graceful shutdown and finish when it's done. On second cancellation,
    // close the connection without waiting for any active streams.
    static inline bool Serve(NSv::IO& io, NSv::THandler h, NSv::TNumber<ui64>* emptyWritesSignal, NSv::TLogFrame log = {}) {
        auto c = H2Server(io, h, emptyWritesSignal, {}, std::move(log));
        // If `Wait` fails, this coroutine has been cancelled.
        return c && (c->Wait() || (c->Shutdown() && c->Wait()));
    }

    // A server-side stream that simply yields a preconstructed request.
    class TConstRequestStream : public IStream {
    public:
        TConstRequestStream(THead head, TStringBuf data, NSv::TLogFrame log)
            : Log_(std::move(log))
            , Head_(head)
            , OriginalHead_(std::move(head))
            , Data_(data)
        {
        }

        void UpdatePayload(TStringBuf data) noexcept {
            Y_ASSERT(data.size() >= Data_.size());
            Data_ = data;
        }

        NSv::TLogFrame& Log() noexcept override {
            return Log_;
        }

        THead* Head() noexcept override {
            return &Head_;
        }

        THeaderVector* Tail() noexcept override {
            return &Tail_;
        }

        TMaybe<TStringBuf> Read() noexcept override {
            return Data_.Tail(Read_);
        }

        bool Consume(size_t n) noexcept override {
            Read_ += n;
            return true;
        }

        bool AtEnd() const noexcept override {
            return Read_ >= Data_.size();
        }

        size_t Replay() noexcept override {
            if (State_ != Headers) {
                return -1;
            }
            Head_ = OriginalHead_;
            Read_ = 0;
            Tail_.clear();
            return 0;
        }

        bool WriteHead(THead& h) noexcept override {
            if (State_ < Body ? !WriteHeadEx(h) : mun_error(EPIPE, "stream closed")) {
                return false;
            }
            State_ = h.IsInformational() ? HeadersCont : Body;
            return true;
        }

        bool Write(TStringBuf data) noexcept override {
            return State_ == Body ? WriteEx(data) : !mun_error(EINVAL, "not writable");
        }

        bool Close(THeaderVector& tail) noexcept override {
            if (State_ == Body ? !CloseEx(tail) : mun_error(EINVAL, "write headers first")) {
                return false;
            }
            State_ = Closed;
            return !mun_error(EREQDONE, "abort request handling");
        }

        virtual bool WriteHeadEx(THead&) noexcept = 0;
        virtual bool WriteEx(TStringBuf) noexcept = 0;
        virtual bool CloseEx(THeaderVector&) noexcept = 0;

    private:
        NSv::TLogFrame Log_;
        THead Head_;
        THead OriginalHead_;
        THeaderVector Tail_;
        TStringBuf Data_;
        size_t Read_ = 0;
        enum { Headers, HeadersCont, Body, Closed } State_ = Headers;
    };

    template <typename TWriteHead, typename TWrite, typename TClose>
    class TLambdaStream : public TConstRequestStream, TWriteHead, TWrite, TClose {
    public:
        TLambdaStream(THead head, TStringBuf data, TWriteHead wh, TWrite w, TClose c, NSv::TLogFrame log, TMaybe<NSv::IP> peer)
            : TConstRequestStream(std::move(head), data, std::move(log))
            , TWriteHead(std::move(wh))
            , TWrite(std::move(w))
            , TClose(std::move(c))
            , Peer_(std::move(peer))
        {
        }

        NSv::IP Peer() const noexcept override {
            return Peer_.GetOrElse(NSv::IP{});
        }

        bool WriteHeadEx(THead& h) noexcept override {
            return TWriteHead::operator()(h);
        }
        bool WriteEx(TStringBuf data) noexcept override {
            return TWrite::operator()(data);
        }
        bool CloseEx(THeaderVector& tail) noexcept override {
            return TClose::operator()(tail);
        }

    private:
        TMaybe<NSv::IP> Peer_;
    };

    template <typename TWriteHead, typename TWrite, typename TClose>
    IStreamPtr ConstRequestStream(THead head, TStringBuf data, TWriteHead wh = {}, TWrite w = {},
                                TClose c = {}, NSv::TLogFrame log = {}, TMaybe<NSv::IP> peer = {}) {
        // TODO make logging non-optional (and move to first argument)?
        // TODO a version with custom trailers? requests don't normally have them...
        return std::make_shared<TLambdaStream<TWriteHead, TWrite, TClose>>(
            std::move(head), data, std::move(wh), std::move(w), std::move(c), std::move(log), std::move(peer)
        );
    }
}
