//
// Copyright (c) Double Symmetry GmbH
// Commercial use requires a license. See https://rntp.dev/pricing
//

import Foundation

/// Warms the cache ahead of playback. For plain media a single byte-bounded
/// blob is fetched; for HLS the manifest is resolved and the actual media
/// SEGMENTS (plus the init map and decryption key) are downloaded so the proxy
/// can serve a fully offline playlist — matching Android's segment-level
/// preload rather than caching only the manifest text.
final class Preloader {
  private let coordinator: DownloadCoordinator
  private let cache: AudioCache

  /// One logical preload spawns many downloads (manifest, an optional variant
  /// manifest, N segments, init, key) created across async completion
  /// callbacks. All of a job's tokens are tracked under the ROOT preloaded
  /// url's cache key so `cancel(url:)` cancels the whole job. Coordinator
  /// completions fire on a background queue and append tokens from there, so
  /// every access to this store is serialized on `stateQueue`.
  ///
  /// Cancellation is signalled purely by removing the job from `jobs`: both
  /// `cancel(url:)` and `cancelAll()` `removeValue` the job, and every
  /// follow-up callback re-looks-up `jobs[rootKey]` on `stateQueue` and bails
  /// when it is gone. There is no separate `cancelled` flag because, on the
  /// serial queue, dict-absence already happens-before any later callback.
  private final class PreloadJob {
    var tokens: [DownloadToken] = []
    /// Downloads dispatched but not yet completed. The job entry is removed
    /// once this reaches zero, after each completion's follow-up work has had a
    /// chance to schedule more downloads.
    var pending = 0
    /// Manifest URLs already resolved by this job. Seeds the master→media
    /// recursion's loop guard so a cyclic or self-referential master chain
    /// (A → B → A, or A → A) can never recurse forever.
    var visitedManifests: Set<URL> = []
  }

  private var jobs: [String: PreloadJob] = [:]
  private let stateQueue = DispatchQueue(label: "trackplayer.preloader")

  init(cache: AudioCache, coordinator: DownloadCoordinator) {
    self.cache = cache
    self.coordinator = coordinator
  }

  /// Preload a URL. HLS playlists (`.m3u8`/`.m3u`) warm their segments; any
  /// other URL is downloaded as a single blob bounded by `maxBytes` (downloads
  /// the entire file when `maxBytes` is nil). `duration` only bounds HLS: it is
  /// ignored on the non-HLS path.
  func preload(url: URL, headers: [String: String]?, maxBytes: Int64? = nil, duration: Double = -1) {
    let rootKey = cache.cacheKey(for: url)

    stateQueue.async { [weak self] in
      guard let self = self else { return }
      // Re-entrancy/dedup: a job for this root is already active.
      guard self.jobs[rootKey] == nil else { return }

      if self.isHLS(url) {
        self.jobs[rootKey] = PreloadJob()
        self.preloadHLS(manifestURL: url, headers: headers, duration: duration, rootKey: rootKey)
      } else {
        // Non-HLS: single-blob download.
        if self.cache.isFullyCached(key: rootKey) { return }
        self.jobs[rootKey] = PreloadJob()
        self.startDownload(url: url, headers: headers, maxBytes: maxBytes, rootKey: rootKey)
      }
    }
  }

  /// Cancel preload for a specific URL. Removing the job from `jobs` is what
  /// stops any in-flight callback from scheduling further downloads; cancelling
  /// the recorded tokens tears down the downloads already dispatched.
  func cancel(url: URL) {
    let rootKey = cache.cacheKey(for: url)
    stateQueue.async { [weak self] in
      guard let self = self, let job = self.jobs.removeValue(forKey: rootKey) else { return }
      for token in job.tokens {
        self.coordinator.cancelDownload(token: token)
      }
    }
  }

  /// Cancel all active preloads.
  func cancelAll() {
    stateQueue.async { [weak self] in
      guard let self = self else { return }
      for job in self.jobs.values {
        for token in job.tokens {
          self.coordinator.cancelDownload(token: token)
        }
      }
      self.jobs.removeAll()
    }
  }

  // MARK: - HLS

  /// Whether a URL points at an HLS playlist, detected by its path extension
  /// (`m3u8` / `m3u`, case-insensitive).
  private func isHLS(_ url: URL) -> Bool {
    let ext = url.pathExtension.lowercased()
    return ext == "m3u8" || ext == "m3u"
  }

