// Copyright 2026 Google LLC
//
// 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
//
//    https://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 "async/async_context.h"

#include <memory>
#include <string>
#include <utility>

#include "absl/functional/any_invocable.h"
#include "absl/status/status.h"
#include "absl/status/status_macros.h"  // IWYU pragma: keep
#include "absl/status/status_matchers.h"
#include "absl/status/statusor.h"
#include "async/executor.h"
#include "gmock/gmock.h"
#include "gtest/gtest.h"

namespace credentio {
namespace {

using ::absl_testing::IsOkAndHolds;
using ::absl_testing::StatusIs;

class MockExecutor : public Executor {
 public:
  MOCK_METHOD(void, Schedule, (absl::AnyInvocable<void() &&> task), (override));
};

TEST(AsyncContextTest, ExecutesOnExecutorIfContextProvided) {
  MockExecutor mock_executor;
  AsyncContext context(&mock_executor);

  bool task_scheduled = false;
  EXPECT_CALL(mock_executor, Schedule(testing::_))
      .WillOnce([&task_scheduled](absl::AnyInvocable<void() &&> task) {
        task_scheduled = true;
        std::move(task)();
      });

  bool task_run = false;
  AsyncResolver<int>(context, [&task_run](absl::StatusOr<int> status) {
    task_run = true;
    EXPECT_THAT(status, IsOkAndHolds(42));
  }).Return(42);

  EXPECT_TRUE(task_scheduled);
  EXPECT_TRUE(task_run);
}

TEST(AsyncContextTest, ExecutesInlineIfContextIsNull) {
  bool task_run = false;
  AsyncResolver<int>(AsyncContext(nullptr), [&task_run](
                                                absl::StatusOr<int> status) {
    task_run = true;
    EXPECT_THAT(status, IsOkAndHolds(42));
  }).Return(42);

  EXPECT_TRUE(task_run);
}

TEST(AsyncContextTest, ExecutesInlineIfExecutorIsNull) {
  AsyncContext context(nullptr);

  bool task_run = false;
  AsyncResolver<int>(context, [&task_run](absl::StatusOr<int> status) {
    task_run = true;
    EXPECT_THAT(status, IsOkAndHolds(42));
  }).Return(42);

  EXPECT_TRUE(task_run);
}

TEST(AsyncContextTest, CancellationTokensCheck) {
  AsyncContext context(nullptr);
  EXPECT_FALSE(context.GetCancellationToken().IsCancelled());
  context.Cancel();
  EXPECT_TRUE(context.GetCancellationToken().IsCancelled());
}

TEST(AsyncContextTest, ResolveOnAsyncContextStatusCancelledWhenCancelled) {
  MockExecutor mock_executor;
  AsyncContext context(&mock_executor);
  context.Cancel();

  EXPECT_CALL(mock_executor, Schedule(testing::_)).Times(0);

  bool task_run = false;
  AsyncResolver<void>(context, [&task_run](absl::Status status) {
    task_run = true;
    EXPECT_THAT(status, StatusIs(absl::StatusCode::kCancelled));
  }).Return(absl::OkStatus());

  EXPECT_TRUE(task_run);
}

TEST(AsyncContextTest, ResolveOnAsyncContextStatusOrCancelledWhenCancelled) {
  MockExecutor mock_executor;
  AsyncContext context(&mock_executor);
  context.Cancel();

  EXPECT_CALL(mock_executor, Schedule(testing::_)).Times(0);

  bool task_run = false;
  AsyncResolver<int>(context, [&task_run](absl::StatusOr<int> status) {
    task_run = true;
    EXPECT_THAT(status.status(), StatusIs(absl::StatusCode::kCancelled));
  }).Return(42);

  EXPECT_TRUE(task_run);
}

TEST(AsyncContextTest, ExecutesOnExecutorIfContextProvidedVoid) {
  MockExecutor mock_executor;
  AsyncContext context(&mock_executor);

  bool task_scheduled = false;
  EXPECT_CALL(mock_executor, Schedule(testing::_))
      .WillOnce([&task_scheduled](absl::AnyInvocable<void() &&> task) {
        task_scheduled = true;
        std::move(task)();
      });

  bool task_run = false;
  AsyncResolver<void>(context, [&task_run](absl::Status status) {
    task_run = true;
    ABSL_EXPECT_OK(status);
  }).Return(absl::OkStatus());

  EXPECT_TRUE(task_scheduled);
  EXPECT_TRUE(task_run);
}

TEST(AsyncContextTest, ExecutesInlineIfExecutorIsNullVoid) {
  AsyncContext context(nullptr);

  bool task_run = false;
  AsyncResolver<void>(context, [&task_run](absl::Status status) {
    task_run = true;
    ABSL_EXPECT_OK(status);
  }).Return(absl::OkStatus());

  EXPECT_TRUE(task_run);
}

struct IntExtension {
  int value = 0;
};

struct StringExtension {
  std::string value;
};

class BaseExtension {
 public:
  virtual ~BaseExtension() = default;
  virtual int id() const = 0;
};

class DerivedExtension : public BaseExtension {
 public:
  int id() const override { return 99; }
};

TEST(AsyncContextTest, ExtensionsStoreRetrieveOverwriteAndClear) {
  AsyncContext context;
  EXPECT_EQ(context.GetExtension<IntExtension>(), nullptr);
  EXPECT_EQ(context.GetExtension<StringExtension>(), nullptr);

  auto int_ext = std::make_shared<IntExtension>(IntExtension{42});
  auto str_ext = std::make_shared<StringExtension>(StringExtension{"hello"});
  context.SetExtension(int_ext);
  context.SetExtension(str_ext);

  EXPECT_EQ(context.GetExtension<IntExtension>(), int_ext);
  EXPECT_EQ(context.GetExtension<StringExtension>(), str_ext);

  // Copying the context preserves attached extensions.
  AsyncContext copy = context;
  EXPECT_EQ(copy.GetExtension<IntExtension>(), int_ext);
  EXPECT_EQ(copy.GetExtension<StringExtension>(), str_ext);

  // Overwriting and clearing on the original does not mutate the copy's map.
  auto new_int_ext = std::make_shared<IntExtension>(IntExtension{7});
  context.SetExtension(new_int_ext);
  EXPECT_EQ(context.GetExtension<IntExtension>(), new_int_ext);
  EXPECT_EQ(copy.GetExtension<IntExtension>(), int_ext);

  context.SetExtension<IntExtension>(nullptr);
  EXPECT_EQ(context.GetExtension<IntExtension>(), nullptr);
  EXPECT_EQ(context.GetExtension<StringExtension>(), str_ext);
}

TEST(AsyncContextTest, InterfaceExtensionKeyedByExplicitBaseType) {
  AsyncContext context;
  auto derived = std::make_shared<DerivedExtension>();
  context.SetExtension<BaseExtension>(derived);

  std::shared_ptr<BaseExtension> retrieved =
      context.GetExtension<BaseExtension>();
  ASSERT_NE(retrieved, nullptr);
  EXPECT_EQ(retrieved->id(), 99);
}

}  // namespace
}  // namespace credentio