#pragma once

#include "function.h"

#include <balancer/serval/contrib/cone/cone.hh>
#include <balancer/serval/core/proto/log.pb.h>
#include <balancer/serval/core/proto/log.ev.pb.h>

#include <util/datetime/base.h>
#include <util/generic/ptr.h>
#include <util/stream/output.h>

#include <atomic>

namespace NSv {
    class TLogFrame;

    // Like eventlog, but tree-structured and streamed.
    //
    // The normal eventlog messages (.ev files, supported by `ya make`) are used. Those
    // are just protobuf messages with unique values of the `message_id` option. Writing
    // is done by a separate thread, with a bounded queue in between. Despite there being
    // a class named `TLogFrame`, events aren't actually grouped into frames; they are
    // written in the order defined by that queue. The delay between pushing and writing
    // is ~5 seconds, see `TLog::Reopen`. For context, instead of a frame id, an event can
    // be assigned a parent event, e.g. "request received" could be a child of "connection
    // accepted", itself under "listening". The file format is explained in `proto/log.proto`.
    class TLog : TNonCopyable {
    public:
        // Just 60 random bytes. See `NSv::TLogEvent::Sync` in `proto/log.proto`.
        static constexpr TStringBuf SyncField =
            "\x2c\x3e\xb8\x48\xa3\xb5\xbd\x72\xd9\xf2\x10\x5e\xf4\x70\x8f\xce\x2e\xf2\x88\x6b"
            "\xc6\xc4\xe3\x70\xae\xe8\xbc\x46\xf5\x84\x06\xd9\xe7\x71\x4a\x15\x46\x32\xbf\x07"
            "\xeb\x15\x22\x59\x7f\x43\xbb\x75\x11\xc9\x65\x52\x67\x72\x26\x27\xa3\xea\xa0\x7f"sv;

        TLog(size_t queueLimit)
            : Remaining_(queueLimit)
        {
        }

        ~TLog() {
            Reopen(nullptr);
        }

        // Flush all pending events to the current output stream, then begin writing
        // to a new one. While swapping the stream, events are buffered, which may
        // cause a bit of overflow if the queue is not large enough. This function is
        // not safe to call from multiple coroutines (might result in no logging at all).
        void Reopen(THolder<IOutputStream> out);

        // Total number of events dropped due to queue overflow.
        ui32 Overflow() const { return Overflow_.load(std::memory_order_relaxed); }

        // Current space remaining in the queue.
        ui32 Capacity() const { return static_cast<ui32>(Max(0, Remaining_.load(std::memory_order_relaxed))); }

    private:
        struct IItem {
            std::atomic<IItem*> Next_{nullptr};
        };

        struct TItem : IItem {
            virtual ~TItem() = default;
            ui64 Time;
            ui32 Context;
            ui32 Current;
            NProtoBuf::Message* Data;
        };

        template <typename T>
        struct TConcreteItem : TItem, T {
            template <typename... Args>
            TConcreteItem(TInstant t, ui32 context, ui32 id, Args&&... args)
                : T(std::forward<Args>(args)...)
            {
                Time = t.MicroSeconds();
                Context = context;
                Current = id;
                Data = this;
            }
        };

        void PushOwned(IItem* item) noexcept {
            item->Next_.store(nullptr, std::memory_order_relaxed);
            // Blocking between xchg and store blocks the consumer, too.
            Tail_.exchange(item)->Next_.store(item, std::memory_order_release);
            More_.wake();
        }

    private:
        friend class TLogFrame;
        std::atomic<ui32> Overflow_{0};
        std::atomic<ui32> LastId_{0};
        std::atomic<int> Remaining_;
        std::atomic<bool> Enabled_{false};
        std::atomic<IItem*> Tail_{&Stub_};
        cone::event More_;
        cone::guard Flush_;
        IItem* Head_{&Stub_};
        IItem Stub_;
        IItem Stop_;
        THolder<IOutputStream> Out_;
    };

    class TLogFrame {
    public:
        using TTracer = NSv::TFunction<void(TInstant, ui32 parent, ui32 id, const NProtoBuf::Message&)>;

        TLogFrame() = default;
        TLogFrame(TLogFrame&&) = default;
        TLogFrame& operator=(TLogFrame&&) = default;

        TLogFrame(TLog& log, ui32 parent = 0, TTracer trace = {})
            : Log_(&log)
            , Trace_(std::move(trace))
            , Context_(parent)
        {
        }

        ~TLogFrame() {
            if (Context_) {
                Push<NSv::NEv::TFrameEnd>();
            }
        }

        // Add a single event as a child of the one that began this subtree. Be careful
        // when using timestamps far in the past/future, as this behaves poorly with
        // `tools/logdump`'s binary search.
        template <typename T, typename... Args>
        void PushAt(TInstant t, Args&&... args) {
            if (Log_ && Log_->Enabled_) {
                PushImpl<T, Args...>(0, t, std::forward<Args>(args)...);
            }
        }

        template <typename T, typename... Args>
        void Push(Args&&... args) {
            return PushAt<T, Args...>(TInstant::Now(), std::forward<Args>(args)...);
        }

        // Add a child event that itself can have children.
        template <typename T, typename... Args>
        TLogFrame ForkAt(TInstant t, Args&&... args) {
            return Log_ && Log_->Enabled_
                 ? TLogFrame(*Log_, PushImpl<T, Args...>(++Log_->LastId_, t, std::forward<Args>(args)...), Trace_)
                 : TLogFrame();
        }

        template <typename T, typename... Args>
        TLogFrame Fork(Args&&... args) {
            return ForkAt<T, Args...>(TInstant::Now(), std::forward<Args>(args)...);
        }

        // Add a subtree, and also make it so that events within that subtree are passed
        // to a callback. Useful for adding debug traces to responses, etc.
        template <typename T, typename... Args>
        TLogFrame ForkTracingAt(TTracer f, TInstant t, Args&&... args) {
            return Log_ && Log_->Enabled_
                 ? TLogFrame(*Log_, PushImpl<T, Args...>(++Log_->LastId_, t, std::forward<Args>(args)...), std::move(f))
                 : TLogFrame();
        }

        template <typename T, typename... Args>
        TLogFrame ForkTracing(TTracer f, Args&&... args) {
            return ForkTracingAt<T, Args...>(std::move(f), TInstant::Now(), std::forward<Args>(args)...);
        }

    private:
        template <typename T, typename... Args>
        ui32 PushImpl(ui32 id, TInstant t, Args&&... args) {
            if (Trace_) {
                Trace_(t, Context_, id, T(args...));
            }
            if (--Log_->Remaining_ < 0) {
                // Here we have a small window where `Remaining_` slightly overestimates
                // the actual number of items (at most by the number of producer threads).
                ++Log_->Remaining_, ++Log_->Overflow_;
            } else {
                // TODO small thread-local buffer to avoid contention on a single atomic?
                Log_->PushOwned(new TLog::TConcreteItem<T>(t, Context_, id, std::forward<Args>(args)...));
            }
            return id;
        }

    private:
        THolder<TLog, TNoAction> Log_;
        TTracer Trace_;
        ui32 Context_ = 0;
    };
}
