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

/// 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")
  private lazy var session: URLSession = {
    let config = URLSessionConfiguration.default
    let opQueue = OperationQueue()
    opQueue.underlyingQueue = sessionQueue
    opQueue.maxConcurrentOperationCount = 4
    return URLSession(configuration: config, delegate: self, delegateQueue: opQueue)
  }()

  /// 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?
    var subscribers: [(token: DownloadToken, completion: (Result<Void, Error>) -> Void)] = []

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

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

  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,
    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 cache.isFullyCached(key: key) {
        DispatchQueue.global().async { completion(.success(())) }
        return
      }

      if let existing = downloads[key] {
        existing.subscribers.append((token, completion))
        return
      }

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

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

      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 }
      if dl.subscribers.isEmpty {
        dl.task?.cancel()
        if let taskId = dl.task?.taskIdentifier {
          taskKeyMap.removeValue(forKey: taskId)
        }
        downloads.removeValue(forKey: token.cacheKey)
      }
    }
  }

  /// Cancel all active downloads.
  func cancelAll() {
    sessionQueue.async { [self] in
      for (_, dl) in downloads {
        dl.task?.cancel()
      }
      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 }
    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
    }

    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 }
    cache.appendData(data, 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
        )
        cache.storeContentInfo(updated, for: key, url: downloads[key]?.url ?? URL(fileURLWithPath: "/"))
      }
    }

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

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