kangkang/康康/Services/ModelDownloadService.swift

import Foundation
import Observation

/// 模型下载编排:遍历 ModelManifest 逐文件串行下载,聚合进度,支持暂停/重试/旁路导入。
/// UI 只观察 `states`,不直接碰 URLSession(§3.1 模块边界)。
/// 核心下载/校验逻辑在 `FileDownloader`,文件路径/就绪判定在 `ModelStore`。
@MainActor
@Observable
final class ModelDownloadService {
    static let shared = ModelDownloadService()

    private(set) var states: [ModelKind: DownloadState] = [:]

    private let store: ModelStore
    private var tasks: [ModelKind: Task<Void, Never>] = [:]
    private var lastSampleTime: [ModelKind: Date] = [:]
    private var lastSampleBytes: [ModelKind: Int] = [:]

    init(store: ModelStore = .shared) {
        self.store = store
        refreshStates()
    }

    /// 根据沙盒现状刷新每个模型的状态(已完整→ready,否则 idle)。
    func refreshStates() {
        for kind in ModelKind.allCases {
            let total = ModelManifest.totalBytes(for: kind)
            if store.isComplete(for: kind) {
                states[kind] = DownloadState(phase: .ready, receivedBytes: total,
                                             totalBytes: total, bytesPerSecond: 0)
            } else if states[kind]?.phase == .downloading {
                continue  // 不打断进行中的下载
            } else {
                states[kind] = DownloadState(phase: .idle, receivedBytes: completedBytes(for: kind),
                                             totalBytes: total, bytesPerSecond: 0)
            }
        }
    }

    var isAnyDownloading: Bool {
        states.values.contains { $0.phase == .downloading }
    }

    /// 下载某个模型。幂等:已在下载或已就绪则忽略。
    func download(_ kind: ModelKind) {
        guard tasks[kind] == nil, states[kind]?.phase != .ready else { return }
        let total = ModelManifest.totalBytes(for: kind)
        states[kind] = DownloadState(phase: .downloading, receivedBytes: completedBytes(for: kind),
                                     totalBytes: total, bytesPerSecond: 0)
        lastSampleTime[kind] = Date()
        lastSampleBytes[kind] = completedBytes(for: kind)

        let task = Task { [weak self] in
            guard let self else { return }
            await self.run(kind)
        }
        tasks[kind] = task
    }

    func downloadAll() {
        for kind in ModelKind.userFacing { download(kind) }
    }

    /// 暂停下载。已下载的 .part 保留,下次从断点续传。
    func cancel(_ kind: ModelKind) {
        tasks[kind]?.cancel()
        tasks[kind] = nil
        let total = ModelManifest.totalBytes(for: kind)
        states[kind] = DownloadState(phase: .idle, receivedBytes: completedBytes(for: kind),
                                     totalBytes: total, bytesPerSecond: 0)
    }

    /// 旁路导入:从用户选择的文件夹拷入模型(现场重装兜底)。
    func importModel(_ kind: ModelKind, from folder: URL) throws {
        try store.importModel(kind, from: folder)
        refreshStates()
    }

    // MARK: - 内部

    private func run(_ kind: ModelKind) async {
        let files = ModelManifest.files(for: kind)
        let downloader = FileDownloader()
        var completedBefore = 0

        do {
            for file in files {
                if Task.isCancelled { return }
                let destination = store.fileURL(for: kind, relativePath: file.path)
                let base = completedBefore

                // 魔搭优先 + hf-mirror 回退:按序尝试候选源,任一成功即止;
                // 换源续传复用同一 .part(两源字节一致,Range 续传穿透)。
                let sources = ModelManifest.fileURLs(for: kind, file: file)
                var lastError: Error?
                var downloaded = false
                for (idx, source) in sources.enumerated() {
                    if Task.isCancelled { return }
                    do {
                        try await downloader.download(
                            from: source,
                            to: destination,
                            expectedBytes: file.bytes,
                            onProgress: { [weak self] received in
                                guard let self else { return }
                                Task { @MainActor in
                                    self.applyProgress(kind, currentTotal: base + received)
                                }
                            }
                        )
                        downloaded = true
                        break
                    } catch {
                        lastError = error
                        #if DEBUG
                        let more = idx + 1 < sources.count ? ",回退下一个源" : ""
                        print("[ModelDownload] \(file.path) 源 \(idx + 1)/\(sources.count) 失败:"
                              + "\(error.localizedDescription)\(more)")
                        #endif
                    }
                }
                if !downloaded { throw lastError ?? DownloadError.badStatus(0) }
                completedBefore += file.bytes
            }
            finish(kind, success: true, message: nil)
        } catch {
            if Task.isCancelled {
                // cancel() 已设置 idle 状态
            } else {
                finish(kind, success: false, message: error.localizedDescription)
            }
        }
    }

    private func applyProgress(_ kind: ModelKind, currentTotal: Int) {
        guard var state = states[kind], state.phase == .downloading else { return }
        let now = Date()
        if let lastTime = lastSampleTime[kind], let lastBytes = lastSampleBytes[kind] {
            let dt = now.timeIntervalSince(lastTime)
            if dt >= 0.5 {
                state.bytesPerSecond = Double(currentTotal - lastBytes) / dt
                lastSampleTime[kind] = now
                lastSampleBytes[kind] = currentTotal
            }
        }
        state.receivedBytes = currentTotal
        states[kind] = state
    }

    private func finish(_ kind: ModelKind, success: Bool, message: String?) {
        tasks[kind] = nil
        let total = ModelManifest.totalBytes(for: kind)
        if success {
            states[kind] = DownloadState(phase: .ready, receivedBytes: total,
                                         totalBytes: total, bytesPerSecond: 0)
        } else {
            states[kind] = DownloadState(phase: .failed(message ?? String(appLoc: "下载失败")),
                                         receivedBytes: completedBytes(for: kind),
                                         totalBytes: total, bytesPerSecond: 0)
        }
    }

    /// 已完整下载的文件字节之和(用于续传时的起始进度)。
    private func completedBytes(for kind: ModelKind) -> Int {
        ModelManifest.files(for: kind).reduce(0) { sum, file in
            store.localBytes(for: kind, relativePath: file.path) == file.bytes ? sum + file.bytes : sum
        }
    }
}