//
//
// Copyright 2025 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.
//
//

#include <grpc/event_engine/event_engine.h>

#include <cstdint>
#include <limits>
#include <utility>
#include <variant>
#include <vector>

#include "src/core/call/message.h"
#include "src/core/call/metadata.h"
#include "src/core/call/metadata_batch.h"
#include "src/core/ext/transport/chttp2/transport/frame.h"
#include "src/core/ext/transport/chttp2/transport/header_assembler.h"
#include "src/core/ext/transport/chttp2/transport/hpack_encoder.h"
#include "src/core/ext/transport/chttp2/transport/hpack_parser.h"
#include "src/core/ext/transport/chttp2/transport/http2_settings.h"
#include "src/core/ext/transport/chttp2/transport/http2_status.h"
#include "src/core/ext/transport/chttp2/transport/message_assembler.h"
#include "src/core/ext/transport/chttp2/transport/stream_data_queue.h"
#include "src/core/ext/transport/chttp2/transport/write_cycle.h"
#include "src/core/lib/promise/if.h"
#include "src/core/lib/promise/loop.h"
#include "src/core/lib/promise/map.h"
#include "src/core/lib/promise/party.h"
#include "src/core/lib/promise/sleep.h"
#include "src/core/lib/promise/status_flag.h"
#include "src/core/lib/promise/try_seq.h"
#include "src/core/lib/resource_quota/arena.h"
#include "src/core/lib/slice/slice.h"
#include "src/core/lib/slice/slice_buffer.h"
#include "src/core/util/crash.h"
#include "src/core/util/grpc_check.h"
#include "src/core/util/ref_counted_ptr.h"
#include "src/core/util/time.h"
#include "test/core/call/yodel/yodel_test.h"
#include "gmock/gmock.h"
#include "gtest/gtest.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/strings/string_view.h"

namespace grpc_core {

using http2::SimpleQueue;
using http2::StreamDataQueue;
using http2::ValueOrHttp2Status;
using ::testing::MockFunction;
using ::testing::StrictMock;

// Fuzzer tests
class SimpleQueueFuzzTest : public YodelTest {
 protected:
  using YodelTest::YodelTest;

  Party* GetParty() { return party_.get(); }
  Arena* GetArena() { return party_->arena(); }

  void InitParty() {
    auto party_arena = SimpleArenaAllocator(0)->MakeArena();
    party_arena->SetContext<grpc_event_engine::experimental::EventEngine>(
        event_engine().get());
    party_ = Party::Make(std::move(party_arena));
  }

  auto EnqueueAndCheckSuccess(SimpleQueue<int>& queue, int data, int tokens) {
    return Map([&queue, data,
                tokens]() mutable { return queue.Enqueue(data, tokens); },
               [data, tokens](bool result) {
                 LOG(INFO) << "Enqueue done for data with tokens: " << data
                           << " tokens: " << tokens << " result: " << result;
               });
  }

  bool DequeueAndCheck(SimpleQueue<int>& queue, int data,
                       bool allow_oversized_dequeue,
                       int allowed_dequeue_tokens) {
    auto result =
        queue.Dequeue(allowed_dequeue_tokens, allow_oversized_dequeue);
    if (!result.has_value()) {
      return false;
    }

    EXPECT_EQ(result.value(), data);
    return true;
  }

 private:
  void InitCoreConfiguration() override {}
  void InitTest() override { InitParty(); }
  void Shutdown() override { party_.reset(); }