  /// Resolve a playlist and warm its dependencies. Must be called on
  /// `stateQueue`; the job entry for `rootKey` is assumed to already exist.
  ///
  /// The manifest is downloaded under its own cache key so the RAW bytes are
  /// cached (the proxy reads them back to serve a rewritten playlist). On
  /// success the bytes are decoded and parsed: a master playlist recurses into
  /// its lowest-bandwidth variant; a media playlist warms the chosen segment
  /// prefix plus its auxiliary init/key resources.
  private func preloadHLS(manifestURL: URL, headers: [String: String]?, duration: Double, rootKey: String) {
    let manifestKey = cache.cacheKey(for: manifestURL)

    startDownload(url: manifestURL, headers: headers, cacheKey: manifestKey, rootKey: rootKey) { [weak self] result in
      // Runs on `stateQueue`, inside the manifest's completion critical section
      // (before the pending-zero cleanup), so any segment downloads scheduled
      // here keep the job alive. A nil `jobs[rootKey]` lookup means the job was
      // cancelled; we bail before scheduling anything further.
      guard let self = self else { return }
      guard case .success = result else { return }
      guard let job = self.jobs[rootKey] else { return }
      job.visitedManifests.insert(manifestURL)

      // `readData` runs here on `stateQueue` while other segments of this (or
      // another) job may be appending on the coordinator's session queue; this
      // is safe because `AudioCache` serializes every access on its own queue.
      guard let data = self.cache.readData(for: manifestKey, offset: 0, length: Int.max),
            let text = String(data: data, encoding: .utf8) else { return }

      if HLSPlaylistParser.isMaster(text) {
        self.handleMaster(text: text, manifestURL: manifestURL, headers: headers, duration: duration, rootKey: rootKey)
      } else {
        self.handleMedia(text: text, manifestURL: manifestURL, headers: headers, duration: duration, rootKey: rootKey)
      }
    }
  }

  /// Pick the lowest-bandwidth variant and recurse one level (master → media).
  /// Real HLS is master → media (one hop); the visited-manifest set guards a
  /// pathological master pointing back at an already-resolved manifest (`A → A`
  /// or a longer `A → B → A` cycle) so the recursion can never loop forever.
  /// Must be called on `stateQueue` with a live job.
  private func handleMaster(
    text: String,
    manifestURL: URL,
    headers: [String: String]?,
    duration: Double,
    rootKey: String
  ) {
    guard let job = jobs[rootKey] else { return }

    let variants = HLSPlaylistParser.variants(manifest: text, baseURL: manifestURL)
    guard !variants.isEmpty else { return }

    // Lowest bandwidth; ties (and bandwidth 0) fall back to declared order.
    let chosen = variants.min { $0.bandwidth < $1.bandwidth } ?? variants[0]
    guard !job.visitedManifests.contains(chosen.url) else { return }

    preloadHLS(manifestURL: chosen.url, headers: headers, duration: duration, rootKey: rootKey)
  }

  /// Warm the chosen segment prefix plus auxiliary init/key resources of a
  /// media playlist. `duration <= 0` warms every segment; `duration > 0` warms
  /// the shortest prefix whose cumulative EXTINF covers it (at least one).
  /// Must be called on `stateQueue` with a live job.
  private func handleMedia(
    text: String,
    manifestURL: URL,
    headers: [String: String]?,
    duration: Double,
    rootKey: String
  ) {
    let segments = HLSPlaylistParser.segments(manifest: text, baseURL: manifestURL)
    let chosen = segmentPrefix(segments, coveringDuration: duration)
    let aux = HLSPlaylistParser.auxiliaryURIs(manifest: text, baseURL: manifestURL)

    for url in aux + chosen.map({ $0.url }) {
      let key = cache.cacheKey(for: url)
      if cache.isFullyCached(key: key) { continue }
      startDownload(url: url, headers: headers, cacheKey: key, rootKey: rootKey)
    }
  }

  /// The shortest leading run of segments whose cumulative EXTINF reaches
  /// `duration`. `duration <= 0` returns all segments; otherwise at least one
  /// segment is returned whenever any exist.
  private func segmentPrefix(_ segments: [HLSPlaylistParser.Segment], coveringDuration duration: Double) -> [HLSPlaylistParser.Segment] {
    guard duration > 0 else { return segments }
    guard !segments.isEmpty else { return [] }

    var cumulative = 0.0
    var count = 0
    for segment in segments {
      cumulative += segment.duration
      count += 1
      if cumulative >= duration { break }
    }
    return Array(segments.prefix(max(count, 1)))
  }

  // MARK: - Download bookkeeping

  /// Dispatch a single download and record its token under `rootKey`. Must be
  /// called on `stateQueue`. The coordinator completion (fired on a background
  /// queue) re-enters `stateQueue` where `onComplete` runs and may schedule
  /// further downloads; only then is the job's pending count settled and the
  /// job entry removed once every download has finished.
  private func startDownload(
    url: URL,
    headers: [String: String]?,
    maxBytes: Int64? = nil,
    cacheKey: String? = nil,
    rootKey: String,
    onComplete: @escaping (Result<Void, Error>) -> Void = { _ in }
  ) {
    guard let job = jobs[rootKey] else { return }
    job.pending += 1

    let token = coordinator.download(url: url, headers: headers, cacheKey: cacheKey, maxBytes: maxBytes) { [weak self] result in
      guard let self = self else { return }
      self.stateQueue.async {
        // Schedule any follow-up downloads first so they keep the job alive,
        // then settle this download and clean up once nothing is pending.
        onComplete(result)
        guard let job = self.jobs[rootKey] else { return }
        job.pending -= 1
        if job.pending <= 0 {
          self.jobs.removeValue(forKey: rootKey)
        }
      }
    }

    job.tokens.append(token)
  }
}