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

import Foundation

/// Token returned by download() to allow cancellation of a specific subscriber.
struct DownloadToken: Equatable {
  fileprivate let id: UUID = UUID()
  fileprivate let cacheKey: String
}

/// An idempotent acknowledgement that cancellation can also resolve.
private final class StreamAcknowledgement {
  let id = UUID()
  private let lock = NSLock()
  private var resolved = false
  private let action: () -> Void

  init(action: @escaping () -> Void) { self.action = action }

  func resolve() {
    lock.lock()
    guard !resolved else {
      lock.unlock()
      return
    }
    resolved = true
    lock.unlock()
    action()
  }
}

/// Receives a non-cached streaming response with per-chunk backpressure.
struct DownloadStreamSubscriber {
  let onResponse: (_ contentType: String) -> Void
  let onData: (_ data: Data, _ acknowledge: @escaping () -> Void) -> Void
}

/// Coordinates downloads for audio URLs. Deduplicates concurrent requests
/// for the same cache key so that preload → play transitions share a single
/// upstream connection.
final class DownloadCoordinator: NSObject {

  private let cache: AudioCache
  private let sessionQueue = DispatchQueue(label: "trackplayer.downloads")

  // URLSession retains its delegate until the owner calls invalidate().
  private let sessionLock = NSLock()
  private var _session: URLSession?
  private var session: URLSession {
    sessionLock.lock()
    defer { sessionLock.unlock() }
    if let existing = _session { return existing }
    let config = URLSessionConfiguration.default
    let opQueue = OperationQueue()
    opQueue.underlyingQueue = sessionQueue
    opQueue.maxConcurrentOperationCount = 4
    let created = URLSession(configuration: config, delegate: self, delegateQueue: opQueue)
    _session = created
    return created
  }

  /// Cancel active tasks and release the session's delegate reference.
  func invalidate() {
    sessionLock.lock()
    let toInvalidate = _session
    _session = nil
    sessionLock.unlock()
    toInvalidate?.invalidateAndCancel()
  }

  /// State for a single in-progress download.
  private class ActiveDownload {
    let url: URL
    let cacheKey: String
    var task: URLSessionDataTask?
    var contentInfo: AudioCache.ContentInfo?
    var maxBytes: Int64?
    struct Subscriber {
      let token: DownloadToken
      let completion: (Result<Void, Error>) -> Void
      let stream: DownloadStreamSubscriber?
    }

    var subscribers: [Subscriber] = []
    var isICY = false
    var deinterleaver: ICYStreamDeinterleaver?
    var isLiveRelay = false
    /// Retained for subscribers joining an active relay.
    var liveContentType: String?
    var streamOnly = false
    var pendingAcknowledgements: [UUID: [StreamAcknowledgement]] = [:]

    init(url: URL, cacheKey: String) {
      self.url = url
      self.cacheKey = cacheKey
    }
  }

  /// Signals that the response cannot be safely handled through this download.
  struct ICYNotCacheable: Error {}
  struct UpstreamHTTPError: Error { let statusCode: Int }

  /// Map of cache key → active download.
  private var downloads: [String: ActiveDownload] = [:]
  /// Map of URLSessionTask identifier → cache key (for delegate routing).
  private var taskKeyMap: [Int: String] = [:]

  /// Called when any download completes successfully.
  /// Parameters: cache key, whether the full file was downloaded (false when stopped by maxBytes).
  var onDownloadComplete: ((_ key: String, _ isFull: Bool) -> Void)?

  /// Called on the session queue when an ICY metadata block is decoded.
  var onICYMetadata: ((_ cacheKey: String, _ metadata: StreamMetadata) -> Void)?

  /// Requests direct fallback for a response the proxy cannot safely frame.
  var onDirectFallbackRequired: ((_ cacheKey: String) -> Void)?

  var activeDownloadCount: Int {
    sessionQueue.sync { downloads.count }
  }

  init(cache: AudioCache) {
    self.cache = cache
  }

  // MARK: - Public

