//
//  ExecutionService.swift
//  Zyquo Agent
//
//  Author: Simon-Pierre Boucher
//  Mail: contact@spboucher.ai
//
//  The single place Zyquo Agent spawns processes. Runs `/bin/bash -lc` and
//  `/usr/bin/osascript` with the task workspace as cwd, streams stdout/stderr
//  line by line to the UI as they arrive, enforces per-command timeouts
//  (SIGTERM, then SIGKILL after a grace period), honors Swift Task
//  cancellation, and caps captured output while the live stream still sees
//  everything. Every caller reaches this actor only *after* the PolicyEngine
//  has cleared the action — ExecutionService itself performs no policy checks
//  and must never be invoked with an uncleared payload.
//

import Foundation

/// Tunables for process execution; a task can override the timeout per call.
struct ExecutionConfiguration: Sendable {
    /// Default per-command timeout when the caller does not specify one.
    var defaultTimeout: TimeInterval = 120
    /// Seconds between SIGTERM and SIGKILL when a process must die.
    var killGracePeriod: TimeInterval = 3
    /// Maximum characters kept in the captured stdout/stderr/combined strings.
    /// The streaming callback is NOT capped — the UI terminal sees everything.
    var capturedOutputLimit: Int = 200_000

    static let `default` = ExecutionConfiguration()
}

/// Why a process stopped. Timeout and cancellation are reported distinctly
/// from a normal exit so tools can mark results accurately.
enum ExecutionTermination: Sendable, Equatable {
    /// The process exited (or was signalled) on its own; carries the exit code.
    case exited(Int32)
    /// The per-command timeout expired and the service killed the process.
    case timedOut
    /// The surrounding Swift Task was cancelled and the service killed the process.
    case cancelled
}

/// Everything a caller learns from one process run.
struct ExecutionResult: Sendable {
    var termination: ExecutionTermination
    /// Exit code when available (also populated for killed processes when the
    /// OS reports one; nil only if the process never launched).
    var exitCode: Int32?
    /// Captured stdout, capped at `capturedOutputLimit`.
    var stdout: String
    /// Captured stderr, capped at `capturedOutputLimit`.
    var stderr: String
    /// stdout+stderr interleaved in arrival order, capped, with a truncation
    /// notice appended when the cap was hit.
    var combinedOutput: String
    /// True when any captured string was truncated at the cap.
    var wasTruncated: Bool
    /// Wall-clock seconds from launch to termination.
    var duration: TimeInterval

    var isTimeout: Bool { termination == .timedOut }
    var isCancelled: Bool { termination == .cancelled }
}

/// Thrown when the process cannot even be launched (bad executable, spawn
/// failure) — distinct from a process that ran and failed.
struct ExecutionLaunchError: Error, Sendable {
    var reason: String
}

