/*

 * Copyright 2022, gRPC Authors All rights reserved.

 *

 * 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.

 */

import DequeModule

import NIOCore

internal struct CoalescingLengthPrefixedMessageWriter {

  /// Length of the gRPC message header (1 compression byte, 4 bytes for the length).

  static let metadataLength = 5

  /// Message size above which we emit two buffers: one containing the header and one with the

  /// actual message bytes. At or below the limit we copy the message into a new buffer containing

  /// both the header and the message.

  ///

  /// Using two buffers avoids expensive copies of large messages. For smaller messages the copy

  /// is cheaper than the additional allocations and overhead required to send an extra HTTP/2 DATA

  /// frame.

  ///

  /// The value of 16k was chosen empirically. We subtract the length of the message header

  /// as `ByteBuffer` reserve capacity in powers of two and want to avoid overallocating.

  static let singleBufferSizeLimit = 16384 - Self.metadataLength

  /// The compression algorithm to use, if one should be used.

  private let compression: CompressionAlgorithm?

  /// Any compressor associated with the compression algorithm.

  private let compressor: Zlib.Deflate?

  /// Whether the compression message flag should be set.

  private var supportsCompression: Bool {

    return self.compression != nil

  }

  /// A scratch buffer that we encode messages into: if the buffer isn't held elsewhere then we

  /// can avoid having to allocate a new one.

  private var scratch: ByteBuffer

  /// Outbound buffers waiting to be written.

  private var pending: OneOrManyQueue<Pending>

  private struct Pending {

    var buffer: ByteBuffer

    var promise: EventLoopPromise<Void>?

    var compress: Bool

    init(buffer: ByteBuffer, compress: Bool, promise: EventLoopPromise<Void>?) {

      self.buffer = buffer

      self.promise = promise

      self.compress = compress

    }

    var isSmallEnoughToCoalesce: Bool {

      let limit = CoalescingLengthPrefixedMessageWriter.singleBufferSizeLimit

      return self.buffer.readableBytes <= limit

    }

    var shouldCoalesce: Bool {

      return self.isSmallEnoughToCoalesce || self.compress

    }

  }

  private enum State {

    // Coalescing pending messages.

    case coalescing

    // Emitting a non-coalesced message; the header has been written, the body

    // needs to be written next.

    case emittingLargeFrame(ByteBuffer, EventLoopPromise<Void>?)

  }

  private var state: State

  init(compression: CompressionAlgorithm? = nil, allocator: ByteBufferAllocator) {

    self.compression = compression

    self.scratch = allocator.buffer(capacity: 0)

    self.state = .coalescing

    self.pending = .init()

    switch self.compression?.algorithm {

    case .none, .some(.identity):

      self.compressor = nil

    case .some(.deflate):

      self.compressor = Zlib.Deflate(format: .deflate)

    case .some(.gzip):

      self.compressor = Zlib.Deflate(format: .gzip)

    }

  }

  /// Append a serialized message buffer to the writer.

  mutating func append(buffer: ByteBuffer, compress: Bool, promise: EventLoopPromise<Void>?) {

    let pending = Pending(

      buffer: buffer,

      compress: compress && self.supportsCompression,

      promise: promise

    )

    self.pending.append(pending)

  }

  /// Return a tuple of the next buffer to write and its associated write promise.

  mutating func next() -> (Result<ByteBuffer, Error>, EventLoopPromise<Void>?)? {

    switch self.state {

    case .coalescing:

      // Nothing pending: exit early.

      if self.pending.isEmpty {

        return nil

      }

      // First up we need to work out how many elements we're going to pop off the front

      // and coalesce.

      //

      // At the same time we'll compute how much capacity we'll need in the buffer and cascade

      // their promises.

      var messagesToCoalesce = 0

      var requiredCapacity = 0

      var promise: EventLoopPromise<Void>?

      for element in self.pending {

        if !element.shouldCoalesce {

          break

        }

        messagesToCoalesce &+= 1

        requiredCapacity += element.buffer.readableBytes + Self.metadataLength

        if let existing = promise {

          existing.futureResult.cascade(to: element.promise)

        } else {

          promise = element.promise

        }

      }

      if messagesToCoalesce == 0 {

        // Nothing to coalesce; this means the first element should be emitted with its header in

        // a separate buffer. Note: the force unwrap is okay here: we early exit if `self.pending`

        // is empty.

        let pending = self.pending.pop()!

        // Set the scratch buffer to just be a message header then store the message bytes.

        self.scratch.clear(minimumCapacity: Self.metadataLength)

        self.scratch.writeMultipleIntegers(UInt8(0), UInt32(pending.buffer.readableBytes))

        self.state = .emittingLargeFrame(pending.buffer, pending.promise)

        return (.success(self.scratch), nil)

      } else {

        self.scratch.clear(minimumCapacity: requiredCapacity)

        // Drop and encode the messages.

        while messagesToCoalesce > 0, let next = self.pending.pop() {

          messagesToCoalesce &-= 1

          do {

            try self.encode(next.buffer, compress: next.compress)

          } catch {

            return (.failure(error), promise)

          }

        }

        return (.success(self.scratch), promise)

      }

    case let .emittingLargeFrame(buffer, promise):

      // We just emitted the header, now emit the body.

      self.state = .coalescing

      return (.success(buffer), promise)

    }

  }

