ResumableUpload.swift

//
//  ResumableUpload.swift
//  Storage
//
//  Created by Guilherme Souza on 15/10/26.
//

import ConcurrencyExtras
import Foundation
import HTTPTypes
import Logging

/// A tus upload that can be paused, resumed, cancelled, and continued after the process
/// restarts.
///
/// Created by ``StorageBucket/resumableUpload(path:fileURL:options:)`` and
/// ``StorageBucket/resumableUpload(path:data:options:)``. Nothing is sent until ``start()``:
///
/// ```swift
/// let upload = try bucket.resumableUpload(path: "video.mp4", fileURL: url)
/// try await upload.start()
/// for await progress in upload.progress() { … }
/// let object = try await upload.result
/// ```
///
/// The handle outlives whoever created it, like a `Task`. Cancelling a task that awaits
/// ``start()``, ``result``, ``progress()`` or ``statusChanges()`` only ends that await; the
/// upload keeps going until ``pause()``, ``cancel()``, completion or failure. After a network failure the
/// upload retries, with a `HEAD` for the server's offset before each attempt, until
/// ``ResumableUploadConfiguration/retryTimeout`` is spent; then it fails and its record stays
/// on disk for ``resume()``, in this process or the next.
///
/// - Warning: Experimental. Opt in with `@_spi(Experimental) import Supabase` (or `import
///   Storage`); the API may change in a future release.
@_spi(Experimental)
public actor ResumableUpload {
  /// The identity of an upload: the same file to the same path in the same bucket has the same
  /// `ID`, whichever URL it is picked through.
  public struct ID: RawRepresentable, Hashable, Sendable {
    /// The fingerprint of the endpoint, bucket, path and source.
    public let rawValue: String

    /// Creates an ``ID`` from a fingerprint.
    public init(rawValue: String) {
      self.rawValue = rawValue
    }
  }

  /// Where an upload is.
  public enum Status: Sendable {
    /// Created, not started.
    case pending
    /// Sending chunks, or waiting to retry.
    case uploading
    /// Stopped by ``ResumableUpload/pause()``; ``ResumableUpload/resume()`` continues it.
    case paused
    /// Every byte is on the server.
    case completed(UploadedObject)
    /// Gave up. A transport failure keeps the record, so ``ResumableUpload/resume()`` can try
    /// again; a server rejection removes it.
    case failed(StorageError)
    /// Stopped by ``ResumableUpload/cancel()``; the server's copy and the record are gone, unless
    /// another handle for the same upload drives them.
    case cancelled

    var isTerminal: Bool {
      switch self {
      case .completed, .failed, .cancelled: true
      case .pending, .uploading, .paused: false
      }
    }
  }

  enum Input: Sendable {
    case file(URL)
    case data(Data)
    case prepared(ResumableSource)
  }

  private enum Next {
    case append
    case resync
    case recreate
  }

  /// The identity of this upload. Asking for the same upload again returns the same `ID`.
  public nonisolated let id: ID
  /// The fingerprint behind ``id``, which names this upload's files in the store.
  nonisolated let fingerprint: Fingerprint
  /// The bucket the object goes to.
  public nonisolated let bucketID: String
  /// The path of the object in the bucket.
  public nonisolated let path: String
  /// The size of the upload.
  public nonisolated let totalBytes: Int64

  /// Where the upload is now. ``statusChanges()`` reports every change.
  public private(set) var status: Status {
    didSet {
      if oldValue.isTerminal, !status.isTerminal {
        // resume() revives a failed upload, so new observers follow it again.
        statusBroadcaster.reopen(with: status)
        progressBroadcaster.reopen()
      } else {
        statusBroadcaster.yield(status)
      }
      if status.isTerminal {
        statusBroadcaster.finish()
        progressBroadcaster.finish()
        // cancel() drops the entry itself once its cleanup ends, so no second handle takes the
        // record while it is being deleted.
        if case .cancelled = status {} else { dropLive() }
      }
    }
  }

  private let api: StorageAPI
  private let store: ResumableStore
  private let options: UploadOptions
  private let input: Input
  private var record: ResumableUploadRecord?
  private var loop: Task<Void, Never>?
  private var cancellation: Task<Void, Never>?
  private var pauseRequested = false
  private var startInProgress = false
  private var startWaiters: [CheckedContinuation<Void, Never>] = []
  private let progressBroadcaster = StreamBroadcaster<TransferProgress>(
    bufferingPolicy: .bufferingNewest(1))
  private let statusBroadcaster: StreamBroadcaster<Status>

  /// The handles that are not yet terminal, one per record, so two callers asking for the
  /// same upload drive one tus session instead of two. Held weakly, so a handle that never
  /// starts lives only as long as its caller keeps it.
  private struct Live {
    weak var upload: ResumableUpload?
    var started = false
  }

  private static let live = LockIsolated<[String: Live]>([:])

  private nonisolated var liveKey: String { Self.liveKey(fingerprint: fingerprint, api: api) }

  private static func liveKey(fingerprint: Fingerprint, api: StorageAPI) -> String {
    ResumableStore(configuration: api.configuration.resumable).directory
      .appendingPathComponent(fingerprint.rawValue).path
  }

  /// The live handle for `fingerprint`, or the one `make` builds.
  static func handle(
    fingerprint: Fingerprint, api: StorageAPI, make: @Sendable () -> ResumableUpload
  ) -> ResumableUpload {
    let key = liveKey(fingerprint: fingerprint, api: api)
    return live.withValue { live in
      if let existing = live[key]?.upload { return existing }
      let upload = make()
      live = live.filter { $0.value.upload != nil }
      live[key] = Live(upload: upload)
      return upload
    }
  }

  /// Removes `record` unless a handle for it has started and is not yet terminal in this
  /// process. Holds the lock ``start()`` claims the record under, so a start never loses its
  /// record to a concurrent removal.
  static func removeUnlessLive(
    _ record: ResumableUploadRecord, from store: ResumableStore, api: StorageAPI
  ) {
    live.withValue { live in
      let entry = live[liveKey(fingerprint: record.fingerprint, api: api)]
      if entry?.upload != nil, entry?.started == true { return }
      store.remove(record)
    }
  }

  /// The live handle for `record`, or a pending one built from it.
  static func handle(record: ResumableUploadRecord, api: StorageAPI) -> ResumableUpload {
    handle(fingerprint: record.fingerprint, api: api) {
      ResumableUpload(
        fingerprint: record.fingerprint, bucketID: record.bucketID, path: record.path,
        totalBytes: record.totalBytes, input: .prepared(record.source),
        options: UploadOptions(
          contentType: record.contentType,
          cacheControl: CacheControl(rawValue: record.cacheControl),
          upsert: record.upsert, metadata: record.metadata),
        api: api)
    }
  }

  init(
    fingerprint: Fingerprint, bucketID: String, path: String, totalBytes: Int64, input: Input,
    options: UploadOptions, api: StorageAPI
  ) {
    self.id = ID(rawValue: fingerprint.rawValue)
    self.fingerprint = fingerprint
    self.bucketID = bucketID
    self.path = path
    self.totalBytes = totalBytes
    self.input = input
    self.options = options
    self.api = api
    self.store = ResumableStore(configuration: api.configuration.resumable)
    let initial = Status.pending
    self.status = initial
    self.statusBroadcaster = StreamBroadcaster(initial: initial, bufferingPolicy: .unbounded)
  }

  /// The server-confirmed bytes, as each chunk lands. A late subscriber gets the latest value
  /// first; the stream ends with the upload.
  public nonisolated func progress() -> AsyncStream<TransferProgress> {
    progressBroadcaster.stream()
  }

  /// The current ``status`` first, then every change, ending after a terminal one.
  public nonisolated func statusChanges() -> AsyncStream<Status> { statusBroadcaster.stream() }

  /// The uploaded object once every byte is on the server.
  ///
  /// - Throws: The ``StorageError`` the upload failed with, or `CancellationError` after
  ///   ``cancel()`` or when the awaiting task is cancelled.
  public var result: UploadedObject {
    get async throws {
      for await status in statusChanges() {
        if let outcome = try outcome(of: status) { return outcome }
      }
      throw CancellationError()
    }
  }

  private func outcome(of status: Status) throws -> UploadedObject? {
    switch status {
    case .completed(let object): object
    case .failed(let error): throw error
    case .cancelled: throw CancellationError()
    case .pending, .uploading, .paused: nil
    }
  }

  /// Creates the upload on the server, or finds the one a previous run left, and starts
  /// sending chunks. Returns as soon as the first request is answered; ``result`` awaits the
  /// rest. Does nothing once started.
  ///
  /// - Throws: ``StorageError`` when the server refuses the upload, in which case ``status`` is
  ///   ``Status/failed(_:)``.
  public func start() async throws {
    guard case .pending = status, !startInProgress else { return }
    startInProgress = true
    let key = liveKey
    let claimed = claimLive(key)
    defer {
      startInProgress = false
      startWaiters.forEach { $0.resume() }
      startWaiters.removeAll()
      switch status {
      case .uploading, .cancelled: break
      default: releaseStart(key)
      }
    }
    // Another handle has started on this upload and drives it; a second loop would race it.
    guard claimed else { return }
    do {
      let record = try await loadOrCreateRecord()
      self.record = record
      guard case .pending = status else { return }
      status = .uploading
      startLoop(from: record.lastKnownOffset, next: .append)
    } catch is CancellationError {
      throw CancellationError()
    } catch {
      if case .cancelled = status { return }
      let error = storageError(error)
      status = .failed(error)
      throw error
    }
  }

  /// Marks this handle as the one driving its record, so ``ResumableUploadRegistry/clear()``
  /// keeps the record and a second handle is never made for it. Takes the entry from another
  /// handle that has not started, such as one ``ResumableUploadRegistry/pending()`` made after
  /// this one failed. Returns `false` when another handle has started on the record.
  @discardableResult
  private func claimLive(_ key: String) -> Bool {
    Self.live.withValue { live in
      let entry = live[key]
      guard entry?.upload == nil || entry?.upload === self || entry?.started != true else {
        return false
      }
      live[key] = Live(upload: self, started: true)
      return true
    }
  }

  /// Gives up the claim ``start()`` took on the record when it ends without uploading.
  private func releaseStart(_ key: String) {
    Self.live.withValue { live in
      if live[key]?.upload === self { live[key]?.started = false }
    }
  }

  /// Stops after the chunk in flight lands. The record stays, so ``resume()`` continues from
  /// the server's offset.
  public func pause() {
    guard case .uploading = status else { return }
    pauseRequested = true
    status = .paused
  }

  /// Continues a paused or failed upload from the offset the server reports. When ``start()``
  /// failed with a transport error before the upload existed on the server, creates it again.
  ///
  /// - Throws: The ``StorageError`` the server refused the upload with when ``start()`` failed
  ///   that way, or the error the new attempt fails with.
  public func resume() async throws {
    guard canResume else { return }
    await loop?.value
    guard canResume else { return }
    guard record != nil else { return try await startAgain() }
    claimLive(liveKey)
    status = .uploading
    startLoop(from: 0, next: .resync)
  }

  /// Runs ``start()`` again after it failed before the upload existed on the server. Only a
  /// transport failure is worth another attempt; the server would refuse the same upload again.
  private func startAgain() async throws {
    if case .failed(let error) = status, error.kind != .transport { throw error }
    // Another handle made the upload since; it drives it, and this handle stays failed.
    guard claimLive(liveKey) else { return }
    // .pending, not a terminal status, so cancel() and statusChanges() follow the new attempt.
    status = .pending
    try await start()
  }

  private var canResume: Bool {
    switch status {
    case .paused, .failed: true
    case .pending, .uploading, .completed, .cancelled: false
    }
  }

  /// Aborts the upload: the chunk in flight is dropped, the server's copy is deleted (best
  /// effort), and so are the record and any sandbox copy. When ``start()`` is still creating or
  /// looking up the upload, waits for that request to end first, so the server's copy it makes
  /// is deleted too. A call made while an earlier one is still cleaning up waits for it.
  /// Cancelling the calling task does not stop the cleanup.
  ///
  /// A handle that never loaded the record, such as one ``ResumableUploadRegistry/pending()``
  /// made, leaves the record and the server's copy alone when another handle for the same
  /// upload has started; only this handle is cancelled, and the other one goes on.
  public func cancel() async {
    if let cancellation { return await cancellation.value }
    guard !status.isTerminal else { return }
    loop?.cancel()
    status = .cancelled
    let cleanup = Task { await cleanUpAfterCancel() }
    cancellation = cleanup
    await cleanup.value
  }

  private func cleanUpAfterCancel() async {
    if startInProgress {
      // Waits for the record start() is loading or creating, so the discard below covers it.
      await withCheckedContinuation { startWaiters.append($0) }
    }
    await loop?.value
    if let record {
      await discard(record)
    } else if claimLive(liveKey) {
      // Without a record of its own, this handle deletes the one on disk only when no other
      // handle has started on it; otherwise the record and the server's copy are that handle's.
      if let record = try? store.record(fingerprint: fingerprint) {
        await discard(record)
      } else {
        store.removeFiles(fingerprint: fingerprint)
      }
    }
    dropLive()
  }

  private func dropLive() {
    Self.live.withValue { if $0[liveKey]?.upload === self { $0.removeValue(forKey: liveKey) } }
  }

  /// Deletes the server's copy (best effort), the record, and any sandbox copy.
  private func discard(_ record: ResumableUploadRecord) async {
    await TusSession(api: api).terminate(record.uploadURL)
    store.remove(record)
  }

  var isRunning: Bool { loop != nil }

  private func startLoop(from offset: Int64, next: Next) {
    pauseRequested = false
    loop = Task { await run(from: offset, next: next) }
  }

  private func run(from start: Int64, next first: Next) async {
    let tus = TusSession(api: api)
    var offset = start
    var next = first
    var retry = 0
    var waited = Duration.zero
    while !pauseRequested, let record {
      do {
        switch next {
        case .resync:
          offset = try await tus.offset(of: record.uploadURL)
          next = .append
        case .recreate:
          offset = try await recreate(record).lastKnownOffset
          next = .append
        case .append:
          guard let more = try await append(at: offset, with: tus) else { break }
          offset = more
          retry = 0
          waited = .zero
          continue
        }
      } catch is CancellationError {
        break
      } catch TusSession.Failure.offsetConflict {
        next = .resync
      } catch TusSession.Failure.uploadGone {
        next = .recreate
      } catch let error as StorageError where error.kind == .transport {
        retry += 1
        let delay = RetryPolicy.default.backoffDelay(retry: retry)
        waited += delay
        guard waited <= api.configuration.resumable.retryTimeout else {
          fail(error, keepRecord: true)
          break
        }
        api.configuration.logger.debug(
          "Resumable upload \(id.rawValue) retrying in \(delay) after: \(error.message)")
        try? await api.configuration.clock.sleep(for: delay)
        if Task.isCancelled { break }
        next = .resync
      } catch {
        fail(storageError(error), keepRecord: false)
        break
      }
      if status.isTerminal { break }
    }
    loop = nil
  }

  /// Sends the chunk at `offset` and returns the server's new offset, or `nil` once the upload
  /// is complete.
  private func append(at offset: Int64, with tus: TusSession) async throws -> Int64? {
    guard var record else { return nil }
    if offset < totalBytes {
      let range = FileRange(
        fileURL: store.fileURL(for: record), offset: offset,
        count: min(record.chunkSize, totalBytes - offset))
      let chunk = HTTPBody(range, length: .known(range.count), iterationBehavior: .multiple)
      let appended = try await tus.append(chunk, at: offset, to: record.uploadURL)
      record.lastKnownOffset = appended.offset
      self.record = record
      try store.save(record)
      progressBroadcaster.yield(
        TransferProgress(completedBytes: appended.offset, totalBytes: totalBytes))
      guard appended.complete || appended.offset >= totalBytes else { return appended.offset }
      complete(with: appended.body)
    } else {
      complete(with: Data())
    }
    return nil
  }

  private func complete(with body: Data) {
    let response = try? JSONDecoder.storage.decode(UploadResponse.self, from: body)
    let object = UploadedObject(
      id: response?.id.flatMap(UUID.init(uuidString:)),
      path: path,
      fullPath: response?.key ?? "\(bucketID)/\(path)")
    if let record { store.remove(record) }
    status = .completed(object)
  }

  private func fail(_ error: StorageError, keepRecord: Bool) {
    if !keepRecord, let record { store.remove(record) }
    status = .failed(error)
  }

  private func loadOrCreateRecord() async throws -> ResumableUploadRecord {
    let tus = TusSession(api: api)
    if var existing = try store.record(fingerprint: fingerprint) {
      do {
        existing.lastKnownOffset = try await tus.offset(of: existing.uploadURL)
        try store.save(existing)
        return existing
      } catch TusSession.Failure.uploadGone {
        return try await recreate(existing)
      }
    }
    let source =
      switch input {
      case .file(let fileURL): try await store.prepare(fileURL: fileURL, fingerprint: fingerprint)
      case .data(let data): try store.prepare(data: data, fingerprint: fingerprint)
      case .prepared(let source): source
      }
    return try await createRecord(source: source, with: tus)
  }

  /// Creates the upload again on the server for a record the server no longer knows.
  private func recreate(_ stale: ResumableUploadRecord) async throws -> ResumableUploadRecord {
    try? FileManager.default.removeItem(at: store.recordURL(fingerprint: stale.fingerprint))
    let record = try await createRecord(source: stale.source, with: TusSession(api: api))
    self.record = record
    return record
  }

  private func createRecord(source: ResumableSource, with tus: TusSession) async throws
    -> ResumableUploadRecord
  {
    let key = try ObjectKey(path)
    let (_, modificationDate) = try ResumableStore.attributes(
      of: store.fileURL(for: source, fingerprint: fingerprint))
    let created = try await tus.create(
      bucket: bucketID, key: key, length: totalBytes, options: options)
    let now = Date()
    let record = ResumableUploadRecord(
      fingerprint: fingerprint,
      bucketID: bucketID,
      path: path,
      totalBytes: totalBytes,
      source: source,
      sourceModificationDate: modificationDate,
      uploadURL: created.uploadURL,
      createdAt: now,
      expiresAt: created.expiresAt ?? now.addingTimeInterval(ResumableStore.defaultLifetime),
      lastKnownOffset: created.offset,
      chunkSize: api.configuration.resumable.chunkSize,
      upsert: options.upsert,
      contentType: options.contentType,
      cacheControl: options.cacheControl.rawValue,
      metadata: options.metadata)
    try store.save(record)
    return record
  }

  private nonisolated func storageError(_ error: any Error) -> StorageError {
    error as? StorageError
      ?? StorageError(
        kind: .invalidRequest, message: "The upload source could not be read.",
        underlyingError: error)
  }
}

