//
//  DownloadManager.swift
//  Zyquo MLX
//
//  Author: Simon-Pierre Boucher
//  Mail: contact@spboucher.ai
//

import Foundation

/// Progress of one repo download.
struct DownloadProgress: Sendable {
    var repo: String
    var completedBytes: Int64
    var totalBytes: Int64
    var currentFile: String
    var fraction: Double { totalBytes > 0 ? Double(completedBytes) / Double(totalBytes) : 0 }
}

enum DownloadEvent: Sendable {
    case progress(DownloadProgress)
    case finished(LocalModel)
    case failed(String)
}

/// Downloads Hub repos into the model library with resumable, Range-based
/// transfers (contract verified in docs/MODELS.md §1.3: fresh resolve URL +
/// `Range: bytes=N-` → 206; CDN URLs expire so we never persist them).
actor DownloadManager {

    static let shared = DownloadManager()

    private var activeTasks: [String: Task<Void, Never>] = [:]

    var activeRepos: [String] { Array(activeTasks.keys) }

    /// Download a full repo. Partially-downloaded files resume from their
    /// current byte count.
    func download(repo: String) -> AsyncStream<DownloadEvent> {
        let (stream, continuation) = AsyncStream.makeStream(of: DownloadEvent.self)

        let task = Task {
            let directory = PersistenceService.modelsDirectory
                .appendingPathComponent(repo.replacingOccurrences(of: "/", with: "--"), isDirectory: true)
            do {
                try FileManager.default.createDirectory(at: directory, withIntermediateDirectories: true)
                let files = try await HubService.shared.files(repo: repo)
                    .filter { !$0.path.hasPrefix(".") }
                let totalBytes = files.reduce(0) { $0 + $1.size }
                var completed: Int64 = 0

                for file in files {
                    try Task.checkCancellation()
                    let destination = directory.appendingPathComponent(file.path)
                    try FileManager.default.createDirectory(
                        at: destination.deletingLastPathComponent(),
                        withIntermediateDirectories: true)

                    let existing = (try? destination.resourceValues(forKeys: [.fileSizeKey]))?
                        .fileSize.map(Int64.init) ?? 0
                    if existing == file.size, file.size > 0 {
                        completed += file.size
                        continuation.yield(.progress(DownloadProgress(
                            repo: repo, completedBytes: completed,
                            totalBytes: totalBytes, currentFile: file.path)))
                        continue
                    }

                    let completedSoFar = completed
                    try await downloadFile(
                        repo: repo, path: file.path, to: destination,
                        resumeFrom: existing < file.size ? existing : 0
                    ) { bytesSoFar in
                        continuation.yield(.progress(DownloadProgress(
                            repo: repo, completedBytes: completedSoFar + bytesSoFar,
                            totalBytes: totalBytes, currentFile: file.path)))
                    }
                    completed += file.size
                }

                let model = try await ModelStore.shared.describe(directory: directory)
                continuation.yield(.finished(model))
            } catch is CancellationError {
                continuation.yield(.failed("Download paused — it will resume from where it stopped."))
            } catch {
                continuation.yield(.failed(error.localizedDescription))
            }
            self.clearTask(repo: repo)
            continuation.finish()
        }
        activeTasks[repo] = task
        continuation.onTermination = { termination in
            if case .cancelled = termination { task.cancel() }
        }
        return stream
    }

    func cancel(repo: String) {
        activeTasks[repo]?.cancel()
    }

    private func clearTask(repo: String) {
        activeTasks[repo] = nil
    }

    /// Stream one file to disk with Range resume and periodic progress.
    private func downloadFile(
        repo: String, path: String, to destination: URL,
        resumeFrom: Int64,
        progress: @Sendable (Int64) -> Void
    ) async throws {
        var request = URLRequest(url: HubService.resolveURL(repo: repo, path: path))
        if resumeFrom > 0 {
            request.setValue("bytes=\(resumeFrom)-", forHTTPHeaderField: "Range")
        }
        if let token = HFTokenStore.token, !token.isEmpty {
            request.setValue("Bearer \(token)", forHTTPHeaderField: "Authorization")
        }

        let (bytes, response) = try await URLSession.shared.bytes(for: request)
        guard let http = response as? HTTPURLResponse,
            http.statusCode == 200 || http.statusCode == 206
        else {
            throw HubServiceError.badResponse((response as? HTTPURLResponse)?.statusCode ?? 0)
        }

        let appending = http.statusCode == 206 && resumeFrom > 0
        if !appending {
            FileManager.default.createFile(atPath: destination.path, contents: nil)
        }
        let handle = try FileHandle(forWritingTo: destination)
        defer { try? handle.close() }
        if appending {
            try handle.seekToEnd()
        } else {
            try handle.truncate(atOffset: 0)
        }

        var written: Int64 = appending ? resumeFrom : 0
        var buffer = Data(capacity: 1 << 20)
        var lastReport = Date()
        for try await byte in bytes {
            buffer.append(byte)
            if buffer.count >= 1 << 20 {
                try handle.write(contentsOf: buffer)
                written += Int64(buffer.count)
                buffer.removeAll(keepingCapacity: true)
                if Date().timeIntervalSince(lastReport) > 0.2 {
                    progress(written)
                    lastReport = Date()
                }
                try Task.checkCancellation()
            }
        }
        if !buffer.isEmpty {
            try handle.write(contentsOf: buffer)
            written += Int64(buffer.count)
        }
        progress(written)
    }
}