//
//  Helper.swift
//  SwanDesk — app-side XPC client + SMAppService daemon registration.
//
//  Ported from AppVersionChecker. Registers the privileged daemon (SMAppService),
//  detects a stale/outdated installed daemon and forces re-approval, and keeps a
//  lazily-recreated NSXPCConnection. Adds a reverse-XPC exported object so the
//  daemon can push status/log back to the app.
//

import Foundation
import ServiceManagement

/// Reverse-XPC object the daemon calls to push live status/log. Forwards to
/// closures the ConnectionManager installs.
final class VPNClient: NSObject, VPNClientProtocol {
    var onStatus: ((VPNRuntimeStatus) -> Void)?
    var onLog: ((String) -> Void)?

    func statusChanged(_ data: Data) {
        guard let status = try? JSONDecoder().decode(VPNRuntimeStatus.self, from: data) else { return }
        onStatus?(status)
    }

    func logLine(_ line: String) {
        onLog?(line)
    }
}

final class Helper {
    static let shared = Helper()

    private var xpcConnection: NSXPCConnection?
    private let lock = NSLock()

    /// The reverse-callback object exported on every connection. Held here so its
    /// closures (installed by ConnectionManager) survive connection recreation.
    let client = VPNClient()

    private init() {}

    // MARK: Registration / connection bring-up

    /// Bring the daemon up and confirm the app can actually talk to it. Returns
    /// `true` once a probe round-trips; `false` only after every recovery attempt
    /// has failed — the caller then refuses to run rather than present a UI that
    /// silently can't reach the helper.
    ///
    /// Flow: an idempotent `register()` re-points launchd at the on-disk helper (no
    /// re-approval prompt when already enabled, and it heals a stale record left by
    /// a rebuild that re-signed the helper) → binary-change check (if the on-disk
    /// mtime changed since we last launched, the old process is stale; terminate it
    /// so launchd relaunches with the new binary) → probe. If the probe still fails
    /// the registration is wedged, so escalate to a clean `unregister()` +
    /// `register()` — retried with backoff because the `register()` immediately after
    /// an `unregister()` transiently throws "Operation not permitted" (status
    /// notRegistered) until launchd releases the old job. Re-probe after each round.
    func establishConnection() async -> Bool {
        let service = SMAppService.daemon(plistName: "SwanDeskHelperDaemon.plist")

        registerDaemon(service)
        if helperBinaryChanged() {
            await terminateRunningHelper()
            recordHelperMtime()
        }
        clearConnection()
        if await probeReachable() { return true }

        // If the daemon is merely awaiting the user's approval, churning
        // unregister/register won't make it reachable — only approval will (and it
        // would re-pop System Settings each round). Bail to the caller's guidance.
        if service.status == .requiresApproval {
            Log("Helper: daemon requires approval — unreachable until the user enables it")
            return false
        }

        for attempt in 1...3 {
            Log("Helper: daemon unreachable — clean re-register \(attempt)/3")
            await cleanReRegister(service)
            clearConnection()
            if await probeReachable() { return true }
        }
        Log("Helper: daemon still unreachable after recovery — giving up (status \(describe(service.status)))")
        return false
    }

    /// True if `getVersion` round-trips (the XPC channel is alive). Resolves false
    /// on connection error or after a timeout — never leaks the continuation.
    private func probeReachable(timeout: TimeInterval = 6) async -> Bool {
        await withCheckedContinuation { cont in
            let once = ResumeOnce(cont)
            guard let proxy = connection()?.remoteObjectProxyWithErrorHandler({ error in
                Log("Helper: probe connection error: \(error.localizedDescription)")
                once.resume(false)
            }) as? VPNHelperProtocol else { once.resume(false); return }
            DispatchQueue.global().asyncAfter(deadline: .now() + timeout) { once.resume(false) }
            proxy.getVersion { installed in
                Log("Helper: installed helper version \(installed)")
                once.resume(true)
            }
        }
    }

