import Foundation
import RunnerCore
import Virtualization

/// Locates, downloads, and opens macOS restore images (IPSWs).
///
/// Two distinct notions of "restore image" get conflated easily, so this type
/// keeps them apart:
///
/// * `VZMacOSRestoreImage.latestSupported` returns an image whose `url` is a
///   **network** URL on Apple's CDN. It cannot be handed to `VZMacOSInstaller`.
/// * `VZMacOSRestoreImage.image(from:)` (or `load(from:)`) opens a **local
///   file** URL. That is what the installer needs.
///
/// So the pipeline is always: discover → download → load.
public struct IPSWProvider: Sendable {
    /// Where downloads are written, typically `<storeDir>/ipsw`.
    public let downloadDirectory: URL

    /// Creates a provider.
    ///
    /// - Parameter downloadDirectory: Destination directory for downloads.
    public init(downloadDirectory: URL) {
        self.downloadDirectory = downloadDirectory
    }

    /// Asks Apple for the newest restore image this host can run.
    ///
    /// - Returns: The CDN URL to download and the image's build version (e.g.
    ///   `25A354`), recorded into ``VMBundleConfig/macOSVersion``.
    /// - Throws: ``CoreError/notFound(_:)`` when Apple reports no supported
    ///   image (which also happens with no network).
    public func latestSupported() async throws -> (url: URL, buildVersion: String) {
        let image: VZMacOSRestoreImage
        do {
            image = try await VZMacOSRestoreImage.latestSupported
        } catch {
            // The framework reports "no supported image" and "could not reach
            // the CDN" identically, so the message has to cover both.
            throw CoreError.notFound(
                "no supported macOS restore image available: \(error.localizedDescription) "
                    + "(check network connectivity, or pass --ipsw with a local file)"
            )
        }
        return (image.url, image.buildVersion)
    }

    /// Downloads a restore image to ``downloadDirectory``.
    ///
    /// IPSWs are ~15 GB, so this reports progress and resumes nothing — a failed
    /// download is retried from scratch. The file is written to a `.partial`
    /// name and renamed on completion so an interrupted run never leaves a
    /// truncated file that looks valid.
    ///
    /// - Parameters:
    ///   - remoteURL: The CDN URL from ``latestSupported()``.
    ///   - progress: Called with a fraction in `0...1`. May be called from an
    ///     arbitrary thread.
    /// - Returns: The local file URL.
    public func download(
        from remoteURL: URL,
        progress: (@Sendable (Double) -> Void)? = nil
    ) async throws -> URL {
        // A file URL is already local; nothing to do.
        if remoteURL.isFileURL {
            progress?(1.0)
            return remoteURL
        }

        let fileManager = FileManager.default
        try fileManager.createDirectory(at: downloadDirectory, withIntermediateDirectories: true)

        let fileName = IPSWProvider.localFileName(for: remoteURL)
        let finalURL = downloadDirectory.appendingPathComponent(fileName)

        // A previously completed download is reused: only fully-written files
        // ever get the final name.
        if fileManager.fileExists(atPath: finalURL.path) {
            progress?(1.0)
            return finalURL
        }

        let partialURL = downloadDirectory.appendingPathComponent(fileName + ".partial")
        try? fileManager.removeItem(at: partialURL)

        let configuration = URLSessionConfiguration.default
        // The default 7-day resource timeout is useless as a failure signal and
        // the default 60 s request timeout only bounds the *response start*.
        // Six hours is generous for 15 GB on a slow link and still finite.
        configuration.timeoutIntervalForRequest = 120
        configuration.timeoutIntervalForResource = 6 * 60 * 60
        configuration.waitsForConnectivity = true

        let delegate = IPSWDownloadProgressDelegate(onProgress: progress)
        let session = URLSession(configuration: configuration)
        defer { session.finishTasksAndInvalidate() }

        let temporaryURL: URL
        let response: URLResponse
        do {
            (temporaryURL, response) = try await session.download(from: remoteURL, delegate: delegate)
        } catch {
            throw CoreError.notFound(
                "restore image download failed for \(remoteURL.absoluteString): \(error.localizedDescription)"
            )
        }

        if let http = response as? HTTPURLResponse, !(200..<300).contains(http.statusCode) {
            try? fileManager.removeItem(at: temporaryURL)
            throw CoreError.notFound(
                "restore image download failed: HTTP \(http.statusCode) for \(remoteURL.absoluteString)"
            )
        }

        // Move into `.partial` first, then rename: the final name is the
        // "this file is complete" marker that the reuse check above trusts.
        do {
            try fileManager.moveItem(at: temporaryURL, to: partialURL)
            try fileManager.moveItem(at: partialURL, to: finalURL)
        } catch {
            try? fileManager.removeItem(at: temporaryURL)
            try? fileManager.removeItem(at: partialURL)
            throw CoreError.provisioningFailed(
                "could not store the downloaded restore image at \(finalURL.path): \(error.localizedDescription)"
            )
        }

        progress?(1.0)
        return finalURL
    }

    /// Opens a local IPSW.
    ///
    /// Symlinks are resolved first: `VZMacOSRestoreImage` rejects a symlinked
    /// path, and `~/Downloads` paths handed in by users are frequently symlinked
    /// through `/Users` → `/System/Volumes/Data/Users`.
    ///
    /// - Parameter localPath: Path to an `.ipsw` file.
    /// - Returns: The loaded restore image.
    /// - Throws: ``CoreError/notFound(_:)`` when the path does not exist, or
    ///   ``CoreError/configInvalid(_:)`` when it is not a local file URL.
    public func load(localPath: String) async throws -> VZMacOSRestoreImage {
        let expanded = (localPath as NSString).expandingTildeInPath
        var url = URL(fileURLWithPath: expanded)
        // Must happen before the framework ever sees the URL.
        url.resolveSymlinksInPath()

        guard url.isFileURL else {
            throw CoreError.configInvalid(
                "restore image path must be a local file, got \(url.absoluteString)"
            )
        }

        // `VZMacOSRestoreImage.image(from:)` raises an Objective-C exception —
        // not a Swift error — when handed a non-file or missing path, and an
        // ObjC exception cannot be caught here. So the existence check is not
        // politeness; it is the only thing standing between a typo and a crash.
        var isDirectory: ObjCBool = false
        guard FileManager.default.fileExists(atPath: url.path, isDirectory: &isDirectory),
            !isDirectory.boolValue
        else {
            throw CoreError.notFound("restore image not found at \(url.path)")
        }

        do {
            return try await VZMacOSRestoreImage.image(from: url)
        } catch {
            throw CoreError.provisioningFailed(
                "could not read restore image at \(url.path): \(error.localizedDescription)"
            )
        }
    }

    /// Convenience: discover, download if not already present, and load.
    ///
    /// - Parameter progress: Download progress callback.
    /// - Returns: The loaded image and the local file it came from.
    public func fetchLatest(
        progress: (@Sendable (Double) -> Void)? = nil
    ) async throws -> (image: VZMacOSRestoreImage, localURL: URL) {
        let (remoteURL, _) = try await latestSupported()
        let localURL = try await download(from: remoteURL, progress: progress)
        // Deliberately reloaded from the local file: the image returned by
        // `latestSupported` carries a network URL, and the installer needs one
        // whose `url` is on disk.
        let image = try await load(localPath: localURL.path)
        return (image, localURL)
    }

    // MARK: - Helpers

    /// The on-disk name for a remote restore image.
    ///
    /// Apple's CDN names are already unique (`UniversalMac_15.2_24C101_Restore.ipsw`);
    /// anything else falls back to a name derived from the URL so two different
    /// sources cannot collide.
    static func localFileName(for remoteURL: URL) -> String {
        let candidate = remoteURL.lastPathComponent
        if candidate.lowercased().hasSuffix(".ipsw"), candidate.count > ".ipsw".count {
            return candidate
        }
        let digest = abs(remoteURL.absoluteString.hashValue)
        return "restore-\(String(digest, radix: 16)).ipsw"
    }
}

/// Reports `URLSession` download progress as a fraction.
///
/// A task-scoped delegate is the only way to observe byte progress from the
/// `async` download API; the `didFinishDownloadingTo` callback is deliberately
/// *not* implemented, because the `async` variant owns the temporary file.
private final class IPSWDownloadProgressDelegate: NSObject, URLSessionDownloadDelegate, @unchecked Sendable {
    private let onProgress: (@Sendable (Double) -> Void)?

    init(onProgress: (@Sendable (Double) -> Void)?) {
        self.onProgress = onProgress
    }

    func urlSession(
        _ session: URLSession,
        downloadTask: URLSessionDownloadTask,
        didWriteData bytesWritten: Int64,
        totalBytesWritten: Int64,
        totalBytesExpectedToWrite: Int64
    ) {
        // A chunked response reports -1 for the expected length; report nothing
        // rather than a nonsense fraction.
        guard totalBytesExpectedToWrite > 0 else { return }
        let fraction = Double(totalBytesWritten) / Double(totalBytesExpectedToWrite)
        onProgress?(min(max(fraction, 0), 1))
    }

    func urlSession(
        _ session: URLSession,
        downloadTask: URLSessionDownloadTask,
        didFinishDownloadingTo location: URL
    ) {
        // Intentionally empty. `URLSession.download(from:delegate:)` moves the
        // file itself; doing anything here would race with it.
    }
}