/// Fans one producer out to any number of `AsyncStream` consumers. A new consumer gets the
/// latest element first, or `initial` before the first yield.
final class StreamBroadcaster<Element: Sendable>: Sendable {
  private struct Subscribers {
    var continuations: [UUID: AsyncStream<Element>.Continuation] = [:]
    var latest: Element?
    var finished = false
  }

  private let subscribers: LockIsolated<Subscribers>
  private let bufferingPolicy: AsyncStream<Element>.Continuation.BufferingPolicy

  init(
    initial: Element? = nil,
    bufferingPolicy: AsyncStream<Element>.Continuation.BufferingPolicy
  ) {
    self.subscribers = LockIsolated(Subscribers(latest: initial))
    self.bufferingPolicy = bufferingPolicy
  }

  func stream() -> AsyncStream<Element> {
    let (stream, continuation) = AsyncStream.makeStream(
      of: Element.self, bufferingPolicy: bufferingPolicy)
    let id = UUID()
    subscribers.withValue { subscribers in
      if let latest = subscribers.latest { continuation.yield(latest) }
      if subscribers.finished {
        continuation.finish()
      } else {
        subscribers.continuations[id] = continuation
      }
    }
    continuation.onTermination = { [subscribers] _ in
      subscribers.withValue { _ = $0.continuations.removeValue(forKey: id) }
    }
    return stream
  }