    // MARK: Binary-change detection

    private static let helperMtimeKey = "SwanDesk.helperMtime"

    private func helperBinaryMtime() -> Double? {
        let url = Bundle.main.bundleURL
            .appendingPathComponent("Contents/Library/LaunchServices/fi.juvex.swandeskhelper")
        guard let attrs = try? FileManager.default.attributesOfItem(atPath: url.path),
              let date = attrs[.modificationDate] as? Date else { return nil }
        return date.timeIntervalSince1970
    }

    private func helperBinaryChanged() -> Bool {
        guard let mtime = helperBinaryMtime() else { return false }
        let stored = UserDefaults.standard.double(forKey: Self.helperMtimeKey)
        return stored == 0 || stored != mtime
    }

    private func recordHelperMtime() {
        guard let mtime = helperBinaryMtime() else { return }
        UserDefaults.standard.set(mtime, forKey: Self.helperMtimeKey)
    }

    /// Tell the running helper (if any) to exit cleanly so launchd restarts it with
    /// the newly-registered binary. Best-effort: if nothing is running the XPC call
    /// fails silently and we fall through to the normal probe path.
    private func terminateRunningHelper() async {
        Log("Helper: helper binary changed — terminating stale running instance")
        await withCheckedContinuation { (cont: CheckedContinuation<Void, Never>) in
            let lock = NSLock()
            var fired = false
            let fire = {
                lock.lock(); defer { lock.unlock() }
                guard !fired else { return }
                fired = true
                cont.resume()
            }
            guard let proxy = connection()?.remoteObjectProxyWithErrorHandler({ _ in fire() }) as? VPNHelperProtocol else {
                fire(); return
            }
            // Timeout in case the old helper hangs before exiting.
            DispatchQueue.global().asyncAfter(deadline: .now() + 4) { fire() }
            proxy.terminate { fire() }
        }
        clearConnection()
        // Give the helper's SIGTERM handler time to finish teardown and exit.
        try? await Task.sleep(nanoseconds: 1_000_000_000)

        // If the old process is still reachable (e.g. the pre-terminate() binary
        // didn't recognise the XPC method), escalate: unregister makes launchd
        // SIGTERM it, then re-register so the new binary starts on the next probe.
        if await probeReachable(timeout: 1) {
            Log("Helper: XPC terminate had no effect — forcing via SMAppService unregister")
            let service = SMAppService.daemon(plistName: "SwanDeskHelperDaemon.plist")
            await withCheckedContinuation { (cont: CheckedContinuation<Void, Never>) in
                service.unregister { _ in cont.resume() }
            }
            clearConnection()
            try? await Task.sleep(nanoseconds: 500_000_000)
            registerDaemon(service)
            try? await Task.sleep(nanoseconds: 1_000_000_000)
        }
    }

    /// Unregister, then register — retrying the register with backoff to ride out
    /// the post-unregister "Operation not permitted" window. Unlike a one-shot
    /// re-register, this never leaves the daemon stuck `notRegistered`.
    private func cleanReRegister(_ service: SMAppService) async {
        await withCheckedContinuation { (cont: CheckedContinuation<Void, Never>) in
            service.unregister { error in
                if let error = error { Log("Helper: unregister failed: \(error.localizedDescription)") }
                cont.resume()
            }
        }
        for attempt in 1...4 {
            try? await Task.sleep(nanoseconds: UInt64(attempt) * 500_000_000)   // 0.5,1,1.5,2s
            do {
                try service.register()
                Log("Helper: re-register OK (status \(describe(service.status)))")
                if service.status == .requiresApproval {
                    await MainActor.run { SMAppService.openSystemSettingsLoginItems() }
                }
                return
            } catch {
                Log("Helper: re-register attempt \(attempt)/4 failed (status \(describe(service.status))): \(error.localizedDescription)")
            }
        }
    }

