// Copyright (c) 2017-2022 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 "common.h"
#include "writable.h"

#include <workerd/io/io-context.h>
#include <workerd/io/observer.h>
#include <workerd/util/ring-buffer.h>
#include <workerd/util/state-machine.h>

#include <kj/refcount.h>

namespace workerd::api {

// =======================================================================================
// The ReadableStreamInternalController and WritableStreamInternalController provide the
// internal (original) implementation of the ReadableStream/WritableStream objects and are
// each backed by the ReadableStreamSource and WritableStreamSink respectively. Every stream
// implementation that originates from *within* the Workers runtime will use these.
//
// It is important to understand that the behavior of these are not entirely compliant with
// the streams specification.

// The ReadableStreamInternalController is always in one of three states: Readable, Closed,
// or Errored. When the state is Readable, the controller has an associated ReadableStreamSource.
// When the state is Errored, the ReadableStreamSource has been released and the controller
// stores a jsg::Value with whatever value was used to error. When Closed, the
// ReadableStreamSource has been released.

// Likewise, the WritableStreamInternalController is always either Writable, Closed, or Errored.
// When the state is Writable, the controller has an associated WritableStreamSink. In either of
// the other two states, the sink has been released.

class WritableStreamInternalController;

class ReadableStreamInternalController: public ReadableStreamController {
 public:
  using Readable = IoOwn<ReadableStreamSource>;

  explicit ReadableStreamInternalController(StreamStates::Closed closed)
      : state(State::create<StreamStates::Closed>()) {}
  explicit ReadableStreamInternalController(StreamStates::Errored errored)
      : state(State::create<StreamStates::Errored>(kj::mv(errored))) {}
  explicit ReadableStreamInternalController(Readable readable)
      : state(State::create<Readable>(kj::mv(readable))) {}

  KJ_DISALLOW_COPY_AND_MOVE(ReadableStreamInternalController);

  ~ReadableStreamInternalController() noexcept(false) override;

  void setOwnerRef(ReadableStream& stream) override {
    owner = stream;
  }

  jsg::Ref<ReadableStream> addRef() override;

  bool isByteOriented() const override {
    return true;
  }

  kj::Maybe<jsg::Promise<ReadResult>> read(
      jsg::Lock& js, kj::Maybe<ByobOptions> byobOptions) override;

  kj::Maybe<jsg::Promise<DrainingReadResult>> drainingRead(
      jsg::Lock& js, size_t maxRead = kj::maxValue) override;

  jsg::Promise<void> pipeTo(
      jsg::Lock& js, WritableStreamController& destination, PipeToOptions options) override;

  jsg::Promise<void> cancel(jsg::Lock& js, jsg::Optional<v8::Local<v8::Value>> reason) override;

  Tee tee(jsg::Lock& js) override;

  kj::Maybe<kj::Own<ReadableStreamSource>> removeSource(
      jsg::Lock& js, bool ignoreDisturbed = false);

  bool isClosedOrErrored() const override {
    return state.is<StreamStates::Closed>() || state.is<StreamStates::Errored>();
  }

  bool isClosed() const override {
    return state.is<StreamStates::Closed>();
  }

  bool isDisturbed() override {
    return disturbed;
  }

  bool isLockedToReader() const override {
    return !readState.is<Unlocked>();
  }

  bool lockReader(jsg::Lock& js, Reader& reader) override;

  void releaseReader(Reader& reader, kj::Maybe<jsg::Lock&> maybeJs) override;
  // See the comment for releaseReader in common.h for details on the use of maybeJs

  kj::Maybe<PipeController&> tryPipeLock() override;

  void visitForGc(jsg::GcVisitor& visitor) override;

  jsg::Promise<jsg::BufferSource> readAllBytes(jsg::Lock& js, uint64_t limit) override;
  jsg::Promise<kj::String> readAllText(jsg::Lock& js, uint64_t limit) override;

  kj::Maybe<uint64_t> tryGetLength(StreamEncoding encoding) override;

  kj::Promise<DeferredProxy<void>> pumpTo(
      jsg::Lock& js, kj::Own<WritableStreamSink> sink, bool end) override;

  StreamEncoding getPreferredEncoding() override;

  kj::Own<ReadableStreamController> detach(jsg::Lock& js, bool ignoreDisturbed) override;

  void setPendingClosure() override {
    isPendingClosure = true;
  }

  kj::StringPtr jsgGetMemoryName() const override;
  size_t jsgGetMemorySelfSize() const override;
  void jsgGetMemoryInfo(jsg::MemoryTracker& info) const override;

 private:
  void doCancel(jsg::Lock& js, jsg::Optional<v8::Local<v8::Value>> reason);
  void doClose(jsg::Lock& js);
  void doError(jsg::Lock& js, v8::Local<v8::Value> reason);

