ABS-Client/ABS Client/Audiobookshelf swift/Services/DownloadManager.swift

import Foundation
import Observation

struct DownloadedTrack: Codable, Hashable {
    let ino: String
    let filename: String
    let localPath: String  // relative to AppPaths.downloadsDirectory
    let durationSeconds: Double

    enum CodingKeys: String, CodingKey {
        case ino, filename, localPath, durationSeconds
    }

    init(ino: String, filename: String, localPath: String, durationSeconds: Double) {
        self.ino = ino
        self.filename = filename
        self.localPath = localPath
        self.durationSeconds = durationSeconds
    }

    init(from decoder: Decoder) throws {
        let c = try decoder.container(keyedBy: CodingKeys.self)
        ino = try c.decode(String.self, forKey: .ino)
        filename = try c.decode(String.self, forKey: .filename)
        localPath = try c.decode(String.self, forKey: .localPath)
        durationSeconds = try c.decodeIfPresent(Double.self, forKey: .durationSeconds) ?? 0
    }
}

struct DownloadedItem: Codable, Hashable {
    let itemId: String
    var episodeId: String?
    let title: String
    let author: String
    let durationSeconds: Double
    let tracks: [DownloadedTrack]

    enum CodingKeys: String, CodingKey {
        case itemId, episodeId, title, author, durationSeconds, tracks
    }

    init(itemId: String, episodeId: String? = nil, title: String, author: String, durationSeconds: Double, tracks: [DownloadedTrack]) {
        self.itemId = itemId
        self.episodeId = episodeId
        self.title = title
        self.author = author
        self.durationSeconds = durationSeconds
        self.tracks = tracks
    }

    init(from decoder: Decoder) throws {
        let c = try decoder.container(keyedBy: CodingKeys.self)
        itemId = try c.decode(String.self, forKey: .itemId)
        episodeId = try c.decodeIfPresent(String.self, forKey: .episodeId)
        title = try c.decode(String.self, forKey: .title)
        author = try c.decode(String.self, forKey: .author)
        durationSeconds = try c.decode(Double.self, forKey: .durationSeconds)
        tracks = try c.decode([DownloadedTrack].self, forKey: .tracks)
    }

    var downloadKey: String {
        if let episodeId { return "\(itemId)|\(episodeId)" }
        return itemId
    }
}

@Observable
@MainActor
final class DownloadManager: @unchecked Sendable {
    private let client: ABSClient
    /// Keyed by downloadKey (itemId or "itemId|episodeId").
    private(set) var states: [String: DownloadState] = [:]
    private(set) var downloadedItems: [String: DownloadedItem] = [:]
    /// Items currently being downloaded (removed on completion or cancellation).
    private(set) var pendingItems: [String: LibraryItem] = [:]
    /// Bytes received for the currently in-flight track per downloadKey.
    /// Reset between tracks; cleared when the download finishes/cancels.
    private(set) var inFlightBytes: [String: Int64] = [:]
    /// Current track index (0-based) per active download. Set at the start of
    /// each track iteration, cleared via defer when the download exits.
    private var currentTrackIndex: [String: Int] = [:]
    /// Total number of tracks per active download.
    private var totalTrackCount: [String: Int] = [:]

    private var indexFile: URL { AppPaths.supportDirectory.appendingPathComponent("downloads-index.json") }
    private var activeTasks: [String: Task<Void, Never>] = [:]

    init(client: ABSClient) {
        self.client = client
        try? FileManager.default.createDirectory(at: AppPaths.downloadsDirectory, withIntermediateDirectories: true)
        loadIndex()
    }

    func state(for downloadKey: String) -> DownloadState {
        states[downloadKey] ?? .notDownloaded
    }

    func isDownloaded(downloadKey: String) -> Bool {
        if case .downloaded = state(for: downloadKey) { return true }
        return false
    }

    func downloadedBytes(for downloadKey: String) -> Int64 {
        if let item = downloadedItems[downloadKey] {
            return item.tracks.reduce(Int64(0)) { sum, track in
                let url = AppPaths.downloadsDirectory.appendingPathComponent(track.localPath)
                return sum + fileSize(at: url)
            }
        }
        if let pending = pendingItems[downloadKey] {
            let onDisk = folderSize(at: AppPaths.downloadsDirectory.appendingPathComponent(pending.id))
            let inFlight = inFlightBytes[downloadKey] ?? 0
            return onDisk + inFlight
        }
        return 0
    }