    private func registerDaemon(_ service: SMAppService) {
        do {
            try service.register()
            Log("Helper: SMAppService.register OK (status \(describe(service.status)))")
        } catch {
            Log("Helper: SMAppService.register failed (status \(describe(service.status))): \(error.localizedDescription)")
        }
        if service.status == .requiresApproval {
            Log("Helper: daemon requires approval — opening System Settings")
            DispatchQueue.main.async { SMAppService.openSystemSettingsLoginItems() }
        }
    }

    private func describe(_ status: SMAppService.Status) -> String {
        switch status {
        case .notRegistered:    return "notRegistered"
        case .enabled:          return "enabled"
        case .requiresApproval: return "requiresApproval"
        case .notFound:         return "notFound"
        @unknown default:       return "unknown(\(status.rawValue))"
        }
    }

    // MARK: Connection

    /// Returns a live connection, recreating it (and re-exporting the reverse
    /// client object) if the previous one was invalidated.
    func connection() -> NSXPCConnection? {
        lock.lock(); defer { lock.unlock() }
        if xpcConnection == nil {
            let c = NSXPCConnection(machServiceName: HelperConstants.machServiceName,
                                    options: .privileged)
            c.remoteObjectInterface = NSXPCInterface(with: VPNHelperProtocol.self)
            c.exportedInterface = NSXPCInterface(with: VPNClientProtocol.self)
            c.exportedObject = client
            c.invalidationHandler = { [weak self] in
                guard let self = self else { return }
                self.lock.lock()
                self.xpcConnection?.invalidationHandler = nil
                self.xpcConnection = nil
                self.lock.unlock()
                Log("Helper: XPC connection invalidated")
            }
            c.resume()
            xpcConnection = c
        }
        return xpcConnection
    }

    private func proxy(_ errorHandler: @escaping (Error) -> Void) -> VPNHelperProtocol? {
        connection()?.remoteObjectProxyWithErrorHandler(errorHandler) as? VPNHelperProtocol
    }

    // MARK: Operations

    func subscribe() {
        proxy({ error in Log("Helper: subscribe error: \(error.localizedDescription)") })?
            .subscribe { ok in Log("Helper: subscribed \(ok)") }
    }

    func start(profile: VPNProfile, secrets: VPNSecrets) async -> (Int, String) {
        guard let configData = try? JSONEncoder().encode(profile),
              let secretData = try? JSONEncoder().encode(secrets) else {
            return (-3, "Failed to encode configuration")
        }
        var result = await invokeStart(configData, secretData)
        if isConnectionError(result) {
            Log("Helper: start failed (\(result.1)) — refreshing daemon and retrying")
            await recoverConnection()
            result = await invokeStart(configData, secretData)
        }
        return result
    }

    func stop(profileID: String) async -> (Int, String) {
        var result = await invokeStop(profileID)
        if isConnectionError(result) {
            Log("Helper: stop failed (\(result.1)) — refreshing daemon and retrying")
            await recoverConnection()
            result = await invokeStop(profileID)
        }
        return result
    }

    func stopAll() async -> (Int, String) {
        var result = await invokeStopAll()
        if isConnectionError(result) {
            Log("Helper: stopAll failed (\(result.1)) — refreshing daemon and retrying")
            await recoverConnection()
            result = await invokeStopAll()
        }
        return result
    }

    /// Emergency stop: the helper SIGKILLs every charon/dnsmasq (even strays it
    /// lost track of), restores networking and drops kill-switch firewall rules.
    func forceKillAll() async -> (Int, String) {
        await invokeReply { svc, reply in svc.forceKillAll(reply: reply) }
    }

    /// Push the profile list order (top = highest priority) so the helper can
    /// re-resolve route/DNS/dnsmasq conflicts among the active tunnels.
    func setPriorities(profileIDs: [String]) async {
        await withCheckedContinuation { (cont: CheckedContinuation<Void, Never>) in
            let once = ResumeOnce(cont)
            guard let svc = proxy({ _ in once.resume() }) else { once.resume(); return }
            svc.setPriorities(profileIDs: profileIDs) { once.resume() }
        }
    }