  class PipeLocked: public PipeController {
   public:
    static constexpr kj::StringPtr NAME KJ_UNUSED = "pipe-locked"_kj;
    PipeLocked(ReadableStreamInternalController& inner): inner(inner) {}

    bool isClosed() override;

    kj::Maybe<v8::Local<v8::Value>> tryGetErrored(jsg::Lock& js) override;

    void cancel(jsg::Lock& js, v8::Local<v8::Value> reason) override;

    void close(jsg::Lock& js) override;

    void error(jsg::Lock& js, v8::Local<v8::Value> reason) override;

    void release(jsg::Lock& js, kj::Maybe<v8::Local<v8::Value>> maybeError = kj::none) override;

    kj::Maybe<kj::Promise<void>> tryPumpTo(WritableStreamSink& sink, bool end) override;

    jsg::Promise<ReadResult> read(jsg::Lock& js) override;

   private:
    ReadableStreamInternalController& inner;
  };

  kj::Maybe<ReadableStream&> owner;

  // State machine for ReadableStreamInternalController:
  // Closed is terminal, Errored is implicitly terminal via ErrorState.
  // Readable is the active state (stream has data).
  using State = StateMachine<TerminalStates<StreamStates::Closed>,
      ErrorState<StreamStates::Errored>,
      ActiveState<Readable>,
      StreamStates::Closed,
      StreamStates::Errored,
      Readable>;
  State state;

  // Lock state machine for ReadableStreamInternalController:
  // All states can transition to any other state (no terminal states).
  //   Unlocked -> Locked (removeSink() or pumpTo() called)
  //   Unlocked -> ReaderLocked (lockReader() called)
  //   Unlocked -> PipeLocked (tryPipeLock() called)
  //   ReaderLocked -> Unlocked (releaseReader() called)
  //   PipeLocked -> Unlocked (release() or doClose/doError called)
  //   Locked -> (remains until stream is done)
  using ReadLockState = StateMachine<Unlocked, Locked, PipeLocked, ReaderLocked>;
  ReadLockState readState = ReadLockState::create<Unlocked>();
  bool disturbed = false;
  bool readPending = false;

  // Used by Sockets code to signal to the ReadableStream that it should error when read from
  // because the socket is currently being closed.
  bool isPendingClosure = false;

  friend class ReadableStream;
  friend class WritableStreamInternalController;
  friend class PipeLocked;
};

class WritableStreamInternalController: public WritableStreamController {
 public:
  struct Writable {
    kj::Own<WritableStreamSink> sink;
    kj::Canceler canceler;
    Writable(kj::Own<WritableStreamSink> sink): sink(kj::mv(sink)) {}
    void abort(kj::Exception&& ex);
  };

  explicit WritableStreamInternalController(StreamStates::Closed closed)
      : state(State::create<StreamStates::Closed>()) {}
  explicit WritableStreamInternalController(StreamStates::Errored errored)
      : state(State::create<StreamStates::Errored>(kj::mv(errored))) {}
  explicit WritableStreamInternalController(kj::Own<WritableStreamSink> writable,
      kj::Maybe<kj::Own<ByteStreamObserver>> observer,
      kj::Maybe<uint64_t> maybeHighWaterMark = kj::none,
      kj::Maybe<jsg::Promise<void>> maybeClosureWaitable = kj::none)
      : state(State::create<IoOwn<Writable>>(
            IoContext::current().addObject(kj::heap<Writable>(kj::mv(writable))))),
        observer(kj::mv(observer)),
        maybeHighWaterMark(maybeHighWaterMark),
        maybeClosureWaitable(kj::mv(maybeClosureWaitable)) {}

  WritableStreamInternalController(WritableStreamInternalController&& other) = default;
  WritableStreamInternalController& operator=(WritableStreamInternalController&& other) = default;

  ~WritableStreamInternalController() noexcept(false) override;

  void setOwnerRef(WritableStream& stream) override {
    owner = stream;
  }

  jsg::Ref<WritableStream> addRef() override;

  jsg::Promise<void> write(jsg::Lock& js, jsg::Optional<v8::Local<v8::Value>> value) override;

  jsg::Promise<void> close(jsg::Lock& js, bool markAsHandled = false) override;

  jsg::Promise<void> flush(jsg::Lock& js, bool markAsHandled = false) override;

  jsg::Promise<void> abort(jsg::Lock& js, jsg::Optional<v8::Local<v8::Value>> reason) override;

  kj::Maybe<jsg::Promise<void>> tryPipeFrom(
      jsg::Lock& js, jsg::Ref<ReadableStream> source, PipeToOptions options) override;

  kj::Maybe<kj::Own<WritableStreamSink>> removeSink(jsg::Lock& js) override;
  void detach(jsg::Lock& js) override;

  kj::Maybe<int> getDesiredSize() override;