    /// Called from the URL session delegate (any queue) to update in-flight bytes
    /// for a currently downloading track.
    nonisolated func _updateInFlightBytes(_ bytes: Int64, for key: String) {
        Task { @MainActor [self] in
            self.inFlightBytes[key] = bytes
        }
    }

    /// Called from the URL session delegate with the fraction (0…1) of the
    /// currently downloading track. Combines with the completed-track count to
    /// drive a smooth overall progress ring, even for single-track downloads.
    nonisolated func _reportTrackByteFraction(_ fraction: Double, for downloadKey: String) {
        Task { @MainActor [self] in
            guard let idx = self.currentTrackIndex[downloadKey],
                  let total = self.totalTrackCount[downloadKey],
                  total > 0 else { return }
            let overall = (Double(idx) + max(0, min(1, fraction))) / Double(total)
            // Clamp below 1.0 so `performDownload` is the only place that
            // transitions to `.downloaded` after the final track is persisted.
            self.states[downloadKey] = .downloading(progress: min(0.999, overall))
        }
    }

    private func fileSize(at url: URL) -> Int64 {
        guard let attrs = try? FileManager.default.attributesOfItem(atPath: url.path) else { return 0 }
        if let n = attrs[.size] as? NSNumber { return n.int64Value }
        if let i = attrs[.size] as? Int { return Int64(i) }
        return 0
    }

    func localTrackURLs(for downloadKey: String) -> [URL]? {
        guard let item = downloadedItems[downloadKey] else { return nil }
        return item.tracks.map { AppPaths.downloadsDirectory.appendingPathComponent($0.localPath) }
    }

    /// Downloads a book (whole audioFiles list) or a podcast episode (single audioFile).
    func startDownload(item: LibraryItem) {
        let key = item.downloadKey
        guard activeTasks[key] == nil else { return }
        states[key] = .downloading(progress: 0)
        pendingItems[key] = item

        let task = Task { @MainActor [weak self] in
            guard let self else { return }
            var workItem = item

            if !workItem.isPodcast && workItem.audioFiles.isEmpty {
                do {
                    workItem = try await self.client.fetchItemDetail(itemId: item.id)
                } catch {
                    self.states[key] = .failed(message: "Detail konnte nicht geladen werden: \(error.localizedDescription)")
                    self.pendingItems.removeValue(forKey: key)
                    self.activeTasks[key] = nil
                    return
                }
            }
            if workItem.audioFiles.isEmpty {
                self.states[key] = .failed(message: "Keine herunterladbaren Audiodateien gefunden.")
                self.pendingItems.removeValue(forKey: key)
                self.activeTasks[key] = nil
                return
            }
            await self.performDownload(workItem: workItem, downloadKey: key)
            self.activeTasks[key] = nil
        }
        activeTasks[key] = task
    }

    func cancel(downloadKey: String) {
        activeTasks[downloadKey]?.cancel()
        activeTasks[downloadKey] = nil
        states[downloadKey] = .notDownloaded
        pendingItems.removeValue(forKey: downloadKey)
        inFlightBytes.removeValue(forKey: downloadKey)
    }

    func delete(downloadKey: String) {
        pendingItems.removeValue(forKey: downloadKey)
        cancel(downloadKey: downloadKey)
        if let item = downloadedItems[downloadKey] {
            let dir = directoryURL(itemId: item.itemId, episodeId: item.episodeId)
            try? FileManager.default.removeItem(at: dir)
            if item.episodeId != nil {
                let parent = AppPaths.downloadsDirectory.appendingPathComponent(item.itemId)
                if let contents = try? FileManager.default.contentsOfDirectory(atPath: parent.path), contents.isEmpty {
                    try? FileManager.default.removeItem(at: parent)
                }
            }
        }
        downloadedItems.removeValue(forKey: downloadKey)
        states[downloadKey] = .notDownloaded
        persistIndex()
    }

    private func folderSize(at url: URL) -> Int64 {
        guard let enumerator = FileManager.default.enumerator(
            atPath: url.path
        ) else { return 0 }
        var total: Int64 = 0
        for case let relPath as String in enumerator {
            let fileURL = url.appendingPathComponent(relPath)
            total += fileSize(at: fileURL)
        }
        return total
    }

    private func directoryURL(itemId: String, episodeId: String?) -> URL {
        var dir = AppPaths.downloadsDirectory.appendingPathComponent(itemId, isDirectory: true)
        if let episodeId {
            dir = dir.appendingPathComponent(episodeId, isDirectory: true)
        }
        return dir
    }