  /// Start or attach to a download for the given URL.
  /// Returns a token that can be used to cancel this subscriber.
  /// - Parameter cacheKey: Explicit cache key. If nil, computed from the URL.
  ///   Pass this when the URL has gone through encoding/decoding that might
  ///   change its string representation (e.g. proxy query parameters).
  @discardableResult
  func download(
    url: URL,
    headers: [String: String]?,
    cacheKey: String? = nil,
    maxBytes: Int64? = nil,
    streamOnly: Bool = false,
    streamSubscriber: DownloadStreamSubscriber? = nil,
    completion: @escaping (Result<Void, Error>) -> Void = { _ in }
  ) -> DownloadToken {
    let key = cacheKey ?? cache.cacheKey(for: url)
    let token = DownloadToken(cacheKey: key)

    sessionQueue.async { [self] in
      if let existing = downloads[key] {
        if streamOnly && !existing.streamOnly && existing.contentInfo != nil {
          // A classified disk download cannot be upgraded safely; restart live.
          existing.task?.cancel()
          if let taskID = existing.task?.taskIdentifier {
            taskKeyMap.removeValue(forKey: taskID)
          }
          downloads.removeValue(forKey: key)
          existing.pendingAcknowledgements.values.flatMap { $0 }.forEach { $0.resolve() }
          for subscriber in existing.subscribers {
            DispatchQueue.global().async {
              subscriber.completion(.failure(ICYNotCacheable()))
            }
          }
          cache.purge(key: key)
        } else {
          existing.streamOnly = existing.streamOnly || streamOnly
          existing.subscribers.append(.init(
            token: token,
            completion: completion,
            stream: streamSubscriber
          ))
          if existing.isLiveRelay,
             let streamSubscriber,
             let contentType = existing.liveContentType {
            streamSubscriber.onResponse(contentType)
          }
          return
        }
      }

      // Live playback must start at byte zero and never use stale disk state.
      if streamOnly {
        cache.purge(key: key)
      } else if cache.isFullyCached(key: key) {
        DispatchQueue.global().async { completion(.success(())) }
        return
      }

      // Clean cached ICY offsets do not map to the interleaved origin.
      if let info = cache.contentInfo(for: key), info.isICY == true {
        cache.purge(key: key)
      }

      let dl = ActiveDownload(url: url, cacheKey: key)
      dl.maxBytes = maxBytes
      dl.streamOnly = streamOnly
      dl.subscribers.append(.init(
        token: token,
        completion: completion,
        stream: streamSubscriber
      ))

      var request = URLRequest(url: url)
      headers?.forEach { request.setValue($0.value, forHTTPHeaderField: $0.key) }

      // Request ICY metadata unless the caller supplied its own preference.
      let callerSetIcy = (headers ?? [:]).keys.contains {
        $0.caseInsensitiveCompare("icy-metadata") == .orderedSame
      }
      if !callerSetIcy {
        request.setValue("1", forHTTPHeaderField: "Icy-MetaData")
      }

      let cachedBytes = cache.cachedBytes(for: key)
      if cachedBytes > 0 {
        request.setValue("bytes=\(cachedBytes)-", forHTTPHeaderField: "Range")
      }

      let task = session.dataTask(with: request)
      dl.task = task
      downloads[key] = dl
      taskKeyMap[task.taskIdentifier] = key

      task.resume()
    }

    return token
  }

  /// Cancel a specific subscriber. If no subscribers remain, the download is cancelled.
  func cancelDownload(token: DownloadToken) {
    sessionQueue.async { [self] in
      guard let dl = downloads[token.cacheKey] else { return }
      dl.subscribers.removeAll { $0.token == token }
      dl.pendingAcknowledgements.removeValue(forKey: token.id)?.forEach { $0.resolve() }
      let hasStreamSubscriber = dl.subscribers.contains { $0.stream != nil }
      if dl.subscribers.isEmpty || (dl.isLiveRelay && !hasStreamSubscriber) {
        dl.task?.cancel()
        if let taskId = dl.task?.taskIdentifier {
          taskKeyMap.removeValue(forKey: taskId)
        }
        downloads.removeValue(forKey: token.cacheKey)
        dl.pendingAcknowledgements.values.flatMap { $0 }.forEach { $0.resolve() }
        // Completion-only subscribers cannot consume an endless relay.
        for subscriber in dl.subscribers {
          DispatchQueue.global().async {
            subscriber.completion(.failure(ICYNotCacheable()))
          }
        }
        if dl.isICY { cache.purge(key: dl.cacheKey) }
      }
    }
  }

  /// Cancel all active downloads.
  func cancelAll() {
    sessionQueue.async { [self] in
      for (_, dl) in downloads {
        dl.task?.cancel()
        dl.pendingAcknowledgements.values.flatMap { $0 }.forEach { $0.resolve() }
        if dl.isICY { cache.purge(key: dl.cacheKey) }
      }
      downloads.removeAll()
      taskKeyMap.removeAll()
    }
  }

  // MARK: - Private

  private func activeDownload(for task: URLSessionTask) -> ActiveDownload? {
    guard let key = taskKeyMap[task.taskIdentifier] else { return nil }
    return downloads[key]
  }

