import Foundation
import Testing

@testable import NucleicCore

@Suite("Orchestra worker coordinator")
struct OrchestraWorkerCoordinatorTests {
    private func worker(_ id: String) -> SessionID { SessionID(rawValue: id) }

    /// A completed worker surfaces as a `completed` event and drops out of `remaining_running`.
    @Test func superviseDrainsCompletedEvent() async {
        let coord = OrchestraWorkerCoordinator()
        let w = worker("w1")
        await coord.register(workerID: w, task: "scan")
        await coord.complete(workerID: w, model: "m", output: "done")

        let result = await coord.supervise()
        #expect(result.remainingRunning == 0)
        guard case .completed(let id, let task, _, let model, let output)? = result.events.first else {
            Issue.record("expected a completed event, got \(result.events)")
            return
        }
        #expect(id == w)
        #expect(task == "scan")
        #expect(model == "m")
        #expect(output == "done")
    }

    /// A failed worker surfaces as a `failed` event.
    @Test func superviseDrainsFailedEvent() async {
        let coord = OrchestraWorkerCoordinator()
        let w = worker("w1")
        await coord.register(workerID: w, task: "scan")
        await coord.fail(workerID: w, message: "boom")

        let result = await coord.supervise()
        #expect(result.remainingRunning == 0)
        guard case .failed(let id, _, _, let message)? = result.events.first else {
            Issue.record("expected a failed event, got \(result.events)")
            return
        }
        #expect(id == w)
        #expect(message == "boom")
    }

    /// With no workers registered, `supervise` returns at once — the supervisor learns it is done.
    @Test func superviseReturnsEmptyWhenNothingRunning() async {
        let coord = OrchestraWorkerCoordinator()
        #expect(await coord.needsSupervision() == false)
        let result = await coord.supervise()
        #expect(result.events.isEmpty)
        #expect(result.remainingRunning == 0)
    }

    /// A turn boundary must remain wakeable both while a worker runs and after its terminal event
    /// arrives but before the supervisor drains it. Only draining the event settles the channel.
    @Test func needsSupervisionCoversRunningAndUndrainedTerminalEvents() async {
        let coord = OrchestraWorkerCoordinator()
        let w = worker("w1")

        await coord.register(workerID: w, task: "scan")
        #expect(await coord.needsSupervision())

        await coord.complete(workerID: w, model: "m", output: "done")
        #expect(await coord.needsSupervision())

        _ = await coord.supervise()
        #expect(await coord.needsSupervision() == false)
    }

    /// A worker's question surfaces to the supervisor, and the supervisor's reply resumes that exact
    /// worker with the answer. `supervise` blocks until the question arrives, so no polling is needed.
    @Test func askSurfacesQuestionAndReplyResumesTheWorker() async {
        let coord = OrchestraWorkerCoordinator()
        let w = worker("w1")
        await coord.register(workerID: w, task: "scan")

        let asked = Task { await coord.ask(workerID: w, question: "which path?") }

        let result = await coord.supervise()  // blocks until the question is raised
        guard case .question(let id, let task, _, let question)? = result.events.first else {
            Issue.record("expected a question event, got \(result.events)")
            return
        }
        #expect(id == w)
        #expect(task == "scan")
        #expect(question == "which path?")
        #expect(result.remainingRunning == 1)  // still running — parked awaiting the reply

        let replyResult = await coord.reply(workerID: w, reply: "the left one")
        #expect(replyResult == .ok)

        let answer = await asked.value
        #expect(answer == .answered("the left one"))
    }

    /// Replying to a worker that isn't waiting is a no-op `unknown`, not a crash.
    @Test func replyToUnknownWorkerReportsUnknown() async {
        let coord = OrchestraWorkerCoordinator()
        let result = await coord.reply(workerID: worker("ghost"), reply: "hi")
        guard case .unknown = result else {
            Issue.record("expected .unknown, got \(result)")
            return
        }
    }

    /// Teardown force-drains: a worker parked in `ask` is answered "unavailable" (so its call can
    /// unwind rather than hang), and teardown reports the worker as live so its run task gets
    /// cancelled by the caller.
    @Test func teardownAnswersPendingQuestionUnavailable() async {
        let coord = OrchestraWorkerCoordinator()
        let w = worker("w1")
        await coord.register(workerID: w, task: "scan")

        let asked = Task { await coord.ask(workerID: w, question: "wait for me?") }
        _ = await coord.supervise()  // ensure the question is registered before teardown

        let live = await coord.teardown()
        #expect(live.contains(w))

        let answer = await asked.value
        guard case .unavailable = answer else {
            Issue.record("expected .unavailable after teardown, got \(answer)")
            return
        }
    }

    /// After teardown, a fresh `ask` short-circuits to unavailable rather than parking forever.
    @Test func askAfterTeardownIsImmediatelyUnavailable() async {
        let coord = OrchestraWorkerCoordinator()
        _ = await coord.teardown()
        let answer = await coord.ask(workerID: worker("w1"), question: "anyone?")
        guard case .unavailable = answer else {
            Issue.record("expected .unavailable, got \(answer)")
            return
        }
    }

    // MARK: - Batches