    private func relativePath(itemId: String, episodeId: String?, fileName: String) -> String {
        if let episodeId { return "\(itemId)/\(episodeId)/\(fileName)" }
        return "\(itemId)/\(fileName)"
    }

    private func performDownload(workItem: LibraryItem, downloadKey: String) async {
        defer {
            pendingItems.removeValue(forKey: downloadKey)
            inFlightBytes.removeValue(forKey: downloadKey)
            currentTrackIndex.removeValue(forKey: downloadKey)
            totalTrackCount.removeValue(forKey: downloadKey)
        }
        let itemDir = directoryURL(itemId: workItem.id, episodeId: workItem.episodeId)
        do {
            try FileManager.default.createDirectory(at: itemDir, withIntermediateDirectories: true)
        } catch {
            states[downloadKey] = .failed(message: error.localizedDescription)
            return
        }

        var tracks: [DownloadedTrack] = []
        let total = max(workItem.audioFiles.count, 1)
        totalTrackCount[downloadKey] = total

        for (idx, file) in workItem.audioFiles.enumerated() {
            if Task.isCancelled {
                states[downloadKey] = .notDownloaded
                return
            }
            guard let url = client.audioFileURL(itemId: workItem.id, ino: file.ino) else { continue }
            var request = URLRequest(url: url)
            for (k, v) in client.bearerHeader { request.setValue(v, forHTTPHeaderField: k) }
            currentTrackIndex[downloadKey] = idx

            let tempURL: URL
            do {
                tempURL = try await downloadWithRetry(request: request, filename: file.filename, downloadKey: downloadKey)
            } catch is CancellationError {
                states[downloadKey] = .notDownloaded
                return
            } catch {
                states[downloadKey] = .failed(message: error.localizedDescription)
                return
            }

            let ext = file.ext.isEmpty ? "mp3" : file.ext
            let destName = "\(String(format: "%03d", idx))-\(file.ino).\(ext)"
            let dest = itemDir.appendingPathComponent(destName)
            do {
                try? FileManager.default.removeItem(at: dest)
                try FileManager.default.moveItem(at: tempURL, to: dest)
            } catch {
                states[downloadKey] = .failed(message: error.localizedDescription)
                return
            }
            tracks.append(DownloadedTrack(
                ino: file.ino,
                filename: file.filename,
                localPath: relativePath(itemId: workItem.id, episodeId: workItem.episodeId, fileName: destName),
                durationSeconds: file.durationSeconds
            ))
            states[downloadKey] = .downloading(progress: Double(idx + 1) / Double(total))
            // Track is on disk now (folderSize picks it up). Clear the in-flight
            // counter so the next track's bytes don't double-count.
            inFlightBytes[downloadKey] = 0
        }

        let downloaded = DownloadedItem(
            itemId: workItem.id,
            episodeId: workItem.episodeId,
            title: workItem.title,
            author: workItem.author,
            durationSeconds: workItem.durationSeconds,
            tracks: tracks
        )
        downloadedItems[downloadKey] = downloaded
        states[downloadKey] = .downloaded
        persistIndex()
    }

    /// Downloads with up to `maxAttempts` retries and resume-data support so a brief
    /// network dropout picks up where it left off. Uses a classic URLSessionDownloadTask
    /// with explicit delegate (wrapped in a continuation) so we reliably get
    /// `didWriteData` progress callbacks — the async `download(for:delegate:)` API
    /// often doesn't fire them.
    private func downloadWithRetry(request: URLRequest, filename: String, downloadKey: String, maxAttempts: Int = 5) async throws -> URL {
        var resumeData: Data? = nil
        var lastError: Error = URLError(.unknown)

        for attempt in 0..<maxAttempts {
            try Task.checkCancellation()
            do {
                let (tempURL, response) = try await streamingDownload(
                    request: request,
                    resumeData: resumeData,
                    downloadKey: downloadKey
                )
                if let http = response as? HTTPURLResponse, !(200..<300).contains(http.statusCode) {
                    try? FileManager.default.removeItem(at: tempURL)
                    throw URLError(.badServerResponse)
                }
                return tempURL
            } catch is CancellationError {
                throw CancellationError()
            } catch let error as NSError {
                resumeData = error.userInfo[NSURLSessionDownloadTaskResumeData] as? Data
                lastError = error
                if attempt < maxAttempts - 1 {
                    // Exponential backoff: 1 s, 2 s, 4 s, 8 s …
                    let delay = UInt64(min(pow(2.0, Double(attempt)), 30)) * 1_000_000_000
                    try await Task.sleep(nanoseconds: delay)
                }
            }
        }
        throw lastError
    }