/// Actor that owns all child-process execution. Serialization through the
/// actor is intentional for bookkeeping, but the actual waiting is fully
/// async — long commands do not block other actor work because the heavy
/// lifting happens off-actor in handlers and continuations.
actor ExecutionService {
    private let configuration: ExecutionConfiguration

    init(configuration: ExecutionConfiguration = .default) {
        self.configuration = configuration
    }

    // MARK: - Public entry points

    /// Runs a shell command via `/bin/bash -lc "<command>"` with `cwd` as the
    /// working directory. `onOutput` receives every stdout/stderr line live.
    func runBash(
        command: String,
        cwd: URL,
        timeout: TimeInterval? = nil,
        onOutput: @escaping @Sendable (ToolOutputChunk) -> Void
    ) async throws -> ExecutionResult {
        try await runProcess(
            executablePath: "/bin/bash",
            arguments: ["-lc", command],
            cwd: cwd,
            timeout: timeout,
            onOutput: onOutput
        )
    }

    /// Runs an inline AppleScript via `/usr/bin/osascript -e <line> [-e <line>…]`.
    func runOSAScript(
        lines: [String],
        cwd: URL,
        timeout: TimeInterval? = nil,
        onOutput: @escaping @Sendable (ToolOutputChunk) -> Void
    ) async throws -> ExecutionResult {
        var arguments: [String] = []
        for line in lines {
            arguments.append("-e")
            arguments.append(line)
        }
        return try await runProcess(
            executablePath: "/usr/bin/osascript",
            arguments: arguments,
            cwd: cwd,
            timeout: timeout,
            onOutput: onOutput
        )
    }

    /// Runs an AppleScript file via `/usr/bin/osascript <file>` (the path used
    /// for multi-line scripts, written under the workspace's `.zyquo/` dir).
    func runOSAScriptFile(
        scriptFile: URL,
        cwd: URL,
        timeout: TimeInterval? = nil,
        onOutput: @escaping @Sendable (ToolOutputChunk) -> Void
    ) async throws -> ExecutionResult {
        try await runProcess(
            executablePath: "/usr/bin/osascript",
            arguments: [scriptFile.path],
            cwd: cwd,
            timeout: timeout,
            onOutput: onOutput
        )
    }

    /// General process runner used by the specialized entry points above.
    /// Environment: inherits the user's environment, sets `PWD` to the cwd and
    /// adds `ZYQUO_AGENT=1` so scripts can detect they run under the agent.
    func runProcess(
        executablePath: String,
        arguments: [String],
        cwd: URL,
        timeout: TimeInterval? = nil,
        onOutput: @escaping @Sendable (ToolOutputChunk) -> Void
    ) async throws -> ExecutionResult {
        // A cancelled task must not spawn anything.
        if Task.isCancelled {
            return ExecutionResult(
                termination: .cancelled, exitCode: nil,
                stdout: "", stderr: "", combinedOutput: "",
                wasTruncated: false, duration: 0
            )
        }

        let supervisor = ProcessSupervisor(
            executablePath: executablePath,
            arguments: arguments,
            cwd: cwd,
            outputLimit: configuration.capturedOutputLimit,
            killGracePeriod: configuration.killGracePeriod,
            onOutput: onOutput
        )
        let effectiveTimeout = timeout ?? configuration.defaultTimeout
        return try await supervisor.run(timeout: effectiveTimeout)
    }
}

// MARK: - ProcessSupervisor

/// Owns exactly one child process for its lifetime: launch, line-streaming,
/// timeout escalation (SIGTERM → SIGKILL), cancellation, and reaping. All
/// mutable state is behind `lock`, so the class is safe to touch from the
/// pipe-reader queues, the termination handler, the timeout task, and the
/// cancellation handler concurrently.
private final class ProcessSupervisor: @unchecked Sendable {
    private let process = Process()
    private let stdoutPipe = Pipe()
    private let stderrPipe = Pipe()
    private let onOutput: @Sendable (ToolOutputChunk) -> Void
    private let outputLimit: Int
    private let killGracePeriod: TimeInterval

    /// Serial queue that funnels both pipes' data so combined output keeps the
    /// best-possible interleave order and capture buffers mutate safely.
    private let outputQueue = DispatchQueue(label: "com.zyquo.agent.execution.output")
    /// Signalled (once per stream) when a pipe reaches EOF.
    private let drainGroup = DispatchGroup()

    private let lock = NSLock()
    // Guarded by `lock`:
    private var overrideTermination: ExecutionTermination?
    private var continuation: CheckedContinuation<Void, Never>?
    private var terminated = false

    // Guarded by `outputQueue`:
    private var stdoutCapture = ""
    private var stderrCapture = ""
    private var combinedCapture = ""
    private var truncated = false
    private var stdoutLineBuffer = ""
    private var stderrLineBuffer = ""