  private func notifySubscribers(for key: String, result: Result<Void, Error>, isFull: Bool = true) {
    guard let dl = downloads.removeValue(forKey: key) else { return }
    dl.pendingAcknowledgements.values.flatMap { $0 }.forEach { $0.resolve() }
    if let taskId = dl.task?.taskIdentifier {
      taskKeyMap.removeValue(forKey: taskId)
    }
    for sub in dl.subscribers {
      DispatchQueue.global().async { sub.completion(result) }
    }
    if case .success = result {
      onDownloadComplete?(key, isFull)
    }
  }
}

// MARK: - URLSessionDataDelegate

extension DownloadCoordinator: URLSessionDataDelegate {

  func urlSession(
    _ session: URLSession,
    dataTask: URLSessionDataTask,
    didReceive response: URLResponse,
    completionHandler: @escaping (URLSession.ResponseDisposition) -> Void
  ) {
    guard let dl = activeDownload(for: dataTask),
          let http = response as? HTTPURLResponse else {
      completionHandler(.allow)
      return
    }

    guard (200...299).contains(http.statusCode) else {
      cache.purge(key: dl.cacheKey)
      completionHandler(.cancel)
      notifySubscribers(
        for: dl.cacheKey,
        result: .failure(UpstreamHTTPError(statusCode: http.statusCode))
      )
      return
    }

    // Extensionless HLS cannot be identified from the proxy URL path.
    let mimeType = (http.mimeType ?? "").lowercased()
    if dl.streamOnly && [
      "application/vnd.apple.mpegurl",
      "application/x-mpegurl",
      "audio/mpegurl",
      "audio/x-mpegurl"
    ].contains(mimeType) {
      completionHandler(.cancel)
      notifySubscribers(for: dl.cacheKey, result: .failure(ICYNotCacheable()))
      onDirectFallbackRequired?(dl.cacheKey)
      return
    }

    let streamSubscribers = dl.subscribers.compactMap(\.stream)

    // Finite ICY is cached clean; endless or explicit live ICY is relayed.
    if let metaintHeader = http.value(forHTTPHeaderField: "icy-metaint") {
      let interval = Int(metaintHeader) ?? 0
      guard interval > 0 else {
        // Raw passthrough would leak metadata blocks into the audio.
        cache.purge(key: dl.cacheKey)
        completionHandler(.cancel)
        notifySubscribers(for: dl.cacheKey, result: .failure(ICYNotCacheable()))
        onDirectFallbackRequired?(dl.cacheKey)
        return
      }

      // A ranged response may begin in the middle of an ICY frame.
      if http.statusCode == 206 {
        cache.purge(key: dl.cacheKey)
        completionHandler(.cancel)
        notifySubscribers(for: dl.cacheKey, result: .failure(ICYNotCacheable()))
        onDirectFallbackRequired?(dl.cacheKey)
        return
      }

      if dl.streamOnly || response.expectedContentLength <= 0 {
        guard !streamSubscribers.isEmpty else {
          // Do not leave a no-consumer stream running forever.
          cache.purge(key: dl.cacheKey)
          completionHandler(.cancel)
          notifySubscribers(for: dl.cacheKey, result: .failure(ICYNotCacheable()))
          return
        }

        cache.purge(key: dl.cacheKey)
        dl.isICY = true
        dl.isLiveRelay = true
        dl.deinterleaver = ICYStreamDeinterleaver(interval: interval)
        let contentType = http.mimeType ?? "audio/mpeg"
        dl.liveContentType = contentType
        for subscriber in streamSubscribers {
          subscriber.onResponse(contentType)
        }
        completionHandler(.allow)
        return
      }

      if response.expectedContentLength > 0 {
        // The origin ignored a speculative Range; discard the stale prefix.
        if cache.cachedBytes(for: dl.cacheKey) > 0 {
          cache.purge(key: dl.cacheKey)
        }
        dl.isICY = true
        dl.deinterleaver = ICYStreamDeinterleaver(interval: interval)
        // Clean length is unknown until completion; finalize will set it.
        let info = AudioCache.ContentInfo(
          contentType: http.mimeType ?? "audio/mpeg",
          contentLength: -1,
          isByteRangeAccessSupported: false,
          isICY: true
        )
        cache.storeContentInfo(info, for: dl.cacheKey, url: dl.url)
        dl.contentInfo = info
        completionHandler(.allow)
        return
      }
    }

    // Explicitly-live non-ICY responses also bypass disk.
    if dl.streamOnly {
      guard !streamSubscribers.isEmpty else {
        completionHandler(.cancel)
        notifySubscribers(for: dl.cacheKey, result: .failure(ICYNotCacheable()))
        return
      }
      cache.purge(key: dl.cacheKey)
      dl.isLiveRelay = true
      let contentType = http.mimeType ?? "audio/mpeg"
      dl.liveContentType = contentType
      for subscriber in streamSubscribers {
        subscriber.onResponse(contentType)
      }
      completionHandler(.allow)
      return
    }

    if cache.contentInfo(for: dl.cacheKey) == nil {
      let mime = http.mimeType ?? "audio/mpeg"
      var totalLength: Int64 = -1

      if http.statusCode == 206,
         let rangeHeader = http.value(forHTTPHeaderField: "Content-Range"),
         let slashIdx = rangeHeader.lastIndex(of: "/") {
        totalLength = Int64(rangeHeader[rangeHeader.index(after: slashIdx)...]) ?? -1
      } else {
        totalLength = response.expectedContentLength
      }

      let byteRange = http.statusCode == 206
      let info = AudioCache.ContentInfo(
        contentType: mime,
        contentLength: totalLength,
        isByteRangeAccessSupported: byteRange
      )
      cache.storeContentInfo(info, for: dl.cacheKey, url: dl.url)
      downloads[dl.cacheKey]?.contentInfo = info
    }

    completionHandler(.allow)
  }