    /// Runs a single download attempt via URLSessionDownloadTask, reporting byte
    /// progress to `inFlightBytes` and moving the system-temp file to a stable
    /// location before returning.
    private func streamingDownload(
        request: URLRequest,
        resumeData: Data?,
        downloadKey: String
    ) async throws -> (URL, URLResponse) {
        let session = client.session
        return try await withCheckedThrowingContinuation { continuation in
            let delegate = DownloadProgressDelegate(
                manager: self,
                downloadKey: downloadKey,
                continuation: continuation
            )
            let task: URLSessionDownloadTask
            if let resumeData {
                task = session.downloadTask(withResumeData: resumeData)
            } else {
                task = session.downloadTask(with: request)
            }
            task.delegate = delegate
            task.resume()
        }
    }

    private func loadIndex() {
        guard let data = try? Data(contentsOf: indexFile),
              let decoded = try? JSONDecoder().decode([String: DownloadedItem].self, from: data) else { return }
        var rekeyed: [String: DownloadedItem] = [:]
        for (_, item) in decoded {
            if item.tracks.isEmpty { continue }
            rekeyed[item.downloadKey] = item
        }
        downloadedItems = rekeyed
        for k in rekeyed.keys {
            states[k] = .downloaded
        }
        for (oldKey, item) in decoded where item.tracks.isEmpty {
            let dir = AppPaths.downloadsDirectory.appendingPathComponent(oldKey)
            try? FileManager.default.removeItem(at: dir)
        }
        if rekeyed.count != decoded.count {
            persistIndex()
        }
    }

    private func persistIndex() {
        do {
            let data = try JSONEncoder().encode(downloadedItems)
            try data.write(to: indexFile, options: .atomic)
        } catch {
            // non-fatal
        }
    }
}

/// Per-task delegate for a single `URLSessionDownloadTask`. Forwards live byte
/// progress to the manager and bridges the delegate callbacks back to async/await
/// via a checked continuation. The system deletes the `didFinishDownloadingTo`
/// temp URL immediately after the callback returns, so we move it to a stable
/// location before resuming.
private final class DownloadProgressDelegate: NSObject, URLSessionDownloadDelegate, @unchecked Sendable {
    private let manager: DownloadManager
    private let downloadKey: String
    private let continuation: CheckedContinuation<(URL, URLResponse), Error>
    private var stableTempURL: URL?
    private var didResume = false
    private let lock = NSLock()

    init(
        manager: DownloadManager,
        downloadKey: String,
        continuation: CheckedContinuation<(URL, URLResponse), Error>
    ) {
        self.manager = manager
        self.downloadKey = downloadKey
        self.continuation = continuation
        super.init()
    }

    private func resumeOnce(with result: Result<(URL, URLResponse), Error>) {
        lock.lock(); defer { lock.unlock() }
        guard !didResume else { return }
        didResume = true
        continuation.resume(with: result)
    }

    func urlSession(
        _ session: URLSession,
        downloadTask: URLSessionDownloadTask,
        didWriteData bytesWritten: Int64,
        totalBytesWritten: Int64,
        totalBytesExpectedToWrite: Int64
    ) {
        manager._updateInFlightBytes(totalBytesWritten, for: downloadKey)
        if totalBytesExpectedToWrite > 0 {
            let fraction = Double(totalBytesWritten) / Double(totalBytesExpectedToWrite)
            manager._reportTrackByteFraction(fraction, for: downloadKey)
        }
    }

    func urlSession(
        _ session: URLSession,
        downloadTask: URLSessionDownloadTask,
        didFinishDownloadingTo location: URL
    ) {
        // Move out of the system-temp folder before this delegate method returns
        // (otherwise the OS deletes it).
        let target = FileManager.default.temporaryDirectory
            .appendingPathComponent(UUID().uuidString + ".tmp")
        do {
            try FileManager.default.moveItem(at: location, to: target)
            stableTempURL = target
        } catch {
            stableTempURL = nil
        }
    }

    func urlSession(
        _ session: URLSession,
        task: URLSessionTask,
        didCompleteWithError error: Error?
    ) {
        if let error {
            resumeOnce(with: .failure(error))
            return
        }
        guard let url = stableTempURL, let response = task.response else {
            resumeOnce(with: .failure(URLError(.cannotCreateFile)))
            return
        }
        resumeOnce(with: .success((url, response)))
    }
}