// Copyright (c) 2023 Cloudflare, Inc.
// Licensed under the Apache 2.0 license found in the LICENSE file or at:
//     https://opensource.org/licenses/Apache-2.0

#pragma once

#include <workerd/api/memory-cache.h>
#include <workerd/io/io-context.h>
#include <workerd/io/worker.h>
#include <workerd/jsg/jsg.h>
#include <workerd/server/workerd.capnp.h>

#include <capnp/message.h>
#include <kj/function.h>
#include <kj/test.h>

namespace workerd {

// TestFixture is responsible for creating workerd environment during tests.
// All the infrastructure is started in the constructor. It is accessed through run() method.
struct TestFixture {
  struct SetupParams {
    // waitScope of outer IO loop. New IO will be set up if missing.
    kj::Maybe<kj::WaitScope&> waitScope;
    kj::Maybe<CompatibilityFlags::Reader> featureFlags;
    kj::Maybe<kj::StringPtr> mainModuleSource;
    // If set, make a stub of an Actor with the given id.
    kj::Maybe<Worker::Actor::Id> actorId;
    // If true, use real timers instead of mock timers that never advance.
    // Requires waitScope to be kj::none (so that the fixture creates its own AsyncIoContext).
    bool useRealTimers;
    // If set, used instead of the default DummyIoChannelFactory when creating incoming requests.
    // The factory receives the TimerChannel reference.
    kj::Maybe<kj::Function<kj::Own<IoChannelFactory>(TimerChannel&)>> ioChannelFactory;
  };

  TestFixture(SetupParams&& params = {.useRealTimers = false});

  struct V8Environment {
    v8::Isolate* isolate;
  };

  struct Environment: public V8Environment {
    IoContext& context;
    Worker::Lock& lock;
    jsg::Lock& js;
    CompatibilityFlags::Reader features;
  };

  template <typename T>
  struct RunReturnType {
    using Type = T;
  };
  template <typename T>
  struct RunReturnType<kj::Promise<T>> {
    using Type = T;
  };

  // Setup the incoming request and run given callback in worker's IO context.
  // callback should accept const Environment& parameter and return Promise<T>|void.
  // For void callbacks run waits for their completion, for promises waits for their resolution
  // and returns the result.
  template <typename CallBack>
  auto runInIoContext(CallBack&& callback)
      -> RunReturnType<decltype(callback(kj::instance<const Environment&>()))>::Type {
    auto request = createIncomingRequest();
    kj::WaitScope* waitScope;
    KJ_IF_SOME(ws, this->waitScope) {
      waitScope = &ws;
    } else {
      waitScope = &KJ_REQUIRE_NONNULL(io).waitScope;
    }

    auto& context = request->getContext();
    return context
        .run([&](Worker::Lock& lock) {
      // auto features = workerBundle.getFeatureFlags();
      auto& js = jsg::Lock::from(lock.getIsolate());
      Environment env = {{.isolate = lock.getIsolate()}, context, lock, js};
      KJ_ASSERT(env.isolate == v8::Isolate::TryGetCurrent());
      return callback(env);
    }).wait(*waitScope);
  }

  // Special void version of runInIoContext that ignores exceptions with given descriptions.
  void runInIoContext(kj::Function<kj::Promise<void>(const Environment&)>&& callback,
      kj::ArrayPtr<const kj::StringPtr> errorsToIgnore);

  struct Response {
    uint statusCode;
    kj::String body;
  };

  // Performs HTTP request on the default module handler, and waits for full response.
  Response runRequest(kj::HttpMethod method, kj::StringPtr url, kj::StringPtr body);

 private:
  kj::Maybe<kj::WaitScope&> waitScope;
  capnp::MallocMessageBuilder configArena;
  workerd::server::config::Worker::Reader config;
  kj::Maybe<kj::AsyncIoContext> io;
  capnp::MallocMessageBuilder workerBundleArena;
  kj::Own<kj::Timer> timer;
  kj::Own<TimerChannel> timerChannel;
  kj::Own<kj::EntropySource> entropySource;
  kj::Maybe<kj::Own<Worker::Actor>> actor;
  capnp::ByteStreamFactory byteStreamFactory;
  kj::HttpHeaderTable::Builder headerTableBuilder;
  ThreadContext::HeaderIdBundle threadContextHeaderBundle;
  capnp::HttpOverCapnpFactory httpOverCapnpFactory;
  ThreadContext threadContext;
  kj::Own<Worker::ValidationErrorReporter> errorReporter;
  kj::Own<api::MemoryCacheProvider> memoryCacheProvider;
  v8::IsolateGroup isolateGroup;
  kj::Own<Worker::Api> api;
  kj::Own<Worker::Isolate> workerIsolate;
  kj::Own<Worker::Script> workerScript;
  kj::Own<Worker> worker;
  kj::Own<kj::TaskSet::ErrorHandler> errorHandler;
  kj::TaskSet waitUntilTasks;
  kj::Own<kj::HttpHeaderTable> headerTable;
  kj::Maybe<kj::Function<kj::Own<IoChannelFactory>(TimerChannel&)>> ioChannelFactory;

  kj::Own<IoContext::IncomingRequest> createIncomingRequest();

 public:
  // Default IoChannelFactory used by tests. Exposed so tests can subclass it
  // and override individual methods (e.g. startSubrequest for socket connect tests).
  struct DummyIoChannelFactory: public IoChannelFactory {
    virtual ~DummyIoChannelFactory() = default;
    DummyIoChannelFactory(TimerChannel& timer): timer(timer) {}

    void abortIsolate(kj::StringPtr reason) override {
      KJ_FAIL_ASSERT("no abortIsolate");
    }

    kj::Own<WorkerInterface> startSubrequest(uint channel, SubrequestMetadata metadata) override {
      KJ_FAIL_ASSERT("no subrequests");
    }
    kj::Own<SubrequestChannel> getSubrequestChannel(uint channel,
        kj::Maybe<Frankenvalue> props,
        kj::Maybe<VersionRequest> versionRequest) override {
      KJ_FAIL_ASSERT("no subrequests");
    }
    capnp::Capability::Client getCapability(uint channel) override {
      KJ_FAIL_ASSERT("no capabilities");
    }
    // Out-of-line because it references file-local MockCacheClient in test-fixture.c++.
    kj::Own<CacheClient> getCache() override;
    TimerChannel& getTimer() override {
      return timer;
    }
    kj::Promise<void> writeLogfwdr(
        uint channel, kj::FunctionParam<void(capnp::AnyPointer::Builder)> buildMessage) override {
      KJ_FAIL_ASSERT("no log channels");
    }
    kj::Own<ActorChannel> getGlobalActor(uint channel,
        const ActorIdFactory::ActorId& id,
        kj::Maybe<kj::String> locationHint,
        ActorGetMode mode,
        bool enableReplicaRouting,
        ActorRoutingMode routingMode,
        SpanParent parentSpan,
        kj::Maybe<ActorVersion> version) override {
      KJ_FAIL_REQUIRE("no actor channels");
    }
    kj::Own<ActorChannel> getColoLocalActor(
        uint channel, kj::StringPtr id, SpanParent parentSpan) override {
      KJ_FAIL_REQUIRE("no actor channels");
    }

    TimerChannel& timer;
  };
};

}  // namespace workerd