  bool isLockedToWriter() const override {
    return !writeState.is<Unlocked>();
  }

  bool lockWriter(jsg::Lock& js, Writer& writer) override;

  void releaseWriter(Writer& writer, kj::Maybe<jsg::Lock&> maybeJs) override;
  // See the comment for releaseWriter in common.h for details on the use of maybeJs

  kj::Maybe<v8::Local<v8::Value>> isErroring(jsg::Lock& js) override {
    // TODO(later): The internal controller has no concept of an "erroring"
    // state, so for now we just return kj::none here.
    return kj::none;
  }

  void visitForGc(jsg::GcVisitor& visitor) override;

  void setHighWaterMark(uint64_t highWaterMark);

  bool isClosedOrClosing() override;
  bool isPiping();
  bool isErrored() override;

  inline bool isByteOriented() const override {
    return true;
  }

  void setPendingClosure() override {
    isPendingClosure = true;
  }

  kj::StringPtr jsgGetMemoryName() const override;
  size_t jsgGetMemorySelfSize() const override;
  void jsgGetMemoryInfo(jsg::MemoryTracker& info) const override;

 private:
  struct AbortOptions {
    bool reject = false;
    bool handled = false;
  };

  jsg::Promise<void> doAbort(jsg::Lock& js,
      v8::Local<v8::Value> reason,
      AbortOptions options = {.reject = false, .handled = false});
  void doClose(jsg::Lock& js);
  void doError(jsg::Lock& js, v8::Local<v8::Value> reason);
  void ensureWriting(jsg::Lock& js);
  jsg::Promise<void> writeLoop(jsg::Lock& js, IoContext& ioContext);
  jsg::Promise<void> writeLoopAfterFrontOutputLock(jsg::Lock& js);

  void drain(jsg::Lock& js, v8::Local<v8::Value> reason);
  void finishClose(jsg::Lock& js);
  void finishError(jsg::Lock& js, v8::Local<v8::Value> reason);
  jsg::Promise<void> closeImpl(jsg::Lock& js, bool markAsHandled);

  struct PipeLocked {
    static constexpr kj::StringPtr NAME KJ_UNUSED = "pipe-locked"_kj;
    ReadableStream& ref;
  };

  kj::Maybe<WritableStream&> owner;

  // State machine for WritableStreamInternalController:
  // Closed is terminal, Errored is implicitly terminal via ErrorState.
  // IoOwn<Writable> is the active state (stream is writable).
  using State = StateMachine<TerminalStates<StreamStates::Closed>,
      ErrorState<StreamStates::Errored>,
      ActiveState<IoOwn<Writable>>,
      StreamStates::Closed,
      StreamStates::Errored,
      IoOwn<Writable>>;
  State state;

  // Lock state machine for WritableStreamInternalController:
  // All states can transition to any other state (no terminal states).
  //   Unlocked -> Locked (removeSink() or detach() called)
  //   Unlocked -> WriterLocked (lockWriter() called)
  //   Unlocked -> PipeLocked (tryPipeFrom() called)
  //   WriterLocked -> Unlocked (releaseWriter() called)
  //   WriterLocked -> Locked (doClose/doError called - stream closed but writer still attached)
  //   PipeLocked -> Unlocked (pipe completes)
  using WriteLockState = StateMachine<Unlocked, Locked, PipeLocked, WriterLocked>;
  WriteLockState writeState = WriteLockState::create<Unlocked>();

  kj::Maybe<kj::Own<ByteStreamObserver>> observer;

  kj::Maybe<kj::Own<PendingAbort>> maybePendingAbort;

  uint64_t currentWriteBufferSize = 0;

  // The highWaterMark is the total amount of data currently buffered in
  // the controller waiting to be flushed out to the underlying WritableStreamSink.
  // It is used to implement backpressure signaling using desiredSize and the ready
  // promise on the writer.
  kj::Maybe<uint64_t> maybeHighWaterMark;

  // Used by Sockets code to ensure the connection is established before the associated
  // WritableStream is closed.
  kj::Maybe<jsg::Promise<void>> maybeClosureWaitable;
  bool waitingOnClosureWritableAlready = false;

  // Used by Sockets code to signal to the WritableStream that it should error when written to
  // because the socket is currently being closed.
  bool isPendingClosure = false;

  void adjustWriteBufferSize(jsg::Lock& js, int64_t amount);
  void updateBackpressure(jsg::Lock& js, bool backpressure);

  struct Write {
    kj::Maybe<jsg::Promise<void>::Resolver> promise;
    size_t totalBytes;
    kj::Array<kj::byte> ownBytes;
    kj::ArrayPtr<const kj::byte> bytes;

