import Foundation

/// Thresholds for ``ProcessStallMonitor``. Two independent conditions must both hold before the user
/// is bothered: the process has produced no output for at least `idleThreshold`, **and** (when the
/// pid is host-samplable) it has been *resource-quiet* — burning less than `cpuFractionThreshold` of
/// one core **and** moving less than `ioByteThreshold` of disk I/O per sample — continuously for at
/// least `lowCPUSustainNanos`. The idle threshold is the "no output" timer; the resource-sustain
/// window is a confirmation guard so a command that goes quiet on stdout but is still busy — whether
/// CPU-bound (a compile) or I/O-bound (a registry push streaming blobs at near-zero CPU) — never
/// trips it. Defaults aim at 300s of genuine silence with ≥60s of near-zero CPU *and* I/O.
public struct ProcessStallConfig: Sendable {
    public var sampleIntervalNanos: UInt64
    public var idleThresholdNanos: UInt64
    /// Fraction of a single core (0…1). 0.02 ≈ 2% — high enough to ignore a mostly-idle process
    /// that still ticks (a poll loop), low enough that a genuinely working command clears it.
    public var cpuFractionThreshold: Double
    /// Disk-I/O bytes (read + written) per sample interval below which the process tree counts as
    /// I/O-quiet. A command moving real data at near-zero CPU — a `cctl images push` reading blobs
    /// from the content store, a slow copy — clears this and is spared, which CPU alone cannot do.
    /// When the pid isn't I/O-samplable (guest/VM handles) this signal is simply absent.
    public var ioByteThreshold: UInt64
    /// The process must stay resource-quiet (under both `cpuFractionThreshold` and `ioByteThreshold`)
    /// continuously for at least this long before an alert fires. A single busy sample — CPU *or*
    /// I/O — resets the window. When the pid isn't host-samplable (guest/VM handles) this is trivially
    /// satisfied and detection falls back to output-idle alone.
    public var lowCPUSustainNanos: UInt64

    public init(
        sampleIntervalNanos: UInt64 = 5 * 1_000_000_000,
        idleThresholdNanos: UInt64 = 300 * 1_000_000_000,
        cpuFractionThreshold: Double = 0.02,
        ioByteThreshold: UInt64 = 64 * 1024,
        lowCPUSustainNanos: UInt64 = 60 * 1_000_000_000
    ) {
        self.sampleIntervalNanos = sampleIntervalNanos
        self.idleThresholdNanos = idleThresholdNanos
        self.cpuFractionThreshold = cpuFractionThreshold
        self.ioByteThreshold = ioByteThreshold
        self.lowCPUSustainNanos = lowCPUSustainNanos
    }

    public static let `default` = ProcessStallConfig()
}

/// Why a raised stall alert cleared on its own (as opposed to the user killing/dismissing it).
public enum StallResolution: String, Sendable {
    /// Output (or CPU) picked back up — it wasn't actually wedged.
    case resumed
    /// The process exited before the user acted.
    case exited
}

/// Watches one already-launched process for a hang and reports it, without ever touching it: the
/// caller decides what to do (surface an alert, offer a kill). It samples on a fixed interval until
/// the process exits, combining signals — time since last output (``ProcessHandle/lastActivityNanos``)
/// and, when the pid is host-samplable, CPU **and** disk-I/O bytes since the last sample. `onStall` fires at most
/// once per stalled stretch; if the process recovers it fires `onResolve(.resumed)` and re-arms, and
/// on exit it fires `onResolve(.exited)` — but only if it had alerted, so a healthy command that
/// simply finishes reports nothing.
///
/// `@unchecked Sendable`: all mutable state is guarded by `lock`; the two driver tasks (sampling loop
/// + exit watcher) coordinate their one-shot conclusion through it.
public final class ProcessStallMonitor: @unchecked Sendable {
    private let handle: any ProcessHandle
    private let config: ProcessStallConfig
    /// Cumulative CPU-nanos sampler for the process tree, or `nil` when the pid isn't host-samplable
    /// (guest/VM handles) — then detection uses output-idle alone.
    private let cpuSampler: (@Sendable () -> UInt64)?
    /// Cumulative disk-I/O-bytes sampler for the process tree, or `nil` when the pid isn't
    /// host-samplable — the second resource signal alongside CPU, so an I/O-bound-but-quiet command
    /// (a registry push at near-zero CPU) isn't mistaken for wedged.
    private let ioSampler: (@Sendable () -> UInt64)?
    /// Whether time-since-output counts toward the stall verdict. `false` when the command's own
    /// terminal pipeline stage withholds output by construction — a non-follow `tail`/`head` buffers
    /// its whole input until EOF, so the captured stream is silent no matter how alive the command
    /// is. Then the verdict rests on the resource signals (CPU/I/O) alone. Ignored — output-idle is
    /// always honored — when there is no resource sampler at all, or the check would have no signal.
    private let outputIdleCounts: Bool
    private let onStall: @Sendable (Int) async -> Void
    private let onResolve: @Sendable (StallResolution) async -> Void

    private let lock = NSLock()
    private var loopTask: Task<Void, Never>?
    private var exitTask: Task<Void, Never>?
    private var finished = false
    private var alerted = false

