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

/// `ByteBuffer` is the interface that stores the data for a `Flatbuffers` object
/// it allows users to write and read data directly from memory thus the use of its
/// functions should be used
@frozen
public struct ByteBuffer {

  /// 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 {
    @usableFromInline
    enum Blob {
      #if !os(WASI)
      case data(Data)
      case bytes(ContiguousBytes)
      #endif

      case byteBuffer(_InternalByteBuffer)
      case array([UInt8])
      case pointer(UnsafeMutableRawPointer)
    }

    /// This storage doesn't own the memory, therefore, we won't deallocate on deinit.
    private let isOwned: Bool
    /// Capacity of UInt8 the buffer can hold
    private let capacity: Int
    /// Retained blob of data that requires the storage to retain a pointer to.
    @usableFromInline
    var retainedBlob: Blob

    @usableFromInline
    init(count: Int) {
      let memory = UnsafeMutableRawPointer.allocate(
        byteCount: count,
        alignment: MemoryLayout<UInt8>.alignment)
      capacity = count
      retainedBlob = .pointer(memory)
      isOwned = true
    }

    @usableFromInline
    init(blob: Blob, capacity count: Int) {
      capacity = count
      retainedBlob = blob
      isOwned = false
    }

    deinit {
      guard isOwned else { return }
      switch retainedBlob {
      case .pointer(let unsafeMutableRawPointer):
        unsafeMutableRawPointer.deallocate()
      default: break
      }
    }

    @usableFromInline
    func copy(from ptr: UnsafeRawPointer, count: Int) {
      assert(
        isOwned,
        "copy should NOT be called on a buffer that is built by assumingMemoryBound")
      withUnsafeRawPointer {
        $0.copyMemory(from: ptr, byteCount: count)
      }
    }

    @usableFromInline
    func initialize(for size: Int) {
      assert(
        isOwned,
        "initalize should NOT be called on a buffer that is built by assumingMemoryBound")
      withUnsafeRawPointer {
        memset($0, 0, size)
      }
    }

    @discardableResult
    @inline(__always)
    func withUnsafeBytes<T>(
      _ body: (UnsafeRawBufferPointer) throws
        -> T) rethrows -> T
    {
      switch retainedBlob {
      case .byteBuffer(let byteBuffer):
        return try byteBuffer.withUnsafeBytes(body)
      #if !os(WASI)
      case .data(let data):
        return try data.withUnsafeBytes(body)
      case .bytes(let contiguousBytes):
        return try contiguousBytes.withUnsafeBytes(body)
      #endif
      case .array(let array):
        return try array.withUnsafeBytes(body)
      case .pointer(let ptr):
        return try body(UnsafeRawBufferPointer(start: ptr, count: capacity))
      }
    }

    @discardableResult
    @inline(__always)
    func withUnsafeRawPointer<T>(
      _ body: (UnsafeMutableRawPointer) throws
        -> T) rethrows -> T
    {
      switch retainedBlob {
      case .byteBuffer(let byteBuffer):
        return try byteBuffer.withUnsafeRawPointer(body)
      #if !os(WASI)
      case .data(let data):
        return try data
          .withUnsafeBytes {
            try body(UnsafeMutableRawPointer(mutating: $0.baseAddress!))
          }
      case .bytes(let contiguousBytes):
        return try contiguousBytes
          .withUnsafeBytes {
            try body(UnsafeMutableRawPointer(mutating: $0.baseAddress!))
          }
      #endif
      case .array(let array):
        return try array
          .withUnsafeBytes {
            try body(UnsafeMutableRawPointer(mutating: $0.baseAddress!))
          }
      case .pointer(let ptr):
        return try body(ptr)
      }
    }

    @discardableResult
    @inline(__always)
    func readWithUnsafeRawPointer<T>(
      position: Int,
      _ body: (UnsafeRawPointer) throws -> T) rethrows -> T
    {
      switch retainedBlob {
      case .byteBuffer(let byteBuffer):
        return try byteBuffer.readWithUnsafeRawPointer(position: position, body)
      #if !os(WASI)
      case .data(let data):
        return try data.withUnsafeBytes {
          try body($0.baseAddress!.advanced(by: position))
        }
      case .bytes(let contiguousBytes):
        return try contiguousBytes.withUnsafeBytes {
          try body($0.baseAddress!.advanced(by: position))
        }
      #endif
      case .array(let array):
        return try array.withUnsafeBytes {
          try body($0.baseAddress!.advanced(by: position))
        }
      case .pointer(let ptr):
        return try body(ptr.advanced(by: position))
      }
    }
  }

  @usableFromInline var _storage: Storage

  /// The size of the elements written to the buffer + their paddings
  private var _readerIndex: Int = 0
  /// Reader is the position of the current Writer Index (capacity - size)
  public var reader: Int { capacity &- _readerIndex }
  /// Current size of the buffer
  public var size: UOffset { UOffset(_readerIndex) }
  /// Current capacity for the buffer
  public let capacity: Int

  /// Constructor that creates a Flatbuffer object from an InternalByteBuffer
  /// - Parameter
  ///   - bytes: Array of UInt8
  @inline(__always)
  init(byteBuffer: _InternalByteBuffer) {
    _storage = Storage(
      blob: .byteBuffer(byteBuffer),
      capacity: byteBuffer.capacity)
    _readerIndex = Int(byteBuffer.size)
    self.capacity = byteBuffer.capacity
  }

  /// Constructor that creates a Flatbuffer from unsafe memory region by copying
  /// the underlying data to a new pointer
  ///
  /// - Parameters:
  ///   - copyingMemoryBound: The unsafe memory region
  ///   - capacity: The size of the given memory region
  @inline(__always)
  public init(
    copyingMemoryBound memory: UnsafeRawPointer,
    capacity: Int)
  {
    _storage = Storage(count: capacity)
    _storage.copy(from: memory, count: capacity)
    _readerIndex = capacity
    self.capacity = capacity
  }

  /// Constructor that creates a Flatbuffer object from a UInt8
  /// - Parameter
  ///   - bytes: Array of UInt8
  @inline(__always)
  public init(bytes: [UInt8]) {
    _storage = Storage(blob: .array(bytes), capacity: bytes.count)
    _readerIndex = bytes.count
    capacity = bytes.count
  }

  #if !os(WASI)
  /// Constructor that creates a Flatbuffer from the Swift Data type object
  /// - Parameter
  ///   - data: Swift data Object
  @inline(__always)
  public init(data: Data) {
    _storage = Storage(blob: .data(data), capacity: data.count)
    _readerIndex = data.count
    capacity = data.count
  }

  /// Constructor that creates a Flatbuffer object from a ContiguousBytes
  /// - Parameters:
  ///   - contiguousBytes: Binary stripe to use as the buffer
  ///   - count: amount of readable bytes
  @inline(__always)
  public init<Bytes: ContiguousBytes>(
    contiguousBytes: Bytes,
    count: Int)
  {
    _storage = Storage(blob: .bytes(contiguousBytes), capacity: count)
    _readerIndex = count
    self.capacity = count
  }
  #endif

  /// Constructor that creates a Flatbuffer from unsafe memory region without copying
  /// **NOTE** Needs a call to `memory.deallocate()` later on to free the memory
  ///
  /// - Parameters:
  ///   - assumingMemoryBound: The unsafe memory region
  ///   - capacity: The size of the given memory region
  @inline(__always)
  public init(
    assumingMemoryBound memory: UnsafeMutableRawPointer,
    capacity: Int)
  {
    _storage = Storage(
      blob: .pointer(memory),
      capacity: capacity)
    _readerIndex = capacity
    self.capacity = capacity
  }

  /// Creates a copy of the existing flatbuffer, by copying it to a different memory.
  /// - Parameters:
  ///   - memory: Current memory of the buffer
  ///   - count: count of bytes
  ///   - removeBytes: Removes a number of bytes from the current size
  @inline(__always)
  init(
    blob: Storage.Blob,
    count: Int,
    removing removeBytes: Int)
  {
    _storage = Storage(blob: blob, capacity: count)
    _readerIndex = removeBytes
    capacity = count
  }

  /// 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) { ptr in
      _storage.withUnsafeRawPointer {
        memcpy(
          $0.advanced(by: index),
          ptr,
          MemoryLayout<T>.size)
      }
    }
  }

  /// 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.readWithUnsafeRawPointer(position: position) {
      $0.bindMemory(to: T.self, capacity: 1)
        .pointee
    }
  }

  /// Reads a slice from the memory assuming a type of T
  /// - Parameters:
  ///   - index: index of the object to be read from the buffer
  ///   - count: count of bytes in memory
  @inline(__always)
  public func readSlice<T>(
    index: Int,
    count: Int) -> [T]
  {
    assert(
      index + count <= capacity,
      "Reading out of bounds is illegal")

    return _storage.readWithUnsafeRawPointer(position: index) {
      let buf = UnsafeBufferPointer(
        start: $0.bindMemory(to: T.self, capacity: count),
        count: count)
      return Array(buf)
    }
  }

  /// Provides a pointer towards the underlying primitive types
  /// - Parameters:
  ///   - index: index of the object to be read from the buffer
  ///   - count: count of bytes in memory
  @discardableResult
  @inline(__always)
  public func withUnsafePointerToSlice<T>(
    index: Int,
    count: Int,
    body: (UnsafeRawBufferPointer) throws -> T) rethrows -> T
  {
    assert(
      index + count <= capacity,
      "Reading out of bounds is illegal")
    return try _storage.readWithUnsafeRawPointer(position: index) {
      try body(UnsafeRawBufferPointer(start: $0, count: count))
    }
  }

  #if !os(WASI)
  /// Reads a string from the buffer and encodes it to a swift string
  /// - Parameters:
  ///   - index: index of the string in the buffer
  ///   - count: length of the string
  ///   - type: Encoding of the string
  @inline(__always)
  public func readString(
    at index: Int,
    count: Int,
    type: String.Encoding = .utf8) -> String?
  {
    assert(
      index + count <= capacity,
      "Reading out of bounds is illegal")
    return _storage.readWithUnsafeRawPointer(position: index) {
      let buf = UnsafeBufferPointer(
        start: $0.bindMemory(to: UInt8.self, capacity: count),
        count: count)
      return String(
        bytes: buf,
        encoding: type)
    }
  }
  #else
  /// Reads a string from the buffer and encodes it to a swift string
  /// - Parameters:
  ///   - index: index of the string in the buffer
  ///   - count: length of the string
  @inline(__always)
  public func readString(
    at index: Int,
    count: Int) -> String?
  {
    assert(
      index + count <= capacity,
      "Reading out of bounds is illegal")
    return _storage.readWithUnsafeRawPointer(position: index) {
      String(cString: $0.bindMemory(to: UInt8.self, capacity: count))
    }
  }
  #endif

  /// Creates a new Flatbuffer object that's duplicated from the current one
  /// - Parameter removeBytes: the amount of bytes to remove from the current Size
  @inline(__always)
  public func duplicate(removing removeBytes: Int = 0) -> ByteBuffer {
    assert(removeBytes > 0, "Can NOT remove negative bytes")
    assert(
      removeBytes < capacity,
      "Can NOT remove more bytes than the ones allocated")
    return ByteBuffer(
      blob: _storage.retainedBlob,
      count: capacity,
      removing: _readerIndex &- removeBytes)
  }

  /// SkipPrefix Skips the first 4 bytes in case one of the following
  /// functions are called `getPrefixedSizeCheckedRoot` & `getPrefixedSizeRoot`
  /// which allows us to skip the first 4 bytes instead of recreating the buffer
  @discardableResult
  @usableFromInline
  @inline(__always)
  mutating func skipPrefix() -> Int32 {
    _readerIndex = _readerIndex &- MemoryLayout<Int32>.size
    return read(def: Int32.self, position: 0)
  }

  @discardableResult
  @inline(__always)
  public func withUnsafeBytes<T>(
    body: (UnsafeRawBufferPointer) throws
      -> T) rethrows -> T
  {
    try _storage.withUnsafeBytes(body)
  }

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

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

extension ByteBuffer: CustomDebugStringConvertible {

  public var debugDescription: String {
    """
    buffer located at: \(_storage.retainedBlob), 
    with capacity of \(capacity),
    { writtenSize: \(_readerIndex), readerSize: \(reader), 
    size: \(size) }
    """
  }
}