  RefCountedPtr<Party> party_;
};

YODEL_TEST(SimpleQueueFuzzTest, NoOp) {}

YODEL_TEST(SimpleQueueFuzzTest, EnqueueAndDequeueMultiPartyTest) {
  // Test to enqueue and dequeue in a loop. This test enqueues 100 entries and
  // dequeues 100 entries. This test asserts the following:
  // 1. All enqueues and dequeues are successful.
  // 2. The dequeue data is the same as the enqueue data.
  SimpleQueue<int> queue(GetArena(), /*max_tokens=*/100);
  StrictMock<MockFunction<void(absl::Status)>> on_done;
  StrictMock<MockFunction<void(absl::Status)>> on_dequeue_done;
  EXPECT_CALL(on_done, Call(absl::OkStatus()));
  EXPECT_CALL(on_dequeue_done, Call(absl::OkStatus()));
  constexpr int count = 100;
  constexpr int dequeue_count = 100;
  int current_enqueue_count = 0;
  int current_dequeue_count = 0;

  GetParty()->Spawn(
      "EnqueueTest",
      Loop([this, &queue, &on_done, &current_enqueue_count]() mutable {
        return If(
            current_enqueue_count < count,
            [this, &queue, &current_enqueue_count] {
              return Map(
                  EnqueueAndCheckSuccess(queue,
                                         /*data=*/current_enqueue_count++,
                                         /*tokens=*/10),
                  [](auto) -> LoopCtl<StatusFlag> { return Continue(); });
            },
            [&on_done]() -> LoopCtl<StatusFlag> {
              on_done.Call(absl::OkStatus());
              return Success{};
            });
      }),
      [](auto) { LOG(INFO) << "Reached end of EnqueueTest"; });

  GetParty()->Spawn(
      "DequeueTest",
      Loop([&on_dequeue_done, this, &queue, &current_dequeue_count] {
        return If(
            DequeueAndCheck(queue, /*data=*/current_dequeue_count,
                            /*allow_oversized_dequeue=*/false,
                            /*allowed_dequeue_tokens=*/10),
            [&current_dequeue_count, &on_dequeue_done,
             &queue]() -> LoopCtl<absl::Status> {
              if (++current_dequeue_count == dequeue_count) {
                on_dequeue_done.Call(absl::OkStatus());
                EXPECT_TRUE(queue.IsEmpty());
                return absl::OkStatus();
              } else {
                return Continue();
              }
            },
            [] {
              return Map(
                  Sleep(Duration::Seconds(1)),
                  [](auto) -> LoopCtl<absl::Status> { return Continue(); });
            });
      }),
      [](auto) { LOG(INFO) << "Reached end of DequeueTest"; });

  WaitForAllPendingWork();
  event_engine()->TickUntilIdle();
  event_engine()->UnsetGlobalHooks();
}

class StreamDataQueueFuzzTest : public YodelTest {
 protected:
  using YodelTest::YodelTest;

  Party* GetParty() { return party_.get(); }
  Party* GetParty2() { return party2_.get(); }
  Arena* GetArena() { return party_->arena(); }

  void InitParty() {
    auto party_arena = SimpleArenaAllocator(0)->MakeArena();
    party_arena->SetContext<grpc_event_engine::experimental::EventEngine>(
        event_engine().get());
    party_ = Party::Make(std::move(party_arena));
  }

  void InitParty2() {
    auto party_arena = SimpleArenaAllocator(0)->MakeArena();
    party_arena->SetContext<grpc_event_engine::experimental::EventEngine>(
        event_engine().get());
    party2_ = Party::Make(std::move(party_arena));
  }

  HPackCompressor& GetEncoder() { return encoder_; }

  // Helper functions to create test data.
  ClientMetadataHandle TestClientInitialMetadata() {
    auto md = Arena::MakePooledForOverwrite<ClientMetadata>();
    md->Set(HttpPathMetadata(), Slice::FromStaticString("/demo.Service/Step"));
    return md;
  }

  MessageHandle TestMessage(SliceBuffer payload, const uint32_t flags) {
    return Arena::MakePooled<Message>(std::move(payload), flags);
  }
  std::vector<MessageHandle> TestMessages(int num_messages) {
    std::vector<MessageHandle> messages;
    for (int i = 0; i < num_messages; ++i) {
      messages.push_back(
          TestMessage(SliceBuffer(Slice::ZeroContentsWithLength(i * 10)), 0));
    }
    return messages;
  }

  http2::WriteCycle& GetWriteCycle() {
    return transport_write_context_.GetWriteCycle();
  }