  func urlSession(_ session: URLSession, dataTask: URLSessionDataTask, didReceive data: Data) {
    guard let dl = activeDownload(for: dataTask) else { return }

    var output = data
    if dl.isICY, var deinterleaver = dl.deinterleaver {
      let (audio, metadata) = deinterleaver.consume(data)
      dl.deinterleaver = deinterleaver
      output = audio
      for block in metadata {
        if let parsed = ICYMetadataParser.streamMetadata(from: block.value) {
          if !dl.isLiveRelay {
            cache.appendICYMetadata(
              parsed,
              at: block.audioByteOffset,
              for: dl.cacheKey
            )
          }
          onICYMetadata?(dl.cacheKey, parsed)
        }
      }
    }

    if dl.isLiveRelay {
      let subscribers = dl.subscribers.compactMap { subscriber in
        subscriber.stream.map { (subscriber.token, $0) }
      }
      guard !output.isEmpty, !subscribers.isEmpty else { return }

      // Bound buffering to downstream transport throughput.
      dataTask.suspend()
      let acknowledgements = DispatchGroup()
      for (token, stream) in subscribers {
        acknowledgements.enter()
        let acknowledgement = StreamAcknowledgement {
          acknowledgements.leave()
        }
        dl.pendingAcknowledgements[token.id, default: []].append(acknowledgement)
        stream.onData(output) { [weak self, weak dl] in
          acknowledgement.resolve()
          self?.sessionQueue.async {
            guard let dl else { return }
            dl.pendingAcknowledgements[token.id]?.removeAll {
              $0.id == acknowledgement.id
            }
          }
        }
      }
      acknowledgements.notify(queue: sessionQueue) { [weak self, weak dataTask] in
        guard let self, let dataTask,
              self.activeDownload(for: dataTask) != nil else { return }
        dataTask.resume()
      }
      return
    }

    if !output.isEmpty {
      cache.appendData(output, for: dl.cacheKey)
    }

    // Stop download if byte limit reached.
    if let maxBytes = dl.maxBytes, cache.cachedBytes(for: dl.cacheKey) >= maxBytes {
      cache.touchAccessDate(for: dl.cacheKey)
      cache.evictIfNeeded()
      dataTask.cancel()
      notifySubscribers(for: dl.cacheKey, result: .success(()), isFull: false)
    }
  }

  func urlSession(_ session: URLSession, task: URLSessionTask, didCompleteWithError error: Error?) {
    guard let key = taskKeyMap[task.taskIdentifier] else { return }

    if let error = error as? URLError, error.code == .cancelled { return }

    if let error = error {
      notifySubscribers(for: key, result: .failure(error))
      return
    }

    // Finalize content length if the server didn't provide it.
    if let info = cache.contentInfo(for: key), info.contentLength <= 0 {
      let actualLength = cache.cachedBytes(for: key)
      if actualLength > 0 {
        let updated = AudioCache.ContentInfo(
          contentType: info.contentType,
          contentLength: actualLength,
          isByteRangeAccessSupported: info.isByteRangeAccessSupported,
          isICY: info.isICY
        )
        cache.storeContentInfo(updated, for: key, url: downloads[key]?.url ?? URL(fileURLWithPath: "/"))
      }
    }

    cache.touchAccessDate(for: key)
    cache.evictIfNeeded()

    notifySubscribers(for: key, result: .success(()))
  }
}