import Foundation
import NucleicProtocol

// Compiled ONLY on Windows (this directory is excluded elsewhere — Package.swift).

/// The hostd side of the broker pipe (docs/WINDOWS_PORT.md §2.3, §3.3): spawns
/// `nucleic-brokerd.exe` as a supervised child, speaks NDJSON JSON-RPC 2.0 over its stdio
/// via the shared ``JSONRPCConnection``, routes its notifications (process stdio, pull
/// progress, session death) to their consumers, and restarts it with backoff when it dies.
/// wslc state is service-backed, so a restarted broker re-opens the same named session and
/// re-enumerates containers; after a successful reattach `onReattach` fires so the owner can
/// run `ContainerManager.reconcile(activeSessions:)` against the fresh listing.
///
/// The wire contract is pinned from the C# side by `windows/NucleicBroker.Tests` — every
/// method/key used here appears in those tests.
public actor WslcBrokerClient {
    public struct BrokerLost: Error, CustomStringConvertible {
        public let description = "the wslc broker process is not available"
    }

    /// Per-process notification sinks, registered by ``WslcProcessHandle``.
    struct ProcSinks {
        let stdout: @Sendable (Data) -> Void
        let stderr: @Sendable (Data) -> Void
        let exit: @Sendable (Int32) -> Void
    }

    /// A process event that arrived before its ``WslcProcessHandle`` existed.
    ///
    /// This is not a theoretical window. The broker enqueues the `proc.exec` response *before*
    /// starting the process precisely so the response cannot be overtaken — but a client still
    /// cannot register on time, because resolving the request's continuation only SCHEDULES it,
    /// while the notification drain is a separate task already reading the next line. A command
    /// that finishes immediately therefore lands `proc.stdout` and `proc.exit` before
    /// ``register(procId:sinks:)`` runs, and the old code silently dropped both — so `wait()`
    /// never returned and the caller hung forever.
    ///
    /// `runCapturing` is the most exposed path, since every short probe it runs (cgroup reads,
    /// shim re-seeding) is exactly the fast case. Observed on hardware through the C# spike, which
    /// hung twice on it (docs/WINDOWS_PORT.md §13.3).
    private enum PendingProcEvent {
        case stdout(Data)
        case stderr(Data)
        case exit(Int32)
    }

    /// Buffered events per not-yet-registered procId, drained by ``register(procId:sinks:)``.
    private var pendingProcEvents: [Int64: [PendingProcEvent]] = [:]

    /// The window is a continuation hop, so a handful of events is the realistic maximum; the cap
    /// only exists so a procId that never registers (a broker bug) cannot grow without bound.
    private static let maxPendingEventsPerProc = 512

    private func buffer(_ event: PendingProcEvent, for procId: Int64) {
        var events = pendingProcEvents[procId] ?? []
        let isExit: Bool = if case .exit = event { true } else { false }
        // Always keep the exit — losing it is the hang this mechanism exists to prevent, whereas
        // dropping overflow output only truncates a transcript.
        guard isExit || events.count < Self.maxPendingEventsPerProc else { return }
        events.append(event)
        pendingProcEvents[procId] = events
    }

    private let brokerdPath: String
    private var child: ChildProcess?
    private var connection: JSONRPCConnection?
    private var drainTask: Task<Void, Never>?
    private var superviseTask: Task<Void, Never>?
    private var procSinks: [Int64: ProcSinks] = [:]
    private var restartBackoffNanos: UInt64 = 500_000_000
    private var shutdownRequested = false

    /// Capabilities from the broker's `hello` (empty until the first successful start).
    public private(set) var capabilities: Set<String> = []
    public private(set) var wslcVersion: String?

    /// Fired after a broker crash+restart once the new broker answered `hello` — the owner
    /// re-ensures the session and reconciles containers.
    private let onReattach: @Sendable () async -> Void
    /// Forwarded `session.down {reason}` notifications (the wslc VM died underneath us).
    private let onSessionDown: @Sendable (String) async -> Void
    /// Forwarded `image.pullProgress` notifications, folded into the download-progress UI.
    private let onPullProgress: @Sendable (_ ref: String, _ status: String, _ current: Int64, _ total: Int64) async -> Void

    public init(
        brokerdPath: String? = nil,
        onReattach: @Sendable @escaping () async -> Void = {},
        onSessionDown: @Sendable @escaping (String) async -> Void = { _ in },
        onPullProgress: @Sendable @escaping (String, String, Int64, Int64) async -> Void = { _, _, _, _ in }
    ) {
        // Default: the exe ships beside nucleic-hostd in the MSIX payload (§2.4).
        self.brokerdPath = brokerdPath
            ?? ProcessInfo.processInfo.environment["NUCLEIC_BROKERD_PATH"]
            ?? URL(fileURLWithPath: CommandLine.arguments[0])
                .deletingLastPathComponent().appendingPathComponent("nucleic-brokerd.exe").path
        self.onReattach = onReattach
        self.onSessionDown = onSessionDown
        self.onPullProgress = onPullProgress
    }

    // MARK: - Lifecycle

    /// Spawn the broker (if not already running) and complete the `hello` exchange.
    public func start() async throws {
        guard connection == nil else { return }
        shutdownRequested = false
        let child = try ChildProcess(spec: ProcessSpec(
            executable: brokerdPath, args: [],
            cwd: FileManager.default.temporaryDirectory.path))
        let connection = JSONRPCConnection(handle: child, includeVersionHeader: true)
        self.child = child
        self.connection = connection
        // The broker never sends client-bound requests; answer any with an error.
        await connection.start { request in
            try? await connection.replyError(
                to: request.id, code: -32601, message: "hostd accepts no requests")
        }
        drainTask = Task { [weak self] in
            for await note in connection.notifications {
                await self?.route(note)
            }
        }
        superviseTask = Task { [weak self] in
            let code = await child.wait()
            await self?.brokerDied(exitCode: code)
        }

        let hello = try await connection.request("hello")
        capabilities = Set((hello["capabilities"]?.arrayValue ?? []).compactMap(\.stringValue))
        wslcVersion = hello["wslcVersion"]?.stringValue
        restartBackoffNanos = 500_000_000  // healthy start → reset the backoff ladder
    }

    public func shutdown() async {
        shutdownRequested = true
        superviseTask?.cancel()
        drainTask?.cancel()
        await connection?.close()
        await child?.terminate()
        connection = nil
        child = nil
    }

    private func brokerDied(exitCode: Int32) async {
        await connection?.close()
        connection = nil
        child = nil
        drainTask?.cancel()
        // Every registered process is dead with the broker: fail their streams so backends
        // hit the same recovery path a forced stream close serves (docs/WINDOWS_PORT.md §2.3).
        for (_, sinks) in procSinks { sinks.exit(-1) }
        procSinks.removeAll()
        // Events buffered for handles that never registered belong to the dead broker's procId
        // space, which the next broker reuses from 1 — keeping them would misdeliver.
        pendingProcEvents.removeAll()
        // Exponential-backoff restart loop (§2.3); a successful start() resets the ladder.
        while !shutdownRequested {
            let backoff = restartBackoffNanos
            restartBackoffNanos = min(restartBackoffNanos * 2, 30_000_000_000)
            try? await Task.sleep(nanoseconds: backoff)
            guard !shutdownRequested else { return }
            do {
                try await start()
                await onReattach()
                return
            } catch {
                continue
            }
        }
    }

    // MARK: - Requests

    /// One RPC round-trip; `BrokerLost` when the broker is down (a distinguishable error the
    /// backends map to the stream-recovery path).
    public func call(_ method: String, _ params: JSONValue? = nil) async throws -> JSONValue {
        guard let connection else { throw BrokerLost() }
        do {
            return try await connection.request(method, params: params)
        } catch JSONRPCConnection.RPCError.connectionClosed {
            throw BrokerLost()
        }
    }

    // MARK: - Process event routing

    func register(procId: Int64, sinks: ProcSinks) {
        procSinks[procId] = sinks
        // Deliver anything that arrived before this handle existed — see `pendingProcEvents`.
        guard let buffered = pendingProcEvents.removeValue(forKey: procId) else { return }
        for event in buffered {
            switch event {
            case .stdout(let data): sinks.stdout(data)
            case .stderr(let data): sinks.stderr(data)
            case .exit(let code):
                // The process is already over. Drop the sink as `route` would have.
                procSinks[procId] = nil
                sinks.exit(code)
                return
            }
        }
    }

    func unregister(procId: Int64) {
        procSinks[procId] = nil
        pendingProcEvents[procId] = nil
    }

    private func route(_ note: JSONRPCConnection.Notification) async {
        let params = note.params
        switch note.method {
        case "proc.stdout", "proc.stderr":
            guard let rawId = params["procId"]?.intValue,
                  let b64 = params["b64"]?.stringValue,
                  let bytes = Data(base64Encoded: b64)
            else { return }
            let procId = Int64(rawId)
            guard let sinks = procSinks[procId] else {
                buffer(note.method == "proc.stderr" ? .stderr(bytes) : .stdout(bytes), for: procId)
                return
            }
            (note.method == "proc.stderr" ? sinks.stderr : sinks.stdout)(bytes)
        case "proc.exit":
            guard let rawId = params["procId"]?.intValue else { return }
            let procId = Int64(rawId)
            let code = Int32(params["code"]?.intValue ?? -1)
            guard let sinks = procSinks.removeValue(forKey: procId) else {
                buffer(.exit(code), for: procId)
                return
            }
            sinks.exit(code)
        case "session.down":
            await onSessionDown(params["reason"]?.stringValue ?? "unknown")
        case "image.pullProgress":
            await onPullProgress(
                params["ref"]?.stringValue ?? "",
                params["status"]?.stringValue ?? "",
                Int64(params["current"]?.intValue ?? 0),
                Int64(params["total"]?.intValue ?? 0))
        default:
            break  // components.installProgress etc. — routed once onboarding consumes them
        }
    }
}