    /// Ask the running helper to exit cleanly (it SIGTERMs itself, tearing down
    /// any tunnels). launchd will only relaunch it on the next XPC connection, so
    /// this effectively stops the daemon until SwanDesk next needs it.
    func shutdownHelper() async {
        _ = await stopAll()   // restore networking before the process goes away
        await withCheckedContinuation { (cont: CheckedContinuation<Void, Never>) in
            let once = ResumeOnce(cont)
            guard let svc = proxy({ _ in once.resume() }) else { once.resume(); return }
            DispatchQueue.global().asyncAfter(deadline: .now() + 4) { once.resume() }
            svc.terminate { once.resume() }
        }
        clearConnection()
    }

    /// The one wrapper for every (Int, String)-reply XPC call: builds the proxy
    /// (its error handler resolves the continuation) and forwards the reply.
    private func invokeReply(_ call: @escaping (VPNHelperProtocol, @escaping (Int, String) -> Void) -> Void) async -> (Int, String) {
        await withCheckedContinuation { cont in
            guard let svc = proxy({ error in cont.resume(returning: (-1, "XPC error: \(error.localizedDescription)")) }) else {
                cont.resume(returning: (-1, "XPC error: no connection")); return
            }
            call(svc) { status, message in cont.resume(returning: (status, message)) }
        }
    }

    private func invokeStart(_ configData: Data, _ secretData: Data) async -> (Int, String) {
        await invokeReply { svc, reply in svc.start(configData: configData, secrets: secretData, reply: reply) }
    }

    private func invokeStop(_ profileID: String) async -> (Int, String) {
        await invokeReply { svc, reply in svc.stop(profileID: profileID, reply: reply) }
    }

    private func invokeStopAll() async -> (Int, String) {
        await invokeReply { svc, reply in svc.stopAll(reply: reply) }
    }

    private func isConnectionError(_ result: (Int, String)) -> Bool {
        result.0 == -1 && result.1.hasPrefix("XPC error")
    }

    /// Self-heal a dead/stale daemon: refresh the registration (so launchd points
    /// at the current binary and relaunches it on demand), drop the cached
    /// connection, and re-arm the reverse-XPC subscription on the new one.
    private func recoverConnection() async {
        let service = SMAppService.daemon(plistName: "SwanDeskHelperDaemon.plist")
        registerDaemon(service)
        clearConnection()
        try? await Task.sleep(nanoseconds: 400_000_000)
        subscribe()
    }

    /// Drop the cached connection so the next call recreates it. Synchronous, so
    /// it is safe to call from an async context (NSLock.lock() is not).
    private func clearConnection() {
        lock.lock(); xpcConnection = nil; lock.unlock()
    }

    func fetchStatus() async -> VPNRuntimeStatus {
        await withCheckedContinuation { cont in
            let svc = proxy({ _ in cont.resume(returning: VPNRuntimeStatus()) })
            svc?.status { data in
                cont.resume(returning: (try? JSONDecoder().decode(VPNRuntimeStatus.self, from: data)) ?? VPNRuntimeStatus())
            }
        }
    }
}

/// Resumes a continuation at most once, from whichever of several racing
/// callbacks (XPC reply, XPC error handler, timeout) fires first. Without this a
/// double-resume would crash and a never-fired path would leak the continuation.
private final class ResumeOnce<T> {
    private let lock = NSLock()
    private var done = false
    private let cont: CheckedContinuation<T, Never>
    init(_ cont: CheckedContinuation<T, Never>) { self.cont = cont }
    func resume(_ value: T) {
        lock.lock(); defer { lock.unlock() }
        if done { return }
        done = true
        cont.resume(returning: value)
    }
}

extension ResumeOnce where T == Void {
    func resume() { resume(()) }
}

/// Lightweight logger used across the app (mirrors AVC's `Log`).
func Log(_ message: String) {
    NSLog("SwanDesk: \(message)")
}
