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
        var state: WorkerState
        /// Set while the worker is parked in `nucleic_ask_supervisor` awaiting the supervisor's reply.
        var pendingQuestion: CheckedContinuation<OrchestraAskSupervisorResult, Never>?
    }

    private var handles: [SessionID: WorkerHandle] = [:]
    /// 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. 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.
    func register(workerID: SessionID, task: String) {
        if var existing = handles[workerID] {
            existing.state = .running
            handles[workerID] = existing
        } else {
            handles[workerID] = WorkerHandle(task: task, 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 in
            .completed(workerID: workerID, task: task, model: model, output: output)
        }
    }

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

    private func finish(workerID: SessionID, event: (String) -> OrchestraWorkerEvent) {
        guard !tornDown else { return }
        let task = handles[workerID]?.task ?? ""
        // 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
        enqueue(event(task))
    }

    // 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: "", state: .running, pendingQuestion: nil)
                }
                handles[workerID]?.pendingQuestion = cont
                let task = handles[workerID]?.task ?? ""
                enqueue(.question(workerID: workerID, task: task, 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 at least one event is pending (or nothing remains to wait for), then drain the
    /// batch and report how many workers are still running.
    func supervise() async -> OrchestraSuperviseResult {
        while pendingEvents.isEmpty, !tornDown, runningCount() > 0 {
            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.
        for event in drained {
            switch event {
            case .completed(let workerID, _, _, _), .failed(let workerID, _, _):
                handles.removeValue(forKey: workerID)
            case .question:
                break
            }
        }
        return OrchestraSuperviseResult(events: drained, remainingRunning: runningCount())
    }

    // 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() }
    }
}