  class AssembleFrames {
   public:
    explicit AssembleFrames(const uint32_t stream_id,
                            const bool allow_true_binary_metadata)
        : header_assembler_(/*is_client=*/true) {
      header_assembler_.SetStreamId(stream_id);
      header_assembler_.MaybeSetAllowTrueBinaryMetadataAcked(
          allow_true_binary_metadata);
    }
    void operator()(Http2HeaderFrame frame) {
      auto status = header_assembler_.AppendFrame(frame);
      EXPECT_TRUE(status.IsOk());
    }
    void operator()(Http2ContinuationFrame frame) {
      auto status = header_assembler_.AppendFrame(frame);
      EXPECT_TRUE(status.IsOk());
    }
    void operator()(Http2DataFrame frame) {
      auto status = message_assembler_.AppendNewDataFrame(frame.payload,
                                                          frame.end_stream);
      EXPECT_TRUE(status.IsOk());
    }
    void operator()(Http2RstStreamFrame frame) {
      Crash("RstStreamFrame not expected");
    }
    void operator()(Http2GoawayFrame frame) {
      Crash("GoAwayFrame not expected");
    }
    void operator()(Http2WindowUpdateFrame frame) {
      Crash("WindowUpdateFrame not expected");
    }
    void operator()(Http2SettingsFrame frame) {
      Crash("SettingsFrame not expected");
    }
    void operator()(Http2SecurityFrame frame) {
      Crash("SecurityFrame not expected");
    }
    void operator()(Http2EmptyFrame frame) { Crash("EmptyFrame not expected"); }
    void operator()(Http2PingFrame frame) { Crash("PingFrame not expected"); }
    void operator()(Http2UnknownFrame frame) {
      Crash("UnknownFrame not expected");
    }
    ClientMetadataHandle GetMetadata() {
      GRPC_DCHECK(header_assembler_.IsReady());
      ValueOrHttp2Status<ClientMetadataHandle> status_or_metadata =
          header_assembler_.ReadMetadata(
              parser_, /*is_initial_metadata=*/true,
              /*max_header_list_size_soft_limit=*/
              default_settings_.max_header_list_size(),
              /*max_header_list_size_hard_limit=*/
              default_settings_.max_header_list_size());
      EXPECT_TRUE(status_or_metadata.IsOk());
      return TakeValue(std::move(status_or_metadata));
    }
    void GetMessages(std::vector<MessageHandle>& messages) {
      while (true) {
        ValueOrHttp2Status<MessageHandle> message =
            message_assembler_.ExtractMessage();
        if (message.IsOk() && message.value() != nullptr) {
          messages.push_back(TakeValue(std::move(message)));
        } else {
          break;
        }
      }
    }

   private:
    http2::HeaderAssembler header_assembler_;
    http2::GrpcMessageAssembler message_assembler_;
    HPackParser parser_;
    Http2Settings default_settings_;
  };

  void MaybeFlushWriteBuffer() {
    if (GetWriteCycle().CanSerializeRegularFrames()) {
      bool unused;
      SliceBuffer discard = GetWriteCycle().SerializeRegularFrames({unused});
    }
  }

 private:
  void InitCoreConfiguration() override {}
  void InitTest() override {
    InitParty();
    transport_write_context_.StartWriteCycle();
    bool unused;
    // Discard the connection preface
    SliceBuffer discard =
        transport_write_context_.GetWriteCycle().SerializeRegularFrames(
            {unused});
  }

  void Shutdown() override {
    party_.reset();
    party2_.reset();
    transport_write_context_.EndWriteCycle();
  }