    init(
        executablePath: String,
        arguments: [String],
        cwd: URL,
        outputLimit: Int,
        killGracePeriod: TimeInterval,
        onOutput: @escaping @Sendable (ToolOutputChunk) -> Void
    ) {
        self.onOutput = onOutput
        self.outputLimit = outputLimit
        self.killGracePeriod = killGracePeriod

        process.executableURL = URL(fileURLWithPath: executablePath)
        process.arguments = arguments
        process.currentDirectoryURL = cwd

        var environment = ProcessInfo.processInfo.environment
        environment["PWD"] = cwd.path
        environment["ZYQUO_AGENT"] = "1"
        process.environment = environment

        process.standardOutput = stdoutPipe
        process.standardError = stderrPipe
        process.standardInput = FileHandle.nullDevice
    }

    /// Launches the process and suspends until it terminates, times out, or
    /// the surrounding Task is cancelled. Never throws after a successful
    /// launch — abnormal endings are reported in `ExecutionResult.termination`.
    func run(timeout: TimeInterval) async throws -> ExecutionResult {
        installReader(for: stdoutPipe.fileHandleForReading, isStdout: true)
        installReader(for: stderrPipe.fileHandleForReading, isStdout: false)

        let start = Date()
        process.terminationHandler = { [weak self] _ in
            self?.finishWaiting()
        }

        do {
            try process.run()
        } catch {
            // Detach the readers we installed; the pipes never produced data.
            stdoutPipe.fileHandleForReading.readabilityHandler = nil
            stderrPipe.fileHandleForReading.readabilityHandler = nil
            drainGroup.leave()
            drainGroup.leave()
            throw ExecutionLaunchError(reason: "Could not launch \(process.executableURL?.path ?? "?"): \(error.localizedDescription)")
        }

        // Timeout watchdog: mark the run as timed out, then escalate
        // SIGTERM → (grace) → SIGKILL. The termination handler is the single
        // point that resumes the continuation, so there is no double-resume.
        let watchdog = Task { [weak self] in
            try? await Task.sleep(nanoseconds: UInt64(max(0, timeout) * 1_000_000_000))
            guard let self, !Task.isCancelled else { return }
            self.kill(as: .timedOut)
        }

        await withTaskCancellationHandler {
            await withCheckedContinuation { (cont: CheckedContinuation<Void, Never>) in
                let alreadyDone: Bool
                lock.lock()
                if terminated {
                    alreadyDone = true
                } else {
                    continuation = cont
                    alreadyDone = false
                }
                lock.unlock()
                if alreadyDone { cont.resume() }
            }
        } onCancel: {
            self.kill(as: .cancelled)
        }
        watchdog.cancel()

        // Wait (bounded) for the pipes to drain so trailing output is not
        // lost. The bound matters: a backgrounded grandchild can keep the
        // write end open forever even after bash itself has died. The wait
        // happens on a Dispatch thread — never blocks the cooperative pool.
        await withCheckedContinuation { (cont: CheckedContinuation<Void, Never>) in
            DispatchQueue.global(qos: .utility).async { [drainGroup] in
                _ = drainGroup.wait(timeout: .now() + 2.0)
                cont.resume()
            }
        }
        stdoutPipe.fileHandleForReading.readabilityHandler = nil
        stderrPipe.fileHandleForReading.readabilityHandler = nil

        let duration = Date().timeIntervalSince(start)

        let override = recordedOverride()

        let exitCode = process.processIdentifier > 0 ? process.terminationStatus : nil
        let termination = override ?? .exited(exitCode ?? -1)

        // Flush partial trailing lines and read the final buffers on the
        // output queue so we see everything the readers appended.
        return outputQueue.sync {
            flushLineBuffersLocked()
            var combined = combinedCapture
            if truncated {
                combined += "\n… [output truncated at \(outputLimit / 1000)KB — full stream was shown live]"
            }
            return ExecutionResult(
                termination: termination,
                exitCode: exitCode,
                stdout: stdoutCapture,
                stderr: stderrCapture,
                combinedOutput: combined,
                wasTruncated: truncated,
                duration: duration
            )
        }
    }

    /// Lock-guarded read of the recorded termination override (sync helper —
    /// NSLock must not be taken directly inside an async function).
    private func recordedOverride() -> ExecutionTermination? {
        lock.lock()
        defer { lock.unlock() }
        return overrideTermination
    }

