import Foundation
import NucleicPowerProtocol

private final class FullWakeCoordinator: @unchecked Sendable {
    static let shared = FullWakeCoordinator()

    private let queue = DispatchQueue(label: "xyz.blakeslee.nucleic.power-helper")
    private var activeLeases: Set<UUID> = []
    private var originalSleepDisabled = false
    private let recoveryURL = URL(
        fileURLWithPath: "/var/run/\(NucleicPowerHelperConstants.serviceName).state")

    private init() {
        queue.sync {
            // A prior helper crash may have left the global flag asserted. Restore the value
            // captured before Nucleic changed it; a live client will acquire a fresh lease.
            if let stored = try? String(contentsOf: recoveryURL, encoding: .utf8) {
                _ = setSleepDisabled(
                    stored.trimmingCharacters(in: .whitespacesAndNewlines) == "1")
                try? FileManager.default.removeItem(at: recoveryURL)
            }
        }
    }

    func setActive(
        _ active: Bool, for lease: UUID,
        reply: @escaping @Sendable (NSError?) -> Void
    ) {
        queue.async {
            if active {
                guard !self.activeLeases.contains(lease) else {
                    reply(nil)
                    return
                }
                if self.activeLeases.isEmpty {
                    self.originalSleepDisabled = self.readSleepDisabled()
                    do {
                        try (self.originalSleepDisabled ? "1" : "0").write(
                            to: self.recoveryURL, atomically: true, encoding: .utf8)
                    } catch {
                        reply(error as NSError)
                        return
                    }
                    if let error = self.setSleepDisabled(true) {
                        try? FileManager.default.removeItem(at: self.recoveryURL)
                        reply(error)
                        return
                    }
                }
                self.activeLeases.insert(lease)
                reply(nil)
                return
            }

            self.release(lease)
            reply(nil)
        }
    }

    func connectionInvalidated(_ lease: UUID) {
        queue.async { self.release(lease) }
    }

    private func release(_ lease: UUID) {
        guard activeLeases.remove(lease) != nil, activeLeases.isEmpty else { return }
        _ = setSleepDisabled(originalSleepDisabled)
        try? FileManager.default.removeItem(at: recoveryURL)
    }

    private func readSleepDisabled() -> Bool {
        let result = runPMSet(["-g"])
        guard result.status == 0 else { return false }
        return result.output.split(separator: "\n").contains { line in
            let fields = line.split(whereSeparator: { $0 == " " || $0 == "\t" })
            return fields.count >= 2
                && fields[0] == "SleepDisabled"
                && fields[1] == "1"
        }
    }

    private func setSleepDisabled(_ disabled: Bool) -> NSError? {
        let result = runPMSet(["-a", "disablesleep", disabled ? "1" : "0"])
        guard result.status != 0 else { return nil }
        return NSError(
            domain: NucleicPowerHelperConstants.serviceName,
            code: Int(result.status),
            userInfo: [NSLocalizedDescriptionKey:
                result.output.isEmpty ? "pmset failed" : result.output])
    }

    private func runPMSet(_ arguments: [String]) -> (status: Int32, output: String) {
        let process = Process()
        process.executableURL = URL(fileURLWithPath: "/usr/bin/pmset")
        process.arguments = arguments
        let pipe = Pipe()
        process.standardOutput = pipe
        process.standardError = pipe
        do {
            try process.run()
            process.waitUntilExit()
        } catch {
            return (-1, error.localizedDescription)
        }
        let data = pipe.fileHandleForReading.readDataToEndOfFile()
        return (
            process.terminationStatus,
            String(data: data, encoding: .utf8)?
                .trimmingCharacters(in: .whitespacesAndNewlines) ?? "")
    }
}

private final class PowerHelperSession:
    NSObject, NucleicPowerHelperProtocol, @unchecked Sendable
{
    let lease = UUID()
    private let lock = NSLock()
    private var invalidated = false

    func setFullWakeActive(
        _ active: Bool, withReply reply: @escaping @Sendable (NSError?) -> Void
    ) {
        lock.lock()
        let isInvalidated = invalidated
        lock.unlock()
        guard !isInvalidated else {
            reply(NSError(
                domain: NucleicPowerHelperConstants.serviceName,
                code: 1,
                userInfo: [NSLocalizedDescriptionKey: "XPC connection is closed"]))
            return
        }
        FullWakeCoordinator.shared.setActive(active, for: lease, reply: reply)
    }

    func invalidate() {
        lock.lock()
        defer { lock.unlock() }
        guard !invalidated else { return }
        invalidated = true
        FullWakeCoordinator.shared.connectionInvalidated(lease)
    }
}

private final class PowerHelper: NSObject, NSXPCListenerDelegate {
    private let listener = NSXPCListener(
        machServiceName: NucleicPowerHelperConstants.serviceName)

    override init() {
        super.init()
        listener.delegate = self
    }

    func run() {
        _ = FullWakeCoordinator.shared
        listener.resume()
        RunLoop.current.run()
    }

    func listener(
        _ listener: NSXPCListener,
        shouldAcceptNewConnection connection: NSXPCConnection
    ) -> Bool {
        let session = PowerHelperSession()
        connection.exportedInterface = NSXPCInterface(with: NucleicPowerHelperProtocol.self)
        connection.exportedObject = session
        connection.invalidationHandler = { session.invalidate() }
        connection.interruptionHandler = { session.invalidate() }
        connection.resume()
        return true
    }
}

PowerHelper().run()