    public init(
        handle: any ProcessHandle,
        config: ProcessStallConfig = .default,
        cpuSampler: (@Sendable () -> UInt64)? = nil,
        ioSampler: (@Sendable () -> UInt64)? = nil,
        outputIdleCounts: Bool = true,
        onStall: @escaping @Sendable (_ idleSeconds: Int) async -> Void,
        onResolve: @escaping @Sendable (StallResolution) async -> Void
    ) {
        self.handle = handle
        self.config = config
        self.cpuSampler = cpuSampler
        self.ioSampler = ioSampler
        self.outputIdleCounts = outputIdleCounts
        self.onStall = onStall
        self.onResolve = onResolve
    }

    /// Begin sampling. Idempotent-safe to call once. The monitor stops itself when the process exits.
    public func start() {
        loopTask = Task { [weak self] in await self?.runLoop() }
        exitTask = Task { [weak self] in
            guard let self else { return }
            _ = await self.handle.wait()
            await self.conclude(.exited)
        }
    }

    /// Whether the process currently looks wedged: `onStall` has fired for this stretch and nothing has
    /// recovered it yet. Callers that hold a wall-clock ceiling consult this so they only reap a command
    /// that is BOTH past its age AND genuinely quiet — a command that is merely slow (still burning CPU,
    /// still moving I/O, or still emitting output) is never stalled here and must be allowed to continue.
    public var isStalled: Bool { lock.withLock { alerted } }

    /// Stop watching now (the caller resolved the alert its own way — killed or dismissed it). Any
    /// pending self-resolution is suppressed.
    public func cancel() {
        lock.lock()
        finished = true
        let loop = loopTask
        let exit = exitTask
        lock.unlock()
        loop?.cancel()
        exit?.cancel()
    }

    private func runLoop() async {
        // Output-idle is our only signal when nothing is resource-samplable, so honor it then even if
        // the caller asked to discount it (a `tail`-wrapped container/VM command still needs *some*
        // signal). With a sampler present, `outputIdleCounts == false` means the captured stream is
        // silent by construction and the verdict rests on CPU/I/O alone.
        let hasResourceSampler = cpuSampler != nil || ioSampler != nil
        let countOutputIdle = outputIdleCounts || !hasResourceSampler
        // When output-idle is discounted, the resource-quiet window must itself span the idle
        // threshold before we conclude "wedged" — otherwise the shorter sustain window would fire far
        // too eagerly on a command that is simply slow and quiet. Otherwise the normal sustain governs.
        let sustainNanos = countOutputIdle
            ? config.lowCPUSustainNanos
            : max(config.lowCPUSustainNanos, config.idleThresholdNanos)
        // How many consecutive resource-quiet samples make up the sustain window. At least one, so a
        // sub-interval sustain setting still requires a real sample.
        let quietSamplesNeeded = max(1, Int(sustainNanos / max(1, config.sampleIntervalNanos)))
        var quietStreak = 0
        var lastCPU: UInt64? = cpuSampler?()
        var lastIO: UInt64? = ioSampler?()
        while true {
            try? await Task.sleep(nanoseconds: config.sampleIntervalNanos)
            if lock.withLock({ finished }) || Task.isCancelled { return }

            let idleNanos = DispatchTime.now().uptimeNanoseconds &- handle.lastActivityNanos
            // Busy if EITHER resource signal is above its threshold this sample: CPU-bound work and
            // I/O-bound work both count as alive. Absent samplers contribute nothing (stay quiet).
            var busy = false
            if let cpuSampler {
                let now = cpuSampler()
                let delta = now >= (lastCPU ?? now) ? now &- (lastCPU ?? now) : 0
                lastCPU = now
                let fraction = Double(delta) / Double(config.sampleIntervalNanos)
                if fraction >= config.cpuFractionThreshold { busy = true }
            }
            if let ioSampler {
                let now = ioSampler()
                let delta = now >= (lastIO ?? now) ? now &- (lastIO ?? now) : 0
                lastIO = now
                if delta >= config.ioByteThreshold { busy = true }
            }
            // A single busy sample resets the resource-sustain window, so we only alert once the tree
            // has been quiet continuously long enough. Idle is monotonic in real output, so its own
            // threshold needs no such streak.
            quietStreak = busy ? 0 : quietStreak + 1

            let idleStalled = countOutputIdle ? idleNanos >= config.idleThresholdNanos : true
            let stalled = idleStalled && quietStreak >= quietSamplesNeeded
            if stalled {
                let fire = lock.withLock { () -> Bool in
                    guard !finished, !alerted else { return false }
                    alerted = true
                    return true
                }
                if fire { await onStall(Int(idleNanos / 1_000_000_000)) }
            } else {
                let recovered = lock.withLock { () -> Bool in
                    guard alerted else { return false }
                    alerted = false
                    return true
                }
                if recovered { await onResolve(.resumed) }
            }
        }
    }

    /// The process exited: stop the loop and, if an alert was live, report it cleared.
    private func conclude(_ resolution: StallResolution) async {
        let wasAlerted = lock.withLock { () -> Bool? in
            if finished { return nil }
            finished = true
            let a = alerted
            alerted = false
            return a
        }
        guard let wasAlerted else { return }
        loopTask?.cancel()
        if wasAlerted { await onResolve(resolution) }
    }
}