    // MARK: Killing

    /// Records why the process is being killed, sends SIGTERM, and schedules a
    /// SIGKILL after the grace period. Safe to call from any thread; the first
    /// recorded reason wins.
    private func kill(as reason: ExecutionTermination) {
        lock.lock()
        if terminated {
            lock.unlock()
            return
        }
        if overrideTermination == nil { overrideTermination = reason }
        lock.unlock()

        let pid = process.processIdentifier
        guard pid > 0 else { return }
        Foundation.kill(pid, SIGTERM)

        let grace = killGracePeriod
        DispatchQueue.global().asyncAfter(deadline: .now() + grace) { [weak self] in
            guard let self else { return }
            self.lock.lock()
            let stillRunning = !self.terminated
            self.lock.unlock()
            if stillRunning {
                Foundation.kill(pid, SIGKILL)
            }
        }
    }

    /// Termination handler target — marks the process reaped and resumes the
    /// suspended `run(timeout:)`.
    private func finishWaiting() {
        lock.lock()
        terminated = true
        let cont = continuation
        continuation = nil
        lock.unlock()
        cont?.resume()
    }

    // MARK: Output plumbing

    /// Attaches a readability handler that forwards raw data to the serial
    /// output queue where it is captured (capped) and split into lines for
    /// the live stream. An empty read means EOF: detach and signal the drain.
    private func installReader(for handle: FileHandle, isStdout: Bool) {
        drainGroup.enter()
        handle.readabilityHandler = { [weak self] fileHandle in
            let data = fileHandle.availableData
            guard let self else { return }
            if data.isEmpty {
                fileHandle.readabilityHandler = nil
                self.outputQueue.async { self.drainGroup.leave() }
                return
            }
            self.outputQueue.async {
                self.consume(data: data, isStdout: isStdout)
            }
        }
    }

    /// Runs on `outputQueue` only. Captures (with cap) and emits whole lines.
    private func consume(data: Data, isStdout: Bool) {
        let text = String(decoding: data, as: UTF8.self)

        // Capture with cap. The live stream below is never capped.
        appendCapped(text, to: isStdout ? \.stdoutCapture : \.stderrCapture)
        appendCapped(text, to: \.combinedCapture)

        // Line-split for streaming; keep the unterminated tail buffered.
        if isStdout {
            stdoutLineBuffer += text
            emitCompleteLines(from: &stdoutLineBuffer, isStdout: true)
        } else {
            stderrLineBuffer += text
            emitCompleteLines(from: &stderrLineBuffer, isStdout: false)
        }
    }

    private func appendCapped(_ text: String, to keyPath: ReferenceWritableKeyPath<ProcessSupervisor, String>) {
        let current = self[keyPath: keyPath]
        guard current.count < outputLimit else {
            truncated = true
            return
        }
        let remaining = outputLimit - current.count
        if text.count <= remaining {
            self[keyPath: keyPath] = current + text
        } else {
            self[keyPath: keyPath] = current + String(text.prefix(remaining))
            truncated = true
        }
    }

    private func emitCompleteLines(from buffer: inout String, isStdout: Bool) {
        while let newlineRange = buffer.range(of: "\n") {
            let line = String(buffer[buffer.startIndex..<newlineRange.lowerBound])
            buffer.removeSubrange(buffer.startIndex..<newlineRange.upperBound)
            onOutput(isStdout ? .stdout(line) : .stderr(line))
        }
    }

    /// Runs on `outputQueue` only. Emits any unterminated trailing output.
    private func flushLineBuffersLocked() {
        if !stdoutLineBuffer.isEmpty {
            onOutput(.stdout(stdoutLineBuffer))
            stdoutLineBuffer = ""
        }
        if !stderrLineBuffer.isEmpty {
            onOutput(.stderr(stderrLineBuffer))
            stderrLineBuffer = ""
        }
    }
}