    /// The point of a batch: waiting on one group returns when *that* group has finished, while
    /// workers in other batches are still running — so the supervisor can act on the explore results
    /// without waiting for the rest of the plan.
    @Test func superviseOnBatchReturnsWhenThatBatchCompletesWhileOthersRun() async {
        let coord = OrchestraWorkerCoordinator()
        let e1 = worker("e1"), e2 = worker("e2"), impl = worker("i1")
        await coord.register(workerID: e1, task: "read A", batch: "explore")
        await coord.register(workerID: e2, task: "read B", batch: "explore")
        await coord.register(workerID: impl, task: "build", batch: "implement")

        // One of two explore workers done: the batch isn't complete, so a batch wait must keep waiting.
        await coord.complete(workerID: e1, model: "m", output: "found A")
        #expect(await coord.isBatchCompleteForTesting("explore") == false)

        await coord.complete(workerID: e2, model: "m", output: "found B")

        let result = await coord.supervise(batch: "explore")
        #expect(result.awaitedBatch == "explore")
        #expect(result.awaitedBatchComplete == true)
        #expect(result.events.count == 2)
        #expect(result.events.allSatisfy { $0.batch == "explore" })
        // The implement batch is untouched and still running — that's what makes the early wake useful.
        #expect(result.remainingRunning == 1)
        let implement = result.batches.first { $0.name == "implement" }
        #expect(implement?.running == 1)
        #expect(implement?.isComplete == false)
    }

    /// Per-batch tallies survive the drain that prunes terminal workers, so a batch stays reported as
    /// complete (with its counts) after its results have been read.
    @Test func batchProgressSurvivesDrainAndCountsOutcomes() async {
        let coord = OrchestraWorkerCoordinator()
        let a = worker("a"), b = worker("b")
        await coord.register(workerID: a, task: "one", batch: "verify")
        await coord.register(workerID: b, task: "two", batch: "verify")
        await coord.complete(workerID: a, model: "m", output: "ok")
        await coord.fail(workerID: b, message: "boom")

        let first = await coord.supervise(batch: "verify")
        #expect(first.batches.count == 1)
        #expect(first.batches.first?.completed == 1)
        #expect(first.batches.first?.failed == 1)
        #expect(first.batches.first?.running == 0)

        // Waiting on an already-drained batch returns at once rather than parking on work that will
        // never report again — and still reports the batch's cumulative counts.
        let second = await coord.supervise(batch: "verify")
        #expect(second.events.isEmpty)
        #expect(second.awaitedBatchComplete == true)
        #expect(second.batches.first?.completed == 1)
        #expect(second.batches.first?.failed == 1)
    }

    /// A worker question always cuts a batch wait short: the asking worker is blocked until it's
    /// answered, so the supervisor must be woken even though its awaited batch isn't finished.
    @Test func questionInterruptsABatchWait() async {
        let coord = OrchestraWorkerCoordinator()
        let slow = worker("slow"), asker = worker("asker")
        await coord.register(workerID: slow, task: "slow one", batch: "explore")
        await coord.register(workerID: asker, task: "asker", batch: "explore")

        let asked = Task { await coord.ask(workerID: asker, question: "which path?") }
        let result = await coord.supervise(batch: "explore")

        #expect(result.awaitedBatchComplete == false)
        guard case .question(let id, _, let batch, let question)? = result.events.first else {
            Issue.record("expected a question event, got \(result.events)")
            return
        }
        #expect(id == asker)
        #expect(batch == "explore")  // the event carries its batch so the supervisor can place it
        #expect(question == "which path?")

        _ = await coord.reply(workerID: asker, reply: "left")
        _ = await asked.value
    }

    /// Waiting on a batch that was never dispatched must not spin: it degrades to the unbatched wait
    /// (return on the next event) rather than returning empty immediately, which a retrying model
    /// would turn into a busy loop.
    @Test func superviseOnUnknownBatchWaitsForTheNextEventInstead() async {
        let coord = OrchestraWorkerCoordinator()
        let w = worker("w1")
        await coord.register(workerID: w, task: "scan", batch: "explore")

        let waiting = Task { await coord.supervise(batch: "typo-batch") }
        await coord.complete(workerID: w, model: "m", output: "done")

        let result = await waiting.value
        #expect(result.events.count == 1)
        #expect(result.awaitedBatch == "typo-batch")
        #expect(result.awaitedBatchComplete == false)
    }

    /// An unbatched fan-out keeps its original shape: no batch progress is reported, and events carry
    /// no batch, so nothing about the pre-batch behavior changes.
    @Test func unbatchedWorkersReportNoBatchMetadata() async {
        let coord = OrchestraWorkerCoordinator()
        let w = worker("w1")
        await coord.register(workerID: w, task: "scan")
        await coord.complete(workerID: w, model: "m", output: "done")

        let result = await coord.supervise()
        #expect(result.batches.isEmpty)
        #expect(result.awaitedBatch == nil)
        #expect(result.awaitedBatchComplete == nil)
        #expect(result.events.first?.batch == nil)
        // The wire shape omits both, so an unbatched run's result JSON is unchanged.
        let json = result.wireJSON()
        #expect(!json.contains("batches"))
        #expect(!json.contains("awaited_batch"))
    }
}
