/*
 * Copyright 2024 Google Inc. 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 Foundation

#if canImport(Common)
import Common
#endif

/// `ByteBuffer` is the interface that stores the data for a `Flatbuffers` object
/// it allows users to Construct their buffers internally without much cost to performance
@usableFromInline
struct _InternalByteBuffer {

  /// Storage is a container that would hold the memory pointer to solve the issue of
  /// deallocating the memory that was held by (memory: UnsafeMutableRawPointer)
  @usableFromInline
  final class Storage {
    /// pointer to the start of the buffer object in memory
    private(set) var memory: UnsafeMutableRawPointer

    @usableFromInline
    init(count: Int, alignment: Int) {
      memory = UnsafeMutableRawPointer.allocate(
        byteCount: count,
        alignment: alignment)
    }

    deinit {
      memory.deallocate()
    }

    @usableFromInline
    func initialize(for size: Int) {
      memset(memory, 0, size)
    }

    /// Reallocates the buffer incase the object to be written doesnt fit in the current buffer
    /// - Parameter size: Size of the current object
    @usableFromInline
    func reallocate(
      capacity: Int,
      writerSize: Int,
      currentWritingIndex: Int,
      alignment: Int)
    {
      let newData = UnsafeMutableRawPointer.allocate(
        byteCount: capacity,
        alignment: alignment)
      memset(newData, 0, capacity &- writerSize)
      memcpy(
        newData.advanced(by: capacity &- writerSize),
        memory.advanced(by: currentWritingIndex),
        writerSize)
      memory.deallocate()
      memory = newData
    }
  }

  @usableFromInline var _storage: Storage

  private let initialSize: Int
  /// The size of the elements written to the buffer + their paddings
  private var _writerSize: Int = 0
  /// Alignment of the current  memory being written to the buffer
  var alignment = 1
  /// Current Index which is being used to write to the buffer, it is written from the end to the start of the buffer
  var writerIndex: Int { capacity &- _writerSize }

  /// Reader is the position of the current Writer Index (capacity - size)
  public var reader: Int { writerIndex }
  /// Current size of the buffer
  public var size: UOffset { UOffset(_writerSize) }
  /// Public Pointer to the buffer object in memory. This should NOT be modified for any reason
  @usableFromInline
  var memory: UnsafeMutableRawPointer { _storage.memory }
  /// Current capacity for the buffer
  public private(set) var capacity: Int

  /// Constructor that creates a Flatbuffer instance with a size
  /// - Parameter:
  ///   - size: Length of the buffer
  ///   - allowReadingUnalignedBuffers: allow reading from unaligned buffer
  init(initialSize size: Int) {
    initialSize = size.convertToPowerofTwo
    capacity = initialSize
    _storage = Storage(count: initialSize, alignment: alignment)
    _storage.initialize(for: initialSize)
  }

  /// Fills the buffer with padding by adding to the writersize
  /// - Parameter padding: Amount of padding between two to be serialized objects
  @inline(__always)
  @usableFromInline
  mutating func fill(padding: Int) {
    assert(padding >= 0, "Fill should be larger than or equal to zero")
    ensureSpace(size: padding)
    _writerSize = _writerSize &+ (MemoryLayout<UInt8>.size &* padding)
  }

  /// Adds an array of type Scalar to the buffer memory
  /// - Parameter elements: An array of Scalars
  @inline(__always)
  @usableFromInline
  mutating func push<T: Scalar>(elements: [T]) {
    elements.withUnsafeBytes { ptr in
      ensureSpace(size: ptr.count)
      memcpy(
        _storage.memory.advanced(by: writerIndex &- ptr.count),
        ptr.baseAddress!,
        ptr.count)
      _writerSize = _writerSize &+ ptr.count
    }
  }

  /// Adds an array of type Scalar to the buffer memory
  /// - Parameter elements: An array of Scalars
  @inline(__always)
  @usableFromInline
  mutating func push<T: NativeStruct>(elements: [T]) {
    elements.withUnsafeBytes { ptr in
      ensureSpace(size: ptr.count)
      memcpy(
        _storage.memory.advanced(by: writerIndex &- ptr.count),
        ptr.baseAddress!,
        ptr.count)
      _writerSize = _writerSize &+ ptr.count
    }
  }

  /// Adds a `ContiguousBytes` to buffer memory
  /// - Parameter value: bytes to copy
  #if !os(WASI)
  @inline(__always)
  @usableFromInline
  mutating func push(bytes: ContiguousBytes) {
    bytes.withUnsafeBytes { ptr in
      ensureSpace(size: ptr.count)
      memcpy(
        _storage.memory.advanced(by: writerIndex &- ptr.count),
        ptr.baseAddress!,
        ptr.count)
      _writerSize = _writerSize &+ ptr.count
    }
  }
  #endif

  /// Adds an object of type NativeStruct into the buffer
  /// - Parameters:
  ///   - value: Object  that will be written to the buffer
  ///   - size: size to subtract from the WriterIndex
  @usableFromInline
  @inline(__always)
  mutating func push<T: NativeStruct>(struct value: T, size: Int) {
    ensureSpace(size: size)
    withUnsafePointer(to: value) {
      memcpy(
        _storage.memory.advanced(by: writerIndex &- size),
        $0,
        size)
      _writerSize = _writerSize &+ size
    }
  }

  /// Adds an object of type Scalar into the buffer
  /// - Parameters:
  ///   - value: Object  that will be written to the buffer
  ///   - len: Offset to subtract from the WriterIndex
  @inline(__always)
  @usableFromInline
  mutating func push<T: Scalar>(value: T, len: Int) {
    ensureSpace(size: len)
    withUnsafePointer(to: value) {
      memcpy(
        _storage.memory.advanced(by: writerIndex &- len),
        $0,
        len)
      _writerSize = _writerSize &+ len
    }
  }

  /// Adds a RawPointer into the buffer
  /// - Parameter pointer: pointer to be copied into the buffer
  /// - Parameter len: length of the data
  @inline(__always)
  mutating func writeBytes(_ ptr: UnsafeRawPointer, len: Int) {
    ensureSpace(size: len)
    memcpy(
      _storage.memory.advanced(by: writerIndex &- len),
      ptr,
      len)
    _writerSize = _writerSize &+ len
  }

  /// Write stores an object into the buffer directly or indirectly.
  ///
  /// Direct: ignores the capacity of buffer which would mean we are referring to the direct point in memory
  /// indirect: takes into respect the current capacity of the buffer (capacity - index), writing to the buffer from the end
  /// - Parameters:
  ///   - value: Value that needs to be written to the buffer
  ///   - index: index to write to
  ///   - direct: Should take into consideration the capacity of the buffer
  @inline(__always)
  func write<T>(value: T, index: Int, direct: Bool = false) {
    var index = index
    if !direct {
      index = capacity &- index
    }
    assert(index < capacity, "Write index is out of writing bound")
    assert(index >= 0, "Writer index should be above zero")
    _ = withUnsafePointer(to: value) {
      memcpy(
        _storage.memory.advanced(by: index),
        $0,
        MemoryLayout<T>.size)
    }
  }

  /// Makes sure that buffer has enouch space for each of the objects that will be written into it
  /// - Parameter size: size of object
  @discardableResult
  @usableFromInline
  mutating func ensureSpace(size: Int) -> Int {
    let expectedWriterIndex = size &+ _writerSize
    if expectedWriterIndex > capacity {

      let currentWritingIndex = capacity &- _writerSize
      while capacity <= expectedWriterIndex {
        capacity = capacity << 1
      }

      /// solution take from Apple-NIO
      capacity = capacity.convertToPowerofTwo

      _storage.reallocate(
        capacity: capacity,
        writerSize: _writerSize,
        currentWritingIndex: currentWritingIndex,
        alignment: alignment)
    }
    assert(size < FlatBufferMaxSize, "Buffer can't grow beyond 2 Gigabytes")
    return size
  }

  /// pops the written VTable if it's already written into the buffer
  /// - Parameter size: size of the `VTable`
  @usableFromInline
  @inline(__always)
  mutating func pop(_ size: Int) {
    assert(
      (_writerSize &- size) > 0,
      "New size should NOT be a negative number")
    memset(_storage.memory.advanced(by: writerIndex), 0, _writerSize &- size)
    _writerSize = size
  }

  /// Clears the current size of the buffer
  @inline(__always)
  mutating public func clearSize() {
    _writerSize = 0
  }

  /// Clears the current instance of the buffer, replacing it with new memory
  @inline(__always)
  mutating public func clear(keepingCapacity: Bool = false) {
    _writerSize = 0
    alignment = 1
    if keepingCapacity {
      _storage.initialize(for: capacity)
    } else {
      capacity = initialSize
      _storage = Storage(count: initialSize, alignment: alignment)
    }
  }

  /// Reads an object from the buffer
  /// - Parameters:
  ///   - def: Type of the object
  ///   - position: the index of the object in the buffer
  @inline(__always)
  public func read<T>(def: T.Type, position: Int) -> T {
    _storage.memory
      .advanced(by: position)
      .bindMemory(to: T.self, capacity: 1)
      .pointee
  }

  @discardableResult
  @inline(__always)
  func withUnsafeBytes<T>(
    _ body: (UnsafeRawBufferPointer) throws
      -> T) rethrows -> T
  {
    try body(
      UnsafeRawBufferPointer(
        start: _storage.memory,
        count: capacity))
  }

  @discardableResult
  @inline(__always)
  func withUnsafeSlicedBytes<T>(
    _ body: (UnsafeRawBufferPointer) throws
      -> T) rethrows -> T
  {
    try body(
      UnsafeRawBufferPointer(
        start: _storage.memory.advanced(by: writerIndex),
        count: capacity &- writerIndex))
  }

  @discardableResult
  @inline(__always)
  func withUnsafeRawPointer<T>(
    _ body: (UnsafeMutableRawPointer) throws
      -> T) rethrows -> T
  {
    try body(_storage.memory)
  }

  @discardableResult
  @inline(__always)
  func readWithUnsafeRawPointer<T>(
    position: Int,
    _ body: (UnsafeRawPointer) throws -> T) rethrows -> T
  {
    try body(_storage.memory.advanced(by: position))
  }
}

extension _InternalByteBuffer: CustomDebugStringConvertible {

  public var debugDescription: String {
    """
    buffer located at: \(_storage.memory), with capacity of \(capacity)
    { writerSize: \(_writerSize), readerSize: \(reader), writerIndex: \(
      writerIndex) }
    """
  }
}