  private mutating func encode(_ buffer: ByteBuffer, compress: Bool) throws {

    if let compressor = self.compressor, compress {

      try self.encode(buffer, compressor: compressor)

    } else {

      try self.encode(buffer)

    }

  }

  private mutating func encode(_ buffer: ByteBuffer, compressor: Zlib.Deflate) throws {

    // Set the compression byte.

    self.scratch.writeInteger(UInt8(1))

    // Set the length to zero; we'll write the actual value in a moment.

    let payloadSizeIndex = self.scratch.writerIndex

    self.scratch.writeInteger(UInt32(0))

    let bytesWritten: Int

    do {

      var buffer = buffer

      bytesWritten = try compressor.deflate(&buffer, into: &self.scratch)

    } catch {

      throw error

    }

    self.scratch.setInteger(UInt32(bytesWritten), at: payloadSizeIndex)

    // Finally, the compression context should be reset between messages.

    compressor.reset()

  }

  private mutating func encode(_ buffer: ByteBuffer) throws {

    self.scratch.writeMultipleIntegers(UInt8(0), UInt32(buffer.readableBytes))

    self.scratch.writeImmutableBuffer(buffer)

  }

}

/// A FIFO-queue which allows for a single to be stored on the stack and defers to a

/// heap-implementation if further elements are added.

///

/// This is useful when optimising for unary streams where avoiding the cost of a heap

/// allocation is desirable.

internal struct OneOrManyQueue<Element>: Collection {

  private var backing: Backing

  private enum Backing: Collection {

    case none

    case one(Element)

    case many(Deque<Element>)

    var startIndex: Int {

      switch self {

      case .none, .one:

        return 0

      case let .many(elements):

        return elements.startIndex

      }

    }

    var endIndex: Int {

      switch self {

      case .none:

        return 0

      case .one:

        return 1

      case let .many(elements):

        return elements.endIndex

      }

    }

    subscript(index: Int) -> Element {

      switch self {

      case .none:

        fatalError("Invalid index")

      case let .one(element):

        assert(index == 0)

        return element

      case let .many(elements):

        return elements[index]

      }

    }

    func index(after index: Int) -> Int {

      switch self {

      case .none:

        return 0

      case .one:

        return 1

      case let .many(elements):

        return elements.index(after: index)

      }

    }

    var count: Int {

      switch self {

      case .none:

        return 0

      case .one:

        return 1

      case let .many(elements):

        return elements.count

      }

    }

    var isEmpty: Bool {

      switch self {

      case .none:

        return true

      case .one:

        return false

      case let .many(elements):

        return elements.isEmpty

      }

    }

    mutating func append(_ element: Element) {

      switch self {

      case .none:

        self = .one(element)

      case let .one(one):

        var elements = Deque<Element>()

        elements.reserveCapacity(16)

        elements.append(one)

        elements.append(element)

        self = .many(elements)

      case var .many(elements):

        self = .none

        elements.append(element)

        self = .many(elements)

      }

    }

    mutating func pop() -> Element? {

      switch self {

      case .none:

        return nil

      case let .one(element):

        self = .none

        return element

      case var .many(many):

        self = .none

        let element = many.popFirst()

        self = .many(many)

        return element

      }

    }

  }

  init() {

    self.backing = .none

  }

  var isEmpty: Bool {

    return self.backing.isEmpty

  }

  var count: Int {

    return self.backing.count

  }

  var startIndex: Int {

    return self.backing.startIndex

  }

  var endIndex: Int {

    return self.backing.endIndex

  }

  subscript(index: Int) -> Element {

    return self.backing[index]

  }

  func index(after index: Int) -> Int {

    return self.backing.index(after: index)

  }

  mutating func append(_ element: Element) {

    self.backing.append(element)

  }

  mutating func pop() -> Element? {

    return self.backing.pop()

  }

}