    JSG_MEMORY_INFO(Write) {
      tracker.trackField("resolver", promise);
      if (ownBytes != nullptr) {
        tracker.trackFieldWithSize("backing", totalBytes);
      }
    }
  };
  struct Close {
    kj::Maybe<jsg::Promise<void>::Resolver> promise;
    JSG_MEMORY_INFO(Close) {
      tracker.trackField("promise", promise);
    }
  };
  struct Flush {
    kj::Maybe<jsg::Promise<void>::Resolver> promise;
    JSG_MEMORY_INFO(Flush) {
      tracker.trackField("promise", promise);
    }
  };
  struct Pipe {
    // PipeState is ref-counted so that it can be safely captured by lambdas in pipeLoop().
    // When drain() destroys the Pipe, the state survives as long as pending callbacks need it.
    // The `aborted` flag is set when the Pipe is destroyed.
    struct State: public kj::Refcounted {
      WritableStreamInternalController& parent;
      ReadableStreamController::PipeController& source;
      kj::Maybe<jsg::Promise<void>::Resolver> promise;
      kj::Maybe<jsg::Ref<AbortSignal>> maybeSignal;

      bool preventAbort;
      bool preventClose;
      bool preventCancel;

      // True when the Pipe is being destroyed
      bool aborted = false;

      State(WritableStreamInternalController& parent,
          ReadableStreamController::PipeController& source,
          kj::Maybe<jsg::Promise<void>::Resolver> promise,
          bool preventAbort,
          bool preventClose,
          bool preventCancel,
          kj::Maybe<jsg::Ref<AbortSignal>> maybeSignal)
          : parent(parent),
            source(source),
            promise(kj::mv(promise)),
            maybeSignal(kj::mv(maybeSignal)),
            preventAbort(preventAbort),
            preventClose(preventClose),
            preventCancel(preventCancel) {}

      bool checkSignal(jsg::Lock& js);
      jsg::Promise<void> pipeLoop(jsg::Lock& js);
      jsg::Promise<void> write(v8::Local<v8::Value> value);

      JSG_MEMORY_INFO(State) {
        tracker.trackField("resolver", promise);
        tracker.trackField("signal", maybeSignal);
      }
    };

    kj::Own<State> state;

    Pipe(WritableStreamInternalController& parent,
        ReadableStreamController::PipeController& source,
        kj::Maybe<jsg::Promise<void>::Resolver> promise,
        bool preventAbort,
        bool preventClose,
        bool preventCancel,
        kj::Maybe<jsg::Ref<AbortSignal>> maybeSignal)
        : state(kj::refcounted<State>(parent,
              source,
              kj::mv(promise),
              preventAbort,
              preventClose,
              preventCancel,
              kj::mv(maybeSignal))) {}

    ~Pipe() noexcept(false) {
      state->aborted = true;
    }

    WritableStreamInternalController& parent() {
      return state->parent;
    }
    ReadableStreamController::PipeController& source() {
      return state->source;
    }
    kj::Maybe<jsg::Promise<void>::Resolver>& promise() {
      return state->promise;
    }
    bool preventAbort() const {
      return state->preventAbort;
    }
    bool preventClose() const {
      return state->preventClose;
    }
    bool preventCancel() const {
      return state->preventCancel;
    }
    kj::Maybe<jsg::Ref<AbortSignal>>& maybeSignal() {
      return state->maybeSignal;
    }

    bool checkSignal(jsg::Lock& js) {
      return state->checkSignal(js);
    }
    jsg::Promise<void> pipeLoop(jsg::Lock& js) {
      return state->pipeLoop(js);
    }
    jsg::Promise<void> write(v8::Local<v8::Value> value) {
      return state->write(value);
    }

    JSG_MEMORY_INFO(Pipe) {
      tracker.trackField("state", state);
    }
  };
  struct WriteEvent {
    kj::Maybe<IoOwn<kj::Promise<void>>> outputLock;  // must wait for this before actually writing
    kj::OneOf<kj::Own<Write>, kj::Own<Pipe>, kj::Own<Close>, kj::Own<Flush>> event;

    JSG_MEMORY_INFO(WriteEvent) {
      if (outputLock != kj::none) {
        tracker.trackFieldWithSize("outputLock", sizeof(IoOwn<kj::Promise<void>>));
      }
      KJ_SWITCH_ONEOF(event) {
        KJ_CASE_ONEOF(w, kj::Own<Write>) {
          tracker.trackField("inner", w);
        }
        KJ_CASE_ONEOF(p, kj::Own<Pipe>) {
          tracker.trackField("inner", p);
        }
        KJ_CASE_ONEOF(c, kj::Own<Close>) {
          tracker.trackField("inner", c);
        }
        KJ_CASE_ONEOF(f, kj::Own<Flush>) {
          tracker.trackField("inner", f);
        }
      }
    }
  };

  RingBuffer<WriteEvent, 8> queue;
};
}  // namespace workerd::api