#include "identity-transform-stream.h"

#include "common.h"

#include <workerd/util/autogate.h>
#include <workerd/util/state-machine.h>

namespace workerd::api {

namespace {
// An implementation of ReadableStreamSource and WritableStreamSink which communicates read and
// write requests via a StateMachine.
//
// This class is also used as the implementation of FixedLengthStream, in which case `limit` is
// non-nullptr.

struct Idle {
  static constexpr kj::StringPtr NAME KJ_UNUSED = "idle"_kj;
};

struct ReadRequest {
  static constexpr kj::StringPtr NAME KJ_UNUSED = "read-request"_kj;
  kj::ArrayPtr<kj::byte> bytes;
  // WARNING: `bytes` may be invalid if fulfiller->isWaiting() returns false! (This indicates the
  //   read was canceled.)
  kj::Own<kj::PromiseFulfiller<size_t>> fulfiller;
};

struct WriteRequest {
  static constexpr kj::StringPtr NAME KJ_UNUSED = "write-request"_kj;
  kj::ArrayPtr<const kj::byte> bytes;
  kj::Own<kj::PromiseFulfiller<void>> fulfiller;
};

struct Closed {
  static constexpr kj::StringPtr NAME KJ_UNUSED = "closed"_kj;
};

// State machine for IdentityTransformStream:
//   Idle -> ReadRequest (read arrives when no write pending)
//   Idle -> WriteRequest (write arrives when no read pending)
//   Idle -> Closed (empty write = close)
//   ReadRequest -> Idle (write fulfills read completely)
//   WriteRequest -> Idle (read fulfills write completely)
//   ReadRequest -> Closed (empty write closes while read pending)
//   Any -> kj::Exception (cancel/abort)
//   Closed -> kj::Exception (abort can force-transition a closed stream to error)
// Closed is terminal, kj::Exception is implicitly terminal via ErrorState.
// abort() uses forceTransitionTo to allow the exceptional Closed -> Exception transition.
using IdentityTransformState = StateMachine<TerminalStates<Closed>,
    ErrorState<kj::Exception>,
    Idle,
    ReadRequest,
    WriteRequest,
    Closed,
    kj::Exception>;

class IdentityTransformStreamImpl final: public kj::Refcounted,
                                         public ReadableStreamSource,
                                         public WritableStreamSink {
 public:
  // The limit is the maximum number of bytes that can be fed through the stream.
  // If kj::none, there is no limit.
  explicit IdentityTransformStreamImpl(kj::Maybe<uint64_t> limit = kj::none)
      : limit(limit),
        state(IdentityTransformState::create<Idle>()) {}

  ~IdentityTransformStreamImpl() noexcept(false) {
    // Due to the different natures of JS and C++ disposal, there is no point in enforcing the limit
    // for a FixedLengthStream here.
    //
    // 1. Creating but not using a `new FixedLengthStream(n)` should not be an error, and ought not
    //    to logspam us.
    // 2. Chances are high that by the time this object gets destroyed, it's too late to tell the
    //    user about the failure.
  }

  // ReadableStreamSource implementation -------------------------------------------------

  kj::Promise<size_t> tryRead(void* buffer, size_t minBytes, size_t maxBytes) override {
    size_t total = 0;
    while (total < minBytes) {
      // TODO(perf): tryReadInternal was written assuming minBytes would always be 1 but we've now
      // introduced an API for user to specify a larger minBytes. For now, this is implemented as a
      // naive loop dispatching to the 1 byte version but would be better to bake it deeper into
      // the implementation where it can be more efficient.
      auto amount = co_await tryReadInternal(buffer, maxBytes);
      KJ_ASSERT(amount <= maxBytes);
      if (amount == 0) {
        // EOF.
        break;
      }

      total += amount;
      buffer = reinterpret_cast<char*>(buffer) + amount;
      maxBytes -= amount;
    }

    co_return total;
  }

  kj::Promise<size_t> tryReadInternal(void* buffer, size_t maxBytes) {
    auto promise = readHelper(kj::arrayPtr(static_cast<kj::byte*>(buffer), maxBytes));

    KJ_IF_SOME(l, limit) {
      promise = promise.then([this, &l = l](size_t amount) -> kj::Promise<size_t> {
        if (amount > l) {
          auto exception = JSG_KJ_EXCEPTION(
              FAILED, TypeError, "Attempt to write too many bytes through a FixedLengthStream.");
          cancel(exception.clone());
          return kj::mv(exception);
        } else if (amount == 0 && l != 0) {
          auto exception = JSG_KJ_EXCEPTION(FAILED, TypeError,
              "FixedLengthStream did not see all expected bytes before close().");
          cancel(exception.clone());
          return kj::mv(exception);
        }
        l -= amount;
        return amount;
      });
    }

    return promise;
  }

  kj::Promise<DeferredProxy<void>> pumpTo(WritableStreamSink& output, bool end) override {
#ifdef KJ_NO_RTTI
    // Yes, I'm paranoid.
    static_assert(!KJ_NO_RTTI, "Need RTTI for correctness");
#endif

    // HACK: If `output` is another TransformStream, we don't allow pumping to it, in order to
    //   guarantee that we can't create cycles.
    JSG_REQUIRE(!isIdentityTransformStream(output), TypeError,
        "Inter-TransformStream ReadableStream.pipeTo() is not implemented.");

    return ReadableStreamSource::pumpTo(output, end);
  }

  kj::Maybe<uint64_t> tryGetLength(StreamEncoding encoding) override {
    if (encoding == StreamEncoding::IDENTITY) {
      return limit;
    } else {
      return kj::none;
    }
  }

  void cancel(kj::Exception reason) override {
    // Already errored - nothing to do.
    if (state.isErrored()) return;

    // Already closed by writable side - nothing to do.
    if (state.is<Closed>()) return;

    KJ_IF_SOME(request, state.tryGetUnsafe<ReadRequest>()) {
      request.fulfiller->fulfill(static_cast<size_t>(0));
    } else KJ_IF_SOME(request, state.tryGetUnsafe<WriteRequest>()) {
      request.fulfiller->reject(reason.clone());
    }
    // Idle state is fine, just transition to error.

    state.forceTransitionTo<kj::Exception>(kj::mv(reason));

    // TODO(conform): Proactively put WritableStream into Errored state.
  }

  // WritableStreamSink implementation ---------------------------------------------------

  kj::Promise<void> write(kj::ArrayPtr<const byte> buffer) override {
    if (buffer == nullptr) {
      return kj::READY_NOW;
    }
    return writeHelper(buffer);
  }

  kj::Promise<void> write(kj::ArrayPtr<const kj::ArrayPtr<const kj::byte>> pieces) override {
    KJ_UNIMPLEMENTED("IdentityTransformStreamImpl piecewise write() not currently supported");
    // TODO(soon): This will be called by TeeBranch::pumpTo(). We disallow that anyway, since we
    //   disallow inter-TransformStream pumping.
  }

  kj::Promise<void> end() override {
    // If we're already closed, there's nothing else we need to do here.
    if (state.is<Closed>()) return kj::READY_NOW;

    return writeHelper(kj::ArrayPtr<const kj::byte>());
  }

  void abort(kj::Exception reason) override {
    // Already errored - nothing to do.
    if (state.isErrored()) return;

    KJ_IF_SOME(request, state.tryGetUnsafe<ReadRequest>()) {
      request.fulfiller->reject(reason.clone());
    } else KJ_IF_SOME(request, state.tryGetUnsafe<WriteRequest>()) {
      // If the fulfiller is not waiting, the write promise was already
      // canceled and no one is waiting on it.
      KJ_ASSERT(!request.fulfiller->isWaiting(),
          "abort() is supposed to wait for any pending write() to finish");
    }
    // Idle and Closed states are fine, just transition to error.
    // (Closed can transition to error via abort)

    state.forceTransitionTo<kj::Exception>(kj::mv(reason));

    // TODO(conform): Proactively put ReadableStream into Errored state.
  }

 private:
  kj::Promise<size_t> readHelper(kj::ArrayPtr<kj::byte> bytes) {
    // Handle error state first.
    KJ_IF_SOME(exception, state.tryGetErrorUnsafe()) {
      return exception.clone();
    }

    // Handle closed state.
    if (state.is<Closed>()) {
      return static_cast<size_t>(0);
    }

    // Check for already in-flight read.
    if (state.is<ReadRequest>()) {
      KJ_FAIL_ASSERT("read operation already in flight");
    }

    // Check for pending write request.
    KJ_IF_SOME(request, state.tryGetUnsafe<WriteRequest>()) {
      if (bytes.size() >= request.bytes.size()) {
        // The write buffer will entirely fit into our read buffer; fulfill both requests.
        memmove(bytes.begin(), request.bytes.begin(), request.bytes.size());
        auto result = request.bytes.size();
        request.fulfiller->fulfill();

        // Switch to idle state.
        state.transitionTo<Idle>();

        return result;
      }

      // The write buffer won't quite fit into our read buffer; fulfill only the read request.
      memmove(bytes.begin(), request.bytes.begin(), bytes.size());
      request.bytes = request.bytes.slice(bytes.size(), request.bytes.size());
      return bytes.size();
    }

    // Must be idle - no outstanding write request, switch to ReadRequest state.
    KJ_ASSERT(state.is<Idle>());
    auto paf = kj::newPromiseAndFulfiller<size_t>();
    state.transitionTo<ReadRequest>(bytes, kj::mv(paf.fulfiller));
    return kj::mv(paf.promise);
  }

  kj::Promise<void> writeHelper(kj::ArrayPtr<const kj::byte> bytes) {
    // Handle error state first.
    KJ_IF_SOME(exception, state.tryGetErrorUnsafe()) {
      return exception.clone();
    }

    // Handle closed state.
    if (state.is<Closed>()) {
      KJ_FAIL_ASSERT("close operation already in flight");
    }

    // Check for already in-flight write.
    if (state.is<WriteRequest>()) {
      KJ_FAIL_ASSERT("write operation already in flight");
    }

    // Check for pending read request.
    KJ_IF_SOME(request, state.tryGetUnsafe<ReadRequest>()) {
      if (!request.fulfiller->isWaiting()) {
        // Oops, the request was canceled. Currently, this happen in particular when pumping a
        // response body to the client, and the client disconnects, cancelling the pump. In this
        // specific case, we want to propagate the error back to the write end of the transform
        // stream. In theory, though, there could be other cases where propagation is incorrect.
        //
        // TODO(cleanup): This cancellation should probably be handled at a higher level, e.g.
        //   in pumpTo(), but I need a quick fix.
        state.forceTransitionTo<kj::Exception>(KJ_EXCEPTION(DISCONNECTED, "reader canceled"));

        // I was going to use a `goto` but Harris choked on his bagel. Recursion it is.
        return writeHelper(bytes);
      }

      if (bytes.size() == 0) {
        // This is a close operation.
        request.fulfiller->fulfill(static_cast<size_t>(0));
        state.transitionTo<Closed>();
        return kj::READY_NOW;
      }

      KJ_ASSERT(request.bytes.size() > 0);

      if (request.bytes.size() >= bytes.size()) {
        // Our write buffer will entirely fit into the read buffer; fulfill both requests.
        memmove(request.bytes.begin(), bytes.begin(), bytes.size());
        request.fulfiller->fulfill(bytes.size());
        state.transitionTo<Idle>();
        return kj::READY_NOW;
      }

      // Our write buffer won't quite fit into the read buffer; fulfill only the read request.
      memmove(request.bytes.begin(), bytes.begin(), request.bytes.size());
      bytes = bytes.slice(request.bytes.size(), bytes.size());
      request.fulfiller->fulfill(request.bytes.size());

      auto paf = kj::newPromiseAndFulfiller<void>();
      state.transitionTo<WriteRequest>(bytes, kj::mv(paf.fulfiller));
      return kj::mv(paf.promise);
    }

    // Must be idle.
    KJ_ASSERT(state.is<Idle>());
    if (bytes.size() == 0) {
      // This is a close operation.
      state.transitionTo<Closed>();
      return kj::READY_NOW;
    }

    auto paf = kj::newPromiseAndFulfiller<void>();
    state.transitionTo<WriteRequest>(bytes, kj::mv(paf.fulfiller));
    return kj::mv(paf.promise);
  }

  kj::Maybe<uint64_t> limit;
  IdentityTransformState state;
};

struct Pair {
  kj::Own<ReadableStreamSource> readable;
  kj::Own<WritableStreamSink> writable;
};
Pair newIdentityPair(kj::Maybe<uint64_t> expectedLength = kj::none) {
  auto readableSide = kj::refcounted<IdentityTransformStreamImpl>(kj::mv(expectedLength));
  auto writableSide = kj::addRef(*readableSide);
  return Pair{.readable = kj::mv(readableSide), .writable = kj::mv(writableSide)};
}
}  // namespace

jsg::Ref<IdentityTransformStream> IdentityTransformStream::constructor(
    jsg::Lock& js, jsg::Optional<IdentityTransformStream::QueuingStrategy> maybeQueuingStrategy) {

  auto& ioContext = IoContext::current();
  auto pipe = newIdentityPipe();

  kj::Maybe<uint64_t> maybeHighWaterMark = kj::none;
  KJ_IF_SOME(queuingStrategy, maybeQueuingStrategy) {
    maybeHighWaterMark = queuingStrategy.highWaterMark;
  }
  return js.alloc<IdentityTransformStream>(js.alloc<ReadableStream>(ioContext, kj::mv(pipe.in)),
      js.alloc<WritableStream>(ioContext, kj::mv(pipe.out),
          ioContext.getMetrics().tryCreateWritableByteStreamObserver(), maybeHighWaterMark));
}

jsg::Ref<FixedLengthStream> FixedLengthStream::constructor(jsg::Lock& js,
    uint64_t expectedLength,
    jsg::Optional<IdentityTransformStream::QueuingStrategy> maybeQueuingStrategy) {
  constexpr uint64_t MAX_SAFE_INTEGER = (1ull << 53) - 1;

  JSG_REQUIRE(expectedLength <= MAX_SAFE_INTEGER, TypeError,
      "FixedLengthStream requires an integer expected length less than 2^53.");

  auto& ioContext = IoContext::current();
  auto pipe = newIdentityPipe(expectedLength);

  kj::Maybe<uint64_t> maybeHighWaterMark = kj::none;
  // For a FixedLengthStream we do not want a highWaterMark higher than the expectedLength.
  KJ_IF_SOME(queuingStrategy, maybeQueuingStrategy) {
    maybeHighWaterMark = queuingStrategy.highWaterMark.map(
        [&](uint64_t highWaterMark) { return kj::min(expectedLength, highWaterMark); });
  }

  return js.alloc<FixedLengthStream>(js.alloc<ReadableStream>(ioContext, kj::mv(pipe.in)),
      js.alloc<WritableStream>(ioContext, kj::mv(pipe.out),
          ioContext.getMetrics().tryCreateWritableByteStreamObserver(), maybeHighWaterMark));
}

OneWayPipe newIdentityPipe(kj::Maybe<uint64_t> expectedLength) {
  auto pair = newIdentityPair(kj::mv(expectedLength));
  return OneWayPipe{.in = kj::mv(pair.readable), .out = kj::mv(pair.writable)};
}

bool isIdentityTransformStream(WritableStreamSink& sink) {
  return kj::dynamicDowncastIfAvailable<IdentityTransformStreamImpl>(sink) != kj::none;
}

}  // namespace workerd::api