  RefCountedPtr<Party> party_;
  RefCountedPtr<Party> party2_;
  HPackCompressor encoder_;
  http2::TransportWriteContext transport_write_context_{/*is_client=*/true};
};

// TODO(akshitpatel) : [PH2][P3] : Add a test for server side.
YODEL_TEST(StreamDataQueueFuzzTest, EnqueueDequeueMultiParty) {
  // Test to enqueue and dequeue in a loop. This test enqueues 100 messages and
  // dequeues all frames and assembles them back into 100 messages. This test
  // asserts the following:
  // 1. All enqueues and dequeues are successful.
  // 2. The dequeue messages are the same as the enqueue messages.

  InitParty();
  InitParty2();
  constexpr uint32_t max_frame_length = 44;
  constexpr uint32_t max_tokens = 75;
  constexpr uint32_t queue_size = 100;
  constexpr uint32_t stream_id = 1;
  constexpr uint32_t num_messages = 100;
  StrictMock<MockFunction<void()>> on_enqueue_done;
  StrictMock<MockFunction<void()>> on_dequeue_done;
  EXPECT_CALL(on_enqueue_done, Call());
  EXPECT_CALL(on_dequeue_done, Call());
  HPackCompressor encoder;
  StreamDataQueue<ClientMetadataHandle> stream_data_queue(
      GetArena(),
      /*is_client=*/true,
      /*queue_size=*/queue_size);
  stream_data_queue.SetStreamId(stream_id,
                                /*allow_true_binary_metadata_peer=*/true);
  std::vector<MessageHandle> messages_to_be_sent = TestMessages(num_messages);
  std::vector<MessageHandle> messages_copy = TestMessages(num_messages);
  std::vector<MessageHandle> dequeued_messages;
  uint message_index = 0;
  AssembleFrames assembler(stream_id, /*allow_true_binary_metadata=*/true);

  auto validate = [this, &messages_copy, &dequeued_messages, &assembler]() {
    auto metadata = assembler.GetMetadata();
    auto expected_metadata = TestClientInitialMetadata();
    absl::string_view metadata_value =
        metadata->get_pointer(HttpPathMetadata())->as_string_view();
    absl::string_view expected_metadata_value =
        expected_metadata->get_pointer(HttpPathMetadata())->as_string_view();
    EXPECT_EQ(metadata_value, expected_metadata_value);
    auto& messages = dequeued_messages;
    EXPECT_EQ(messages.size(), messages_copy.size());
    for (int i = 0; i < messages.size(); ++i) {
      EXPECT_EQ(messages[i]->flags(), messages_copy[i]->flags());
      EXPECT_EQ(messages[i]->payload()->JoinIntoString(),
                messages_copy[i]->payload()->JoinIntoString());
    }
  };

  GetParty()->Spawn(
      "EnqueuePromise",
      TrySeq(
          [&stream_data_queue, this]() {
            return stream_data_queue.EnqueueInitialMetadata(
                TestClientInitialMetadata());
          },
          [&stream_data_queue, &messages_to_be_sent, &message_index] {
            return Loop([&stream_data_queue, &messages_to_be_sent,
                         &message_index] {
              return If(
                  message_index < messages_to_be_sent.size(),
                  [&stream_data_queue, &messages_to_be_sent, &message_index] {
                    return Map(stream_data_queue.EnqueueMessage(std::move(
                                   messages_to_be_sent[message_index++])),
                               [](auto) -> LoopCtl<absl::Status> {
                                 return Continue{};
                               });
                  },
                  []() -> LoopCtl<absl::Status> { return absl::OkStatus(); });
            });
          },
          [&stream_data_queue] {
            return stream_data_queue.EnqueueHalfClosed();
          }),
      [&on_enqueue_done](auto) { on_enqueue_done.Call(); });

  GetParty2()->Spawn(
      "DequeuePromise",
      Loop([this, &stream_data_queue, &validate, &dequeued_messages,
            &assembler]() {
        return If(
            dequeued_messages.size() >= num_messages,
            [&validate]() -> LoopCtl<absl::Status> {
              validate();
              return absl::OkStatus();
            },
            [this, &stream_data_queue, &assembler, &dequeued_messages] {
              http2::FrameSender frame_sender =
                  GetWriteCycle().GetFrameSender();
              GRPC_UNUSED
              typename StreamDataQueue<ClientMetadataHandle>::DequeueResult
                  result = stream_data_queue.DequeueFrames(
                      max_tokens, max_frame_length, /*stream_fc_tokens=*/
                      std::numeric_limits<uint32_t>::max(), GetEncoder(),
                      frame_sender,
                      /*can_send_reset_stream=*/true);
              for (auto& frame : GetWriteCycle().TestOnlyUrgentFrames()) {
                std::visit(assembler, std::move(frame));
              }
              for (auto& frame : GetWriteCycle().TestOnlyRegularFrames()) {
                std::visit(assembler, std::move(frame));
              }
              assembler.GetMessages(dequeued_messages);
              MaybeFlushWriteBuffer();
              return Map(
                  Sleep(Duration::Seconds(1)),
                  [](auto) -> LoopCtl<absl::Status> { return Continue{}; });
            });
      }),
      [&on_dequeue_done](auto) { on_dequeue_done.Call(); });

  WaitForAllPendingWork();
  event_engine()->TickUntilIdle();
  event_engine()->UnsetGlobalHooks();
}

}  // namespace grpc_core
