shotdeck/Sources/Shotdeck/UpdateChecker.swift

import AppKit
import CryptoKit
import Foundation
import Security
import ShotdeckCore

/// Built-in updater. Checks an appcast, stages a verified payload, and installs
/// only when the user clicks the menu row — never automatically.
@MainActor
final class UpdateChecker {
    static let appcastURLDefaultsKey = "ai.flowmaster.shotdeck.appcastURL"
    static let defaultAppcastURL = URL(string: "https://get.baobab-ts.com/cowork/redline/appcast.json")!
    static let defaultInstallTarget = URL(fileURLWithPath: "/Applications/Redline.app")

    /// Required bundle identifier for any staged or installed payload.
    static let expectedBundleIdentifier = "ai.flowmaster.shotdeck"
    /// Developer team identifiers MMD ships Redline under. Overridable only for the self-test.
    static let allowedTeamIdentifiers: Set<String> = ["PWMCBMX5M8", "L3N9S54CN3"]

    private(set) var availableUpdate: (version: String, notes: String)?
    private(set) var stagedAppURL: URL?
    private(set) var statusMessage: String?
    private(set) var lastCheckedAt: Date?
    private(set) var isCheckingNow: Bool = false

    var onChecked: (() -> Void)?
    /// Fired whenever `isCheckingNow` flips, so a UI can show "Checking…" for the
    /// whole duration of a check rather than only after it lands.
    var onCheckingChanged: ((Bool) -> Void)?

    private let urlSession: URLSession
    private var repeatingTimer: Timer?
    private var firstCheckTask: Task<Void, Never>?
    private var stagingDirectory: URL?
    /// Snapshot-only override for previousVersion; when snapshotUsesPreviousVersionOverride is true,
    /// this value (including nil) is returned instead of checking the file system.
    var snapshotPreviousVersionOverride: String?
    var snapshotUsesPreviousVersionOverride: Bool = false
    /// Test seam: override appcast JSON. When set, returns this instead of fetching from URL.
    var testAppcastJSON: String?

    init() {
        let config = URLSessionConfiguration.ephemeral
        config.timeoutIntervalForRequest = 30
        config.timeoutIntervalForResource = 600
        config.httpCookieAcceptPolicy = .never
        config.httpShouldSetCookies = false
        config.httpCookieStorage = nil
        config.urlCache = nil
        urlSession = URLSession(configuration: config)
    }

    /// First check 10 seconds after start, then every 6 hours. Stages only — never installs.
    func startSchedule() {
        firstCheckTask?.cancel()
        firstCheckTask = Task { [weak self] in
            try? await Task.sleep(for: .seconds(10))
            guard !Task.isCancelled else { return }
            await self?.checkNow()
        }
        repeatingTimer?.invalidate()
        let timer = Timer(timeInterval: 6 * 60 * 60, repeats: true) { [weak self] _ in
            Task { @MainActor in
                await self?.checkNow()
            }
        }
        RunLoop.main.add(timer, forMode: .common)
        repeatingTimer = timer
    }

    /// Checks the appcast and stages a newer, signature-verified payload.
    /// `manual` only affects the status message shown when already up to date —
    /// a user-initiated check says so; the silent background check stays quiet.
    func checkNow(manual: Bool = false) async {
        guard !isCheckingNow else { return }
        isCheckingNow = true
        onCheckingChanged?(true)
        defer {
            isCheckingNow = false
            onCheckingChanged?(false)
        }
        lastCheckedAt = Date()

        let appcast: Appcast
        do {
            appcast = try await fetchAppcast()
        } catch {
            statusMessage = "Could not check for updates."
            onChecked?()
            return
        }

        guard Self.isNewer(appcast.version, than: Self.currentVersion()) else {
            clearOffer()
            if manual {
                let timestamp = DubaiTime.checkTime(lastCheckedAt ?? Date())
                statusMessage = "Redline \(Self.currentVersion()) is up to date, checked \(timestamp)"
            } else {
                statusMessage = nil
            }
            onChecked?()
            return
        }

        do {
            try await downloadAndStage(appcast)
            availableUpdate = (version: appcast.version, notes: appcast.notes ?? "")
            statusMessage = manual ? "Update to \(appcast.version) is ready" : nil
        } catch UpdateCheckError.checksumMismatch {
            discardStaging()
            availableUpdate = nil
            statusMessage = "Update file failed the checksum — not installed."
        } catch UpdateCheckError.signatureInvalid {
            discardStaging()
            availableUpdate = nil
            statusMessage = "Update is not signed by MMD — not installed."
        } catch {
            discardStaging()
            availableUpdate = nil
            statusMessage = "The update could not be prepared."
        }
        onChecked?()
    }