  func yield(_ element: Element) {
    subscribers.withValue { subscribers in
      subscribers.latest = element
      for continuation in subscribers.continuations.values { continuation.yield(element) }
    }
  }

  /// Lets new consumers follow later elements again after ``finish()``. With `element`, it
  /// becomes the latest in the same step, so no new consumer sees the one before it.
  func reopen(with element: Element? = nil) {
    subscribers.withValue { subscribers in
      subscribers.finished = false
      if let element { subscribers.latest = element }
    }
  }

  func finish() {
    subscribers.withValue { subscribers in
      subscribers.finished = true
      for continuation in subscribers.continuations.values { continuation.finish() }
      subscribers.continuations.removeAll()
    }
  }
}

/// `count` bytes of a file from `offset`, in 64 KiB pieces. Each pass re-opens the file, so a
/// retried PATCH replays the same bytes. An in-memory chunk is not an option: URLSession keeps a
/// copy of every `Data` request body after its task ends, so a long upload would hold the whole
/// file.
struct FileRange: AsyncSequence, Sendable {
  typealias Element = ArraySlice<UInt8>

  let fileURL: URL
  let offset: Int64
  let count: Int64

  struct AsyncIterator: AsyncIteratorProtocol {
    let fileURL: URL
    let offset: Int64
    var remaining: Int64
    var handle: FileHandle?

    mutating func next() async throws -> ArraySlice<UInt8>? {
      guard remaining > 0 else {
        try handle?.close()
        handle = nil
        return nil
      }
      let open = try handle ?? Self.open(fileURL, at: offset)
      handle = open
      guard let piece = try open.read(upToCount: Int(Swift.min(remaining, 64 * 1024))),
        !piece.isEmpty
      else {
        remaining = 0
        return nil
      }
      remaining -= Int64(piece.count)
      return ArraySlice(piece)
    }

    private static func open(_ fileURL: URL, at offset: Int64) throws -> FileHandle {
      let handle = try FileHandle(forReadingFrom: fileURL)
      try handle.seek(toOffset: UInt64(offset))
      return handle
    }
  }

  func makeAsyncIterator() -> AsyncIterator {
    AsyncIterator(fileURL: fileURL, offset: offset, remaining: count)
  }
}