// Copyright 2024 gRPC authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef GRPC_SRC_CORE_EXT_TRANSPORT_CHAOTIC_GOOD_LEGACY_CONTROL_ENDPOINT_H
#define GRPC_SRC_CORE_EXT_TRANSPORT_CHAOTIC_GOOD_LEGACY_CONTROL_ENDPOINT_H

#include "src/core/ext/transport/chaotic_good_legacy/legacy_ztrace_collector.h"
#include "src/core/lib/promise/map.h"
#include "src/core/lib/promise/party.h"
#include "src/core/lib/transport/promise_endpoint.h"
#include "src/core/util/sync.h"
#include "absl/cleanup/cleanup.h"

namespace grpc_core {
namespace chaotic_good_legacy {

// Wrapper around PromiseEndpoint.
// Buffers all of the small writes that get enqueued to this endpoint, and then
// uses a separate party to flush them to the wire.
// In doing so we get to batch up effectively all the writes from the transport
// (since party wakeups are sticky), and then flush all the writes in one go.
class ControlEndpoint {
 private:
  class Buffer : public RefCounted<Buffer> {
   public:
    explicit Buffer(std::shared_ptr<LegacyZTraceCollector> ztrace_collector)
        : ztrace_collector_(std::move(ztrace_collector)) {}

    // Queue some buffer to be written.
    // We cap the queue size so that we don't infinitely buffer on one
    // connection - if the cap is hit, this queue operation will not resolve
    // until it empties.
    // Returns a promise that resolves to Empty{} when the data has been queued.
    auto Queue(SliceBuffer&& buffer) {
      return [buffer = std::move(buffer), this]() mutable -> Poll<Empty> {
        Waker waker;
        auto cleanup = absl::MakeCleanup([&waker]() { waker.Wakeup(); });
        MutexLock lock(&mu_);
        if (queued_output_.Length() != 0 &&
            queued_output_.Length() + buffer.Length() > MaxQueued()) {
          write_waker_ = GetContext<Activity>()->MakeNonOwningWaker();
          ztrace_collector_->Append(ControlEndpointQueueWriteTrace{});
          return Pending{};
        }
        ztrace_collector_->Append(ControlEndpointWriteTrace{buffer.Length()});
        waker = std::move(flush_waker_);
        queued_output_.Append(buffer);
        return Empty{};
      };
    }

    auto Pull();

   private:
    size_t MaxQueued() const { return 1024 * 1024; }

    Mutex mu_;
    Waker write_waker_ ABSL_GUARDED_BY(mu_);
    Waker flush_waker_ ABSL_GUARDED_BY(mu_);
    SliceBuffer queued_output_ ABSL_GUARDED_BY(mu_);
    std::shared_ptr<LegacyZTraceCollector> ztrace_collector_;
  };

 public:
  ControlEndpoint(PromiseEndpoint endpoint,
                  grpc_event_engine::experimental::EventEngine* event_engine,
                  std::shared_ptr<LegacyZTraceCollector> ztrace_collector,
                  RefCountedPtr<channelz::SocketNode> socket_node);

  // Write some data to the control endpoint; returns a promise that resolves
  // to Empty{} -- it's not possible to see errors from this api.
  auto Write(SliceBuffer&& bytes) { return buffer_->Queue(std::move(bytes)); }

  // Read operations are simply passthroughs to the underlying promise endpoint.
  auto ReadSlice(size_t length) {
    return Map(
        AddErrorPrefix("CONTROL_CHANNEL: ", endpoint_->ReadSlice(length)),
        [this](absl::StatusOr<Slice> buffer) -> absl::StatusOr<Slice> {
          if (!buffer.ok()) return buffer.status();
          ztrace_collector_->Append(ControlEndpointReadTrace{buffer->size()});
          return buffer;
        });
  }
  auto Read(size_t length) {
    return Map(
        AddErrorPrefix("CONTROL_CHANNEL: ", endpoint_->Read(length)),
        [this](
            absl::StatusOr<SliceBuffer> buffer) -> absl::StatusOr<SliceBuffer> {
          if (!buffer.ok()) return buffer.status();
          ztrace_collector_->Append(ControlEndpointReadTrace{buffer->Length()});
          return buffer;
        });
  }
  auto GetPeerAddress() const { return endpoint_->GetPeerAddress(); }
  auto GetLocalAddress() const { return endpoint_->GetLocalAddress(); }

 private:
  std::shared_ptr<PromiseEndpoint> endpoint_;
  RefCountedPtr<Party> write_party_;
  std::shared_ptr<LegacyZTraceCollector> ztrace_collector_;
  RefCountedPtr<Buffer> buffer_ = MakeRefCounted<Buffer>(ztrace_collector_);
};

}  // namespace chaotic_good_legacy
}  // namespace grpc_core

#endif  // GRPC_SRC_CORE_EXT_TRANSPORT_CHAOTIC_GOOD_LEGACY_CONTROL_ENDPOINT_H