    /// Installs the staged app onto `target` atomically, keeping exactly one rollback
    /// copy (`Redline.app.previous`), then hands off to a relaunch and quits.
    /// Never deletes the old app before the new one is verified in place.
    func installStaged(to target: URL = UpdateChecker.defaultInstallTarget) {
        guard let staged = stagedAppURL else {
            statusMessage = "No update is staged."
            onChecked?()
            return
        }

        let targetDir = target.deletingLastPathComponent()
        let previousURL = targetDir.appendingPathComponent("Redline.app.previous")

        do {
            try FileManager.default.createDirectory(at: targetDir, withIntermediateDirectories: true)

            let replacementDir = try FileManager.default.url(
                for: .itemReplacementDirectory,
                in: .userDomainMask,
                appropriateFor: target,
                create: true
            )
            defer { try? FileManager.default.removeItem(at: replacementDir) }

            let newCopy = replacementDir.appendingPathComponent(target.lastPathComponent)
            try Self.runProcess(executable: "/usr/bin/ditto", arguments: [staged.path, newCopy.path])

            // Exactly one rollback copy is kept — drop any older one before this install.
            if FileManager.default.fileExists(atPath: previousURL.path) {
                try FileManager.default.removeItem(at: previousURL)
            }

            if FileManager.default.fileExists(atPath: target.path) {
                _ = try FileManager.default.replaceItemAt(
                    target,
                    withItemAt: newCopy,
                    backupItemName: previousURL.lastPathComponent,
                    options: [.withoutDeletingBackupItem]
                )
            } else {
                try FileManager.default.moveItem(at: newCopy, to: target)
            }
        } catch {
            statusMessage = "The update could not be installed."
            onChecked?()
            return
        }

        // Defense in depth: re-verify what actually landed on disk, not just the staged copy.
        do {
            try Self.verifySignature(of: target)
        } catch {
            statusMessage = "The update was installed but failed verification."
            onChecked?()
            return
        }

        discardStaging()
        availableUpdate = nil
        relaunch(target: target)
    }

    /// Swaps `Redline.app.previous` back into place, verifying its signature first.
    /// The just-replaced (newer) app becomes the new `.previous` — a revert is
    /// itself reversible.
    func revertToPrevious(target: URL = UpdateChecker.defaultInstallTarget) {
        let targetDir = target.deletingLastPathComponent()
        let previousURL = targetDir.appendingPathComponent("Redline.app.previous")

        guard FileManager.default.fileExists(atPath: previousURL.path) else {
            statusMessage = "No previous version to revert to."
            onChecked?()
            return
        }

        do {
            try Self.verifySignature(of: previousURL)
        } catch {
            statusMessage = "The previous version failed verification and was not restored."
            onChecked?()
            return
        }

        do {
            // `previousURL` cannot be handed to replaceItemAt directly: its own path
            // IS the requested backup name, so the backup step would clobber it
            // before the swap ever reads it. Stage a throwaway copy first, exactly
            // like installStaged does for the forward direction.
            let replacementDir = try FileManager.default.url(
                for: .itemReplacementDirectory,
                in: .userDomainMask,
                appropriateFor: target,
                create: true
            )
            defer { try? FileManager.default.removeItem(at: replacementDir) }

            let newCopy = replacementDir.appendingPathComponent(target.lastPathComponent)
            try Self.runProcess(executable: "/usr/bin/ditto", arguments: [previousURL.path, newCopy.path])
            try FileManager.default.removeItem(at: previousURL)

            _ = try FileManager.default.replaceItemAt(
                target,
                withItemAt: newCopy,
                backupItemName: previousURL.lastPathComponent,
                options: [.withoutDeletingBackupItem]
            )
        } catch {
            statusMessage = "Could not revert to the previous version."
            onChecked?()
            return
        }

        relaunch(target: target)
    }

    /// The version recorded in `Redline.app.previous`'s Info.plist, or nil when no
    /// rollback copy exists. Respects the snapshot-only override for panel rendering.
    func previousVersion(target: URL = UpdateChecker.defaultInstallTarget) -> String? {
        if snapshotUsesPreviousVersionOverride {
            return snapshotPreviousVersionOverride
        }
        let previousURL = target.deletingLastPathComponent().appendingPathComponent("Redline.app.previous")
        let plistURL = previousURL.appendingPathComponent("Contents/Info.plist")
        guard let plist = NSDictionary(contentsOf: plistURL) as? [String: Any] else { return nil }
        return plist["CFBundleShortVersionString"] as? String
    }

