import Foundation

/// The live worker channel for a single Orchestra supervisor. One instance per supervisor session
/// (`AppStore` keys them by supervisor id, creating one on the supervisor's first spawn and tearing
/// it down when the supervisor's turn ends). It is the meeting point between the supervisor's
/// `nucleic_supervise` / `nucleic_reply_to_worker` calls and its workers' `nucleic_ask_supervisor`
/// calls: each worker pushes its questions and its terminal result here as ``OrchestraWorkerEvent``s,
/// and the supervisor drains them.
///
/// An `actor`, so the concurrent worker run tasks and the supervisor's tool handlers can all touch its
/// state at once without racing — and without hopping through the `@MainActor` `AppStore`. The
/// suspend/resume primitives mirror `AppStore.acquireOrchestraWorkerSlot`: a `withCheckedContinuation`
/// registered synchronously on the actor before the caller suspends, resumed later from another call.
actor OrchestraWorkerCoordinator {
    private enum WorkerState { case running, done }

    private struct WorkerHandle {
        let task: String
        /// The batch this worker was dispatched in, or nil when dispatched outside one.
        let batch: String?
        var state: WorkerState
        /// Set while the worker is parked in `nucleic_ask_supervisor` awaiting the supervisor's reply.
        var pendingQuestion: CheckedContinuation<OrchestraAskSupervisorResult, Never>?
    }

    /// Cumulative outcome tally for one batch. Kept separately from `handles` because `supervise`
    /// prunes terminal workers as it drains them — without this, a batch's counts would reset to zero
    /// the moment its results were read, and "is this batch complete?" could never be answered.
    private struct BatchTally {
        var completed = 0
        var failed = 0
    }

    private var handles: [SessionID: WorkerHandle] = [:]
    /// Per-batch tallies, plus `batchOrder` so progress is reported in dispatch order rather than a
    /// dictionary's arbitrary one (the supervisor reads these as a plan's phases).
    private var batchTallies: [String: BatchTally] = [:]
    private var batchOrder: [String] = []
    /// Events not yet drained by `nucleic_supervise`, in the order they occurred.
    private var pendingEvents: [OrchestraWorkerEvent] = []
    /// `nucleic_supervise` callers parked waiting for the next event. Usually at most one (the
    /// supervisor drives the loop serially), but a list keeps a second concurrent call from
    /// stranding the first.
    private var superviseWaiters: [CheckedContinuation<Void, Never>] = []
    /// Set once the supervisor's turn ends: no new work is accepted and parked callers are released.
    private var tornDown = false

    private static let unavailableMessage =
        "The supervisor's run has ended; proceed autonomously on your best judgment."

    // MARK: Registration & completion (called from the worker's run task)

    /// Record a freshly spawned worker as running, in `batch` when it belongs to one. Called before its
    /// run task starts the turn, so the supervisor's `nucleic_supervise` counts it toward
    /// `remaining_running` even while it is still queued for a concurrency slot — and so a batch is
    /// never reported "complete" merely because its workers haven't started yet.
    func register(workerID: SessionID, task: String, batch: String? = nil) {
        if let batch {
            if batchTallies[batch] == nil {
                batchTallies[batch] = BatchTally()
                batchOrder.append(batch)
            }
        }
        if let existing = handles[workerID] {
            // Re-registering an id keeps its original batch; only the state is refreshed.
            handles[workerID] = WorkerHandle(
                task: existing.task, batch: existing.batch ?? batch, state: .running,
                pendingQuestion: existing.pendingQuestion)
        } else {
            handles[workerID] = WorkerHandle(
                task: task, batch: batch, state: .running, pendingQuestion: nil)
        }
    }

    /// The worker finished successfully — surface its final output as a `completed` event.
    func complete(workerID: SessionID, model: String?, output: String) {
        finish(workerID: workerID) { task, batch in
            .completed(workerID: workerID, task: task, batch: batch, model: model, output: output)
        }
    }

    /// The worker errored or was interrupted — surface a `failed` event.
    func fail(workerID: SessionID, message: String) {
        finish(workerID: workerID) { task, batch in
            .failed(workerID: workerID, task: task, batch: batch, message: message)
        }
    }

    private func finish(
        workerID: SessionID, event: (String, String?) -> OrchestraWorkerEvent
    ) {
        guard !tornDown else { return }
        let task = handles[workerID]?.task ?? ""
        let batch = handles[workerID]?.batch
        // A finishing worker shouldn't still be parked in a question, but if it somehow is, release
        // its continuation defensively so nothing leaks before we mark it done.
        if let cont = handles[workerID]?.pendingQuestion {
            handles[workerID]?.pendingQuestion = nil
            cont.resume(returning: .unavailable(message: "The worker ended before a reply arrived."))
        }
        handles[workerID]?.state = .done
        let built = event(task, batch)
        // Tally against the batch before enqueueing, so a `supervise` woken by this event already sees
        // the updated per-batch progress (and, for the last worker of a batch, its completion).
        if let batch, built.isTerminal {
            var tally = batchTallies[batch] ?? BatchTally()
            switch built {
            case .completed: tally.completed += 1
            case .failed: tally.failed += 1
            case .question: break
            }
            if batchTallies[batch] == nil { batchOrder.append(batch) }
            batchTallies[batch] = tally
        }
        enqueue(built)
    }

    // MARK: Worker → supervisor (`nucleic_ask_supervisor`)

    /// Park the worker's question until the supervisor replies (`reply`) or its run ends (`teardown`
    /// or task cancellation). The question is surfaced to the supervisor as a `question` event.
    func ask(workerID: SessionID, question: String) async -> OrchestraAskSupervisorResult {
        if tornDown { return .unavailable(message: Self.unavailableMessage) }
        return await withTaskCancellationHandler {
            await withCheckedContinuation {
                (cont: CheckedContinuation<OrchestraAskSupervisorResult, Never>) in
                // Runs on the actor before this call suspends, so registration happens-before any
                // reply/teardown/cancel can race it.
                if tornDown {
                    cont.resume(returning: .unavailable(message: Self.unavailableMessage))
                    return
                }
                if handles[workerID] == nil {
                    handles[workerID] = WorkerHandle(
                        task: "", batch: nil, state: .running, pendingQuestion: nil)
                }
                handles[workerID]?.pendingQuestion = cont
                let task = handles[workerID]?.task ?? ""
                let batch = handles[workerID]?.batch
                enqueue(.question(workerID: workerID, task: task, batch: batch, question: question))
            }
        } onCancel: {
            // The worker's own run ended (its MCP handler task was cancelled): unwind the parked
            // question so the cancelled task can complete. Guarded against double-resume below.
            Task { await self.resolvePendingQuestionUnavailable(workerID) }
        }
    }

    private func resolvePendingQuestionUnavailable(_ workerID: SessionID) {
        guard let cont = handles[workerID]?.pendingQuestion else { return }
        handles[workerID]?.pendingQuestion = nil
        cont.resume(returning: .unavailable(message: Self.unavailableMessage))
    }

    // MARK: Supervisor → worker (`nucleic_reply_to_worker`)

    /// Deliver the supervisor's reply to a worker parked in `ask`. Returns `.unknown` when no worker
    /// with that id is currently waiting (already answered, finished, or a stale id).
    func reply(workerID: SessionID, reply: String) -> OrchestraReplyResult {
        guard let cont = handles[workerID]?.pendingQuestion else {
            return .unknown(
                message: "No worker \(workerID.rawValue) is currently waiting for a reply.")
        }
        handles[workerID]?.pendingQuestion = nil
        cont.resume(returning: .answered(reply))
        return .ok
    }

    // MARK: Supervisor drive loop (`nucleic_supervise`)

    /// Block until there is something for the supervisor to act on, then drain every pending event and
    /// report how many workers are still running plus per-batch progress.
    ///
    /// `batch` selects the wait condition:
    /// - `nil` — return as soon as *any* event is pending (the original behavior).
    /// - a batch name — keep waiting until that batch has fully reported, so the supervisor is woken
    ///   once with the whole group rather than once per worker. Two things still cut the wait short: a
    ///   worker question (it must be answered for anything to progress, and the asking worker may well
    ///   be in this batch), and having nothing left to wait for at all.
    ///
    /// Either way the drain returns *all* pending events, never only the awaited batch's: events are
    /// delivered exactly once, so withholding a sibling batch's result would strand it.
    func supervise(batch: String? = nil) async -> OrchestraSuperviseResult {
        while !tornDown, !readyToReturn(awaiting: batch) {
            await withCheckedContinuation { (cont: CheckedContinuation<Void, Never>) in
                superviseWaiters.append(cont)
            }
        }
        let drained = pendingEvents
        pendingEvents = []
        // Prune reported terminal workers so they stop counting toward `remaining_running` and the
        // handle set doesn't grow without bound. Their batch tallies live on in `batchTallies`.
        for event in drained where event.isTerminal {
            switch event {
            case .completed(let workerID, _, _, _, _), .failed(let workerID, _, _, _):
                handles.removeValue(forKey: workerID)
            case .question:
                break
            }
        }
        return OrchestraSuperviseResult(
            events: drained,
            remainingRunning: runningCount(),
            batches: batchProgress(),
            awaitedBatch: batch,
            awaitedBatchComplete: batch.map { isBatchComplete($0) })
    }

    /// Whether `supervise` has something worth returning: nothing left to wait for, a pending question
    /// (always urgent — a worker is blocked on it), the awaited batch finished, or — with no batch
    /// awaited — any pending event at all.
    private func readyToReturn(awaiting batch: String?) -> Bool {
        if runningCount() == 0 { return true }
        // A name that was never dispatched degrades to the unbatched wait rather than returning at
        // once: a typo'd batch would otherwise return empty immediately, and a model that retried it
        // in a loop would spin. Waiting for the next event still lets it see real progress and correct.
        guard let batch, batchTallies[batch] != nil else { return !pendingEvents.isEmpty }
        if pendingEvents.contains(where: { !$0.isTerminal }) { return true }
        return isBatchComplete(batch)
    }

    /// A batch is complete once none of its workers is still running. A batch that was dispatched and
    /// has already been fully drained is complete too (its handles were pruned) — so a second wait on
    /// it returns at once instead of parking on work that will never report again. A name that was
    /// never dispatched is *not* complete; see ``readyToReturn(awaiting:)``.
    private func isBatchComplete(_ batch: String) -> Bool {
        guard batchTallies[batch] != nil else { return false }
        return !handles.values.contains { $0.batch == batch && $0.state == .running }
    }

    /// Per-batch progress in dispatch order: live workers from `handles`, cumulative outcomes from
    /// `batchTallies` (which survive drain-time pruning).
    private func batchProgress() -> [OrchestraBatchProgress] {
        batchOrder.map { name in
            let running = handles.values.filter { $0.batch == name && $0.state == .running }.count
            let tally = batchTallies[name] ?? BatchTally()
            return OrchestraBatchProgress(
                name: name, running: running, completed: tally.completed, failed: tally.failed)
        }
    }

    /// Test hook: whether a batch has fully reported, without draining anything. Production code
    /// reaches this through ``supervise(batch:)``'s `awaitedBatchComplete`.
    func isBatchCompleteForTesting(_ batch: String) -> Bool { isBatchComplete(batch) }

    /// Whether the supervisor still has something to receive from this channel. This includes both
    /// live workers and terminal events that arrived after the supervisor's last `nucleic_supervise`
    /// call. The latter is important: once the final worker completes `runningCount()` is zero, but
    /// its result still has to wake a follow-up supervisor turn instead of being discarded at the
    /// preceding turn boundary.
    func needsSupervision() -> Bool {
        !handles.isEmpty || !pendingEvents.isEmpty
    }

    // MARK: Teardown

    /// End the channel: refuse further work, answer every parked question "unavailable", wake every
    /// supervise waiter, and return the ids of all tracked workers so the caller can cancel their run
    /// tasks. Idempotent-ish — a second call finds an empty handle set and returns nothing.
    func teardown() -> [SessionID] {
        tornDown = true
        let live = Array(handles.keys)
        // Collect the parked continuations, clear state, THEN resume — never resume/mutate while
        // iterating `handles`.
        let pending = handles.values.compactMap { $0.pendingQuestion }
        handles.removeAll()
        for cont in pending {
            cont.resume(returning: .unavailable(message: Self.unavailableMessage))
        }
        wakeSuperviseWaiters()
        return live
    }

    // MARK: Helpers

    private func runningCount() -> Int {
        handles.values.filter { $0.state == .running }.count
    }

    private func enqueue(_ event: OrchestraWorkerEvent) {
        pendingEvents.append(event)
        wakeSuperviseWaiters()
    }

    private func wakeSuperviseWaiters() {
        let waiters = superviseWaiters
        superviseWaiters = []
        for waiter in waiters { waiter.resume() }
    }
}
