/*
 * 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
  /// Current size of the buffer
  public var count: Int { _readerIndex }
  /// Current capacity for the buffer including unused space
  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 = byteBuffer.writerIndex
    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
    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
    }
  }

  @inline(__always)
  public func readUInt64(offset: Int, byteWidth: UInt8) -> UInt64 {
    readSizedScalar(
      def: UInt64.self,
      t1: UInt8.self,
      t2: UInt16.self,
      t3: UInt32.self,
      t4: UInt64.self,
      position: offset,
      byteWidth: byteWidth)
  }

  @inline(__always)
  public func readInt64(offset: Int, byteWidth: UInt8) -> Int64 {
    readSizedScalar(
      def: Int64.self,
      t1: Int8.self,
      t2: Int16.self,
      t3: Int32.self,
      t4: Int64.self,
      position: offset,
      byteWidth: byteWidth)
  }

  @inline(__always)
  public func readDouble(offset: Int, byteWidth: UInt8) -> Double {
    switch byteWidth {
    case 4:
      Double(read(def: Float32.self, position: offset))
    default:
      read(def: Double.self, position: offset)
    }
  }

  @inline(__always)
  func readSizedScalar<
    T: BinaryInteger,
    T1: BinaryInteger,
    T2: BinaryInteger,
    T3: BinaryInteger,
    T4: BinaryInteger
  >(
    def: T.Type,
    t1: T1.Type,
    t2: T2.Type,
    t3: T3.Type,
    t4: T4.Type,
    position: Int,
    byteWidth: UInt8) -> T
  {
    switch byteWidth {
    case 1:
      numericCast(read(def: T1.self, position: position))
    case 2:
      numericCast(read(def: T2.self, position: position))
    case 4:
      numericCast(read(def: T3.self, position: position))
    default:
      numericCast(read(def: T4.self, position: position))
    }
  }

  /// 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)
    }
  }

  @inline(__always)
  public func readString(
    at index: Int,
    count: Int,
    type: String.Encoding) -> 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)
    }
  }

  /// 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))
    }
  }

  /// 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))
    }
  }

  @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)
  }
}