    static func resolvedAppcastURL() -> URL {
        if let env = ProcessInfo.processInfo.environment["REDLINE_APPCAST_URL"],
           !env.isEmpty,
           let url = URL(string: env)
        {
            return url
        }
        if let stored = UserDefaults.standard.string(forKey: appcastURLDefaultsKey),
           !stored.isEmpty,
           let url = URL(string: stored)
        {
            return url
        }
        return defaultAppcastURL
    }

    static func currentVersion() -> String {
        Bundle.main.object(forInfoDictionaryKey: "CFBundleShortVersionString") as? String ?? "0.0.0"
    }

    static func isNewer(_ candidate: String, than current: String) -> Bool {
        let a = semverParts(candidate)
        let b = semverParts(current)
        for i in 0..<3 {
            if a[i] != b[i] { return a[i] > b[i] }
        }
        return false
    }

    static func sha256Hex(_ data: Data) -> String {
        SHA256.hash(data: data).map { String(format: "%02x", $0) }.joined()
    }

    /// Validates the code signature of the app at `appURL`: strictly, across all
    /// architectures and nested code, then checks its bundle identifier and team
    /// identifier against `expectedBundleIdentifier` / the allowed-teams set.
    /// `REDLINE_ALLOWED_TEAMS` (comma separated) overrides the allowed set — for
    /// the self-test only, so it can accept a locally re-signed fake bundle.
    static func verifySignature(of appURL: URL) throws {
        var staticCode: SecStaticCode?
        let createStatus = SecStaticCodeCreateWithPath(appURL as CFURL, [], &staticCode)
        guard createStatus == errSecSuccess, let code = staticCode else {
            throw UpdateCheckError.signatureInvalid(
                "could not read a code signature (status \(createStatus))"
            )
        }

        let validityFlags = SecCSFlags(
            rawValue: kSecCSStrictValidate | kSecCSCheckAllArchitectures | kSecCSCheckNestedCode
        )
        var validityError: Unmanaged<CFError>?
        let validityStatus = SecStaticCodeCheckValidityWithErrors(code, validityFlags, nil, &validityError)
        guard validityStatus == errSecSuccess else {
            let detail = (validityError?.takeRetainedValue()).map { String(describing: $0) } ?? "status \(validityStatus)"
            throw UpdateCheckError.signatureInvalid("signature is not valid: \(detail)")
        }

        var signingInfo: CFDictionary?
        let infoStatus = SecCodeCopySigningInformation(
            code,
            SecCSFlags(rawValue: kSecCSSigningInformation),
            &signingInfo
        )
        guard infoStatus == errSecSuccess, let info = signingInfo as? [String: Any] else {
            throw UpdateCheckError.signatureInvalid("could not read signing information (status \(infoStatus))")
        }

        let identifier = info[kSecCodeInfoIdentifier as String] as? String
        guard identifier == expectedBundleIdentifier else {
            throw UpdateCheckError.signatureInvalid(
                "unexpected bundle identifier: \(identifier ?? "nil")"
            )
        }

        let teamIdentifier = info[kSecCodeInfoTeamIdentifier as String] as? String
        guard let teamIdentifier, resolvedAllowedTeamIdentifiers().contains(teamIdentifier) else {
            throw UpdateCheckError.signatureInvalid(
                "unexpected team identifier: \(teamIdentifier ?? "nil")"
            )
        }
    }

    private static func resolvedAllowedTeamIdentifiers() -> Set<String> {
        if let env = ProcessInfo.processInfo.environment["REDLINE_ALLOWED_TEAMS"], !env.isEmpty {
            let parts = env.split(separator: ",")
                .map { $0.trimmingCharacters(in: .whitespaces) }
                .filter { !$0.isEmpty }
            if !parts.isEmpty {
                return Set(parts)
            }
        }
        return allowedTeamIdentifiers
    }

    // MARK: - Private

    private struct Appcast: Decodable {
        var version: String
        var zipURL: URL
        var sha256: String
        var notes: String?
    }

    private enum UpdateCheckError: Error {
        case checksumMismatch
        case invalidPayload
        case httpStatus(Int)
        case processFailed(String)
        case signatureInvalid(String)
    }

    private func fetchAppcast() async throws -> Appcast {
        let data: Data
        if let testJSON = testAppcastJSON {
            data = testJSON.data(using: .utf8) ?? Data()
        } else {
            data = try await fetchData(from: Self.resolvedAppcastURL())
        }
        return try JSONDecoder().decode(Appcast.self, from: data)
    }

    private func fetchData(from url: URL) async throws -> Data {
        if url.isFileURL {
            return try Data(contentsOf: url)
        }
        let (data, response) = try await urlSession.data(from: url)
        if let http = response as? HTTPURLResponse, !(200...299).contains(http.statusCode) {
            throw UpdateCheckError.httpStatus(http.statusCode)
        }
        return data
    }

    private func downloadAndStage(_ appcast: Appcast) async throws {
        let zipData = try await fetchData(from: appcast.zipURL)
        let expected = appcast.sha256.trimmingCharacters(in: .whitespacesAndNewlines)
        let actual = Self.sha256Hex(zipData)
        guard actual.caseInsensitiveCompare(expected) == .orderedSame else {
            throw UpdateCheckError.checksumMismatch
        }

        discardStaging()
        let root = FileManager.default.temporaryDirectory
            .appendingPathComponent("shotdeck-update-\(UUID().uuidString)", isDirectory: true)
        try FileManager.default.createDirectory(at: root, withIntermediateDirectories: true)
        stagingDirectory = root

        let zipURL = root.appendingPathComponent("update.zip")
        try zipData.write(to: zipURL)

        let extracted = root.appendingPathComponent("extracted", isDirectory: true)
        try FileManager.default.createDirectory(at: extracted, withIntermediateDirectories: true)
        try Self.runProcess(
            executable: "/usr/bin/ditto",
            arguments: ["-x", "-k", zipURL.path, extracted.path]
        )

        guard let appURL = Self.findRedlineApp(in: extracted) else {
            throw UpdateCheckError.invalidPayload
        }
        let executable = appURL.appendingPathComponent("Contents/MacOS/Shotdeck")
        guard FileManager.default.fileExists(atPath: executable.path) else {
            throw UpdateCheckError.invalidPayload
        }
        try Self.verifySignature(of: appURL)
        stagedAppURL = appURL
    }

    private func clearOffer() {
        availableUpdate = nil
        discardStaging()
    }

    private func discardStaging() {
        if let stagingDirectory {
            try? FileManager.default.removeItem(at: stagingDirectory)
        }
        stagingDirectory = nil
        stagedAppURL = nil
    }

    /// Spawns a detached watcher that waits for this process to exit, then reopens
    /// `target`, and quits. Never called during the self-test, so the in-process
    /// assertions after `installStaged`/`revertToPrevious` can still run.
    private func relaunch(target: URL) {
        guard ProcessInfo.processInfo.environment["SHOTDECK_UPDATE_SELFTEST"] == nil else { return }

        let ownPID = ProcessInfo.processInfo.processIdentifier
        let script = "while kill -0 \(ownPID) 2>/dev/null; do sleep 0.2; done; " +
            "/usr/bin/open -n \(Self.shellQuoted(target.path))"
        let process = Process()
        process.executableURL = URL(fileURLWithPath: "/bin/sh")
        process.arguments = ["-c", script]
        process.standardInput = FileHandle.nullDevice
        process.standardOutput = FileHandle.nullDevice
        process.standardError = FileHandle.nullDevice
        do {
            try process.run()
        } catch {
            statusMessage = "The update was installed but Redline could not relaunch. Open it from Applications."
            onChecked?()
            return
        }
        NSApp.terminate(nil)
    }

    private static func shellQuoted(_ path: String) -> String {
        "'" + path.replacingOccurrences(of: "'", with: "'\\''") + "'"
    }

    private static func findRedlineApp(in directory: URL) -> URL? {
        let fm = FileManager.default
        let direct = directory.appendingPathComponent("Redline.app")
        if fm.fileExists(atPath: direct.path) { return direct }

        guard let enumerator = fm.enumerator(
            at: directory,
            includingPropertiesForKeys: [.isDirectoryKey],
            options: [.skipsHiddenFiles]
        ) else { return nil }

        while let item = enumerator.nextObject() as? URL {
            if item.lastPathComponent == "Redline.app" {
                return item
            }
            if item.pathExtension == "app" {
                enumerator.skipDescendants()
            }
        }
        return nil
    }

    private static func semverParts(_ string: String) -> [Int] {
        let core = string.split(separator: "-").first.map(String.init) ?? string
        var parts = core.split(separator: ".").prefix(3).map { Int($0) ?? 0 }
        while parts.count < 3 { parts.append(0) }
        return parts
    }

    private static func runProcess(executable: String, arguments: [String]) throws {
        let process = Process()
        process.executableURL = URL(fileURLWithPath: executable)
        process.arguments = arguments
        let err = Pipe()
        process.standardError = err
        process.standardOutput = Pipe()
        try process.run()
        process.waitUntilExit()
        guard process.terminationStatus == 0 else {
            let message = String(data: err.fileHandleForReading.readDataToEndOfFile(), encoding: .utf8) ?? ""
            throw UpdateCheckError.processFailed("\(executable) failed: \(message)")
        }
    }
}