import Foundation
import Testing

@testable import NucleicCore

/// A one-shot gate: `wait()` suspends until `open()` (and returns immediately once open).
private actor SlotTestGate {
    private var opened = false
    private var waiters: [CheckedContinuation<Void, Never>] = []
    func wait() async {
        if opened { return }
        await withCheckedContinuation { waiters.append($0) }
    }
    func open() {
        opened = true
        for waiter in waiters { waiter.resume() }
        waiters.removeAll()
    }
}

/// An `AgentBackend` whose run parks mid-turn until `interrupt()` (or `shutdown()`) arrives, then
/// settles as interrupted — the shape of an Orchestra worker whose supervisor dies mid-`join()`:
/// the spawn path must interrupt it (its result has nowhere to go) rather than let it run on.
private final class InterruptParkingBackend: AgentBackend, @unchecked Sendable {
    static let id = BackendID.claudeCode

    let capabilities = BackendCapabilities(
        interactiveApprovals: true,
        allowAlwaysScopes: [.session, .toolName],
        canModifyToolInput: true,
        partialMessageStreaming: false,
        emitsThinking: false,
        emitsFileChangeEvents: false,
        nativeResume: true,
        sandboxModes: [],
        followUpWhileRunning: false)

    private let gate = SlotTestGate()
    /// Set once `interrupt()` was called — the assertion hook for "the worker was stopped".
    let interrupted = LockedBox(false)

    func start(_ run: RunSpec) -> AsyncThrowingStream<AgentEvent, Error> {
        stream(sessionID: run.sessionID, cwd: run.worktree)
    }

    func resume(_ resume: ResumeSpec) -> AsyncThrowingStream<AgentEvent, Error> {
        stream(sessionID: resume.sessionID, cwd: resume.worktree)
    }

    private func stream(sessionID: SessionID, cwd: String) -> AsyncThrowingStream<AgentEvent, Error> {
        AsyncThrowingStream { continuation in
            let gate = self.gate
            let task = Task {
                var seq: UInt64 = 0
                func emit(_ kind: AgentEvent.Kind) {
                    seq += 1
                    continuation.yield(AgentEvent(
                        sessionID: sessionID, seq: seq, at: Date(timeIntervalSince1970: 1_700_000_000),
                        backend: Self.id, nativeType: nil, kind: kind))
                }
                emit(.sessionStarted(SessionStarted(
                    backendSessionID: "be-parked", model: "m", cwd: cwd, toolNames: [])))
                await gate.wait()  // mid-turn park; only interrupt()/shutdown() releases it
                emit(.runFinished(RunFinished(outcome: .interrupted)))
                continuation.finish()
            }
            continuation.onTermination = { _ in task.cancel() }
        }
    }

    func send(_ input: AgentInput) async throws {}
    func respond(to approvalID: ApprovalID, _ decision: Decision, by responder: String) async throws {}
    func interrupt() async {
        interrupted.set(true)
        await gate.open()
    }
    func shutdown() async { await gate.open() }
}

@MainActor
// Serialized + isolated settings for the same reasons as AppStoreTests: @MainActor timing
// assertions, and controlled-project registration touching the container-service settings.
@Suite("Orchestra worker slots — cancellation", .serialized, .isolatedContainerSettings)
struct OrchestraWorkerSlotTests {
    private func makeStore(
        repo: GitTestRepo, backend: @escaping @Sendable (Session) -> any AgentBackend
    ) -> AppStore {
        let database = try! GRDBMetadataStore(path: nil)
        let worktrees = GitWorktreeManager(now: { Date(timeIntervalSince1970: 1_700_000_000) })
        let transcriptsDir = URL(fileURLWithPath: (repo.container as NSString)
            .appendingPathComponent("transcripts"))
        return AppStore(
            database: database, worktrees: worktrees, transcriptsDir: transcriptsDir,
            now: { Date(timeIntervalSince1970: 1_700_000_000) }, backendFactory: backend)
    }

    private func waitFor(_ condition: @escaping () -> Bool, timeoutMs: Int = 3000) async {
        var elapsed = 0
        while !condition() && elapsed < timeoutMs {
            try? await Task.sleep(for: .milliseconds(10))
            elapsed += 10
        }
    }

    // MARK: - The slot gate itself

    /// A queued spawn whose supervisor dies must leave the queue: cancelling the parked task
    /// removes its waiter, `acquire` reports no slot, and the gate's bookkeeping stays intact
    /// for later acquires.
    @Test func cancelledWaiterIsRemovedAndAcquireReturnsFalse() async throws {
        let repo = try await GitTestRepo()
        defer { repo.cleanup() }
        let store = makeStore(repo: repo) { _ in ScriptedBackend { _, _ in } }
        store.defaultOrchestraMaxConcurrentWorkers = 1

        #expect(await store.acquireOrchestraWorkerSlot())  // occupy the only slot
        let parked = Task { @MainActor in await store.acquireOrchestraWorkerSlot() }
        await waitFor { store.orchestraWorkerWaiterCount == 1 }
        #expect(store.orchestraWorkerWaiterCount == 1)

        parked.cancel()
        #expect(await parked.value == false)  // no slot handed to the dead caller
        #expect(store.orchestraWorkerWaiterCount == 0)  // and its waiter is gone, not a ghost
        #expect(store.orchestraWorkersRunning == 1)  // the running count never moved

        // The gate still works: release the held slot, a fresh acquire succeeds immediately.
        store.releaseOrchestraWorkerSlot()
        #expect(await store.acquireOrchestraWorkerSlot())
        store.releaseOrchestraWorkerSlot()
        #expect(store.orchestraWorkersRunning == 0)
    }

    /// A cancelled waiter must not strand the ones queued behind it: after B is cancelled,
    /// releasing A's slot wakes C, which acquires normally.
    @Test func cancelledWaiterDoesNotStrandLaterWaiters() async throws {
        let repo = try await GitTestRepo()
        defer { repo.cleanup() }
        let store = makeStore(repo: repo) { _ in ScriptedBackend { _, _ in } }
        store.defaultOrchestraMaxConcurrentWorkers = 1

        #expect(await store.acquireOrchestraWorkerSlot())  // A holds the slot
        let b = Task { @MainActor in await store.acquireOrchestraWorkerSlot() }
        await waitFor { store.orchestraWorkerWaiterCount == 1 }
        let c = Task { @MainActor in await store.acquireOrchestraWorkerSlot() }
        await waitFor { store.orchestraWorkerWaiterCount == 2 }

        b.cancel()
        #expect(await b.value == false)
        #expect(store.orchestraWorkerWaiterCount == 1)  // only C remains queued

        store.releaseOrchestraWorkerSlot()  // A ends → the freed slot must reach C, not vanish
        #expect(await c.value == true)
        #expect(store.orchestraWorkersRunning == 1)
        store.releaseOrchestraWorkerSlot()
    }

    /// Cancellation is honored even when no parking would occur: a spawn task cancelled before
    /// (or while) acquiring never takes a slot — including under an unlimited (`<= 0`) cap.
    @Test func preCancelledAcquireNeverTakesASlot() async throws {
        let repo = try await GitTestRepo()
        defer { repo.cleanup() }
        let store = makeStore(repo: repo) { _ in ScriptedBackend { _, _ in } }

        for cap in [4, 0] {
            store.defaultOrchestraMaxConcurrentWorkers = cap
            let task = Task { @MainActor in await store.acquireOrchestraWorkerSlot() }
            task.cancel()
            #expect(await task.value == false)
            #expect(store.orchestraWorkersRunning == 0)
        }
    }

    // MARK: - The spawn path (non-blocking)

    /// A spawn is non-blocking: it returns a `worker_id` immediately (creating the worker session),
    /// and the worker's run parks on the cap *inside its own task* — never blocking the spawn call.
    /// When the supervisor's turn ends before a slot frees, teardown cancels that queued run, so the
    /// worker never starts and is archived rather than left as an active chat.
    @Test func queuedWorkerParksThenTeardownCancelsItBeforeItRuns() async throws {
        let repo = try await GitTestRepo()
        defer { repo.cleanup() }
        let store = makeStore(repo: repo) { _ in ScriptedBackend { _, _ in } }
        store.defaultOrchestraMaxConcurrentWorkers = 1

        let project = try #require(
            await store.addProject(name: "demo", rootPath: repo.root, defaultBranch: "main"))
        // An idle parent (empty prompt starts no run) to hang the spawn off.
        let parentID = try await store.createSession(in: project, title: "supervisor", prompt: "")
        let parent = try #require(await store.liveSnapshot(parentID)).session

        #expect(await store.acquireOrchestraWorkerSlot())  // occupy the only slot

        // The spawn returns at once with a worker_id, even though no slot is free.
        let result = await store.spawnOrchestraSubagent(
            OrchestraSubagentRequest(task: "scan", prompt: "check it"),
            parent: parent, project: project)
        guard case .spawned(let workerID, _, _) = result else {
            Issue.record("expected .spawned, got \(result)")
            return
        }
        #expect(store.summaries.contains { $0.id == workerID })  // the worker session exists
        // Its run task parks on the cap (the manual hold owns the only slot).
        await waitFor { store.orchestraWorkerWaiterCount == 1 }
        #expect(store.orchestraWorkerWaiterCount == 1)

        // The supervisor's turn ends: teardown cancels the queued worker's run task, so it leaves the
        // queue without ever starting, and is archived.
        await store.teardownOrchestra(supervisor: parentID)
        await waitFor { store.orchestraWorkerWaiterCount == 0 }
        #expect(store.orchestraWorkerWaiterCount == 0)
        #expect(store.orchestraWorkersRunning == 1)  // only the manual hold; the worker took no slot
        await waitFor { store.summaries.first { $0.id == workerID }?.archived == true }
        #expect(store.summaries.first { $0.id == workerID }?.archived == true)
        store.releaseOrchestraWorkerSlot()
    }

    /// Teardown interrupts a worker that is mid-turn (its result now has nowhere to go), so it is
    /// stopped rather than left running, and its concurrency slot comes back.
    @Test func teardownInterruptsRunningWorkerAndFreesSlot() async throws {
        let repo = try await GitTestRepo()
        defer { repo.cleanup() }
        let backends = LockedBox([InterruptParkingBackend]())
        let store = makeStore(repo: repo) { _ in
            let backend = InterruptParkingBackend()
            backends.set(backends.get() + [backend])
            return backend
        }

        let project = try #require(
            await store.addProject(name: "demo", rootPath: repo.root, defaultBranch: "main"))
        let parentID = try await store.createSession(in: project, title: "supervisor", prompt: "")
        let parent = try #require(await store.liveSnapshot(parentID)).session

        let result = await store.spawnOrchestraSubagent(
            OrchestraSubagentRequest(task: "long job", prompt: "work forever"),
            parent: parent, project: project)
        guard case .spawned = result else {
            Issue.record("expected .spawned, got \(result)")
            return
        }
        // Wait until the worker is mid-turn (its backend is parked, holding a slot).
        await waitFor { store.orchestraWorkersRunning == 1 && store.summaries.count == 2 }
        #expect(store.orchestraWorkersRunning == 1)
        let worker = try #require(backends.get().last)
        #expect(!worker.interrupted.get())

        // The supervisor's turn ends while the worker runs → teardown interrupts it.
        await store.teardownOrchestra(supervisor: parentID)
        await waitFor { worker.interrupted.get() }
        #expect(worker.interrupted.get())  // the worker was actively stopped, not abandoned
        await waitFor { store.orchestraWorkersRunning == 0 }
        #expect(store.orchestraWorkersRunning == 0)  // the slot came back
    }

    // MARK: - Worker model selection (Intelligent)

    /// With **Intelligent** selected, each worker is routed from its own task rather than pinned to
    /// one SKU — and, load-bearing, the sentinel is resolved to a real model. Passing it through as
    /// `--model` would hand the backend a SKU no provider knows.
    @Test func intelligentWorkerModelRoutesPerTaskAndNeverLeaksTheSentinel() async throws {
        let repo = try await GitTestRepo()
        defer { repo.cleanup() }
        let store = makeStore(repo: repo) { _ in ScriptedBackend { _, _ in } }
        store.defaultOrchestraWorkerModel = OrchestrationMode.intelligentWorkerModel
        // Both lanes connected, so routing can pick either provider's candidate.
        let project = try #require(
            await store.addProject(name: "demo", rootPath: repo.root, defaultBranch: "main"))
        let parentID = try await store.createSession(in: project, title: "supervisor", prompt: "")
        let parent = try #require(await store.liveSnapshot(parentID)).session

        // A one-line typo fix and a substantial backend task must not resolve to the same thing on a
        // per-task router — that difference IS the feature.
        let quick = store.orchestraWorkerSelection(
            task: "fix typo", prompt: "Fix the typo in the README heading.",
            parent: parent, fallbackEffort: OrchestrationMode.resolvedOrchestraEffort)
        let heavy = store.orchestraWorkerSelection(
            task: "build the API",
            prompt: "Implement the server endpoint, its data model, and the algorithm behind it.",
            parent: parent, fallbackEffort: OrchestrationMode.resolvedOrchestraEffort)

        for selection in [quick, heavy] {
            let model = try #require(selection.model)
            #expect(model != OrchestrationMode.intelligentWorkerModel)
            #expect(BackendID.forModel(model) != nil)  // a real, runnable SKU
            #expect(selection.routing != nil)  // provenance recorded for "why this model"
        }
        #expect(quick.routing?.purpose == .quickFix)
        #expect(heavy.routing?.purpose == .backendImpl)

        // A pinned SKU is passed straight through, with the caller's effort and no routing note.
        store.defaultOrchestraWorkerModel = "claude-haiku-4-5"
        let pinned = store.orchestraWorkerSelection(
            task: "fix typo", prompt: "Fix the typo.", parent: parent, fallbackEffort: "xhigh")
        #expect(pinned.model == "claude-haiku-4-5")
        #expect(pinned.effort == "xhigh")
        #expect(pinned.routing == nil)
    }

    /// A spawned worker actually runs the routed model — the selection reaches the created session,
    /// not just the return value.
    @Test func spawnedWorkerSessionUsesTheRoutedModel() async throws {
        let repo = try await GitTestRepo()
        defer { repo.cleanup() }
        let store = makeStore(repo: repo) { _ in ScriptedBackend { _, _ in } }
        store.defaultOrchestraWorkerModel = OrchestrationMode.intelligentWorkerModel

        let project = try #require(
            await store.addProject(name: "demo", rootPath: repo.root, defaultBranch: "main"))
        let parentID = try await store.createSession(in: project, title: "supervisor", prompt: "")
        let parent = try #require(await store.liveSnapshot(parentID)).session

        let result = await store.spawnOrchestraSubagent(
            OrchestraSubagentRequest(task: "fix typo", prompt: "Fix the typo in the heading."),
            parent: parent, project: project)
        guard case .spawned(let workerID, _, let model) = result else {
            Issue.record("expected .spawned, got \(result)")
            return
        }
        let reported = try #require(model)
        #expect(reported != OrchestrationMode.intelligentWorkerModel)
        let worker = try #require(await store.liveSnapshot(workerID)).session
        #expect(worker.model == reported)
        // Routing provenance rides onto the session, so its info panel can say why this model.
        #expect(worker.routedPurpose == PromptPurpose.quickFix.rawValue)
        #expect(worker.routedReason?.isEmpty == false)

        await store.teardownOrchestra(supervisor: parentID)
    }

    // MARK: - Plan delegation

    /// Submitting a plan spawns a real worker session for EVERY task in EVERY batch at once — batches
    /// stagger when the supervisor *consumes* results, not when work is dispatched — and each worker is
    /// tagged with its batch so the supervisor can wait on one group. Workers beyond the cap queue
    /// rather than being dropped, so a plan wider than the cap still runs in full.
    @Test func delegatePlanSpawnsEveryTaskTaggedWithItsBatch() async throws {
        let repo = try await GitTestRepo()
        defer { repo.cleanup() }
        let store = makeStore(repo: repo) { _ in ScriptedBackend { _, _ in } }
        store.defaultOrchestraMaxConcurrentWorkers = 1  // narrower than the plan: nothing may be lost

        let project = try #require(
            await store.addProject(name: "demo", rootPath: repo.root, defaultBranch: "main"))
        let parentID = try await store.createSession(in: project, title: "supervisor", prompt: "")
        let parent = try #require(await store.liveSnapshot(parentID)).session

        let result = await store.delegateOrchestraPlan(
            OrchestraPlanRequest(
                plan: "explore, then build",
                batches: [
                    OrchestraPlanBatch(
                        name: "explore", purpose: "research",
                        tasks: [
                            OrchestraPlanTask(task: "read A", prompt: "read file A"),
                            OrchestraPlanTask(task: "read B", prompt: "read file B"),
                        ]),
                    OrchestraPlanBatch(
                        name: "implement",
                        tasks: [OrchestraPlanTask(task: "build", prompt: "write the code")]),
                ]),
            parent: parent, project: project)

        guard case .dispatched(let batches) = result else {
            Issue.record("expected .dispatched, got \(result)")
            return
        }
        #expect(batches.map(\.name) == ["explore", "implement"])
        #expect(batches[0].workers.count == 2)
        #expect(batches[1].workers.count == 1)
        #expect(batches.allSatisfy { $0.rejected.isEmpty })

        // Every task became a real worker session, cap notwithstanding (3 workers + the supervisor).
        let workerIDs = batches.flatMap(\.workers).map(\.workerID)
        #expect(Set(workerIDs).count == 3)
        for workerID in workerIDs {
            #expect(store.summaries.contains { $0.id == workerID })
        }
        #expect(store.subagentSummaries(for: parentID).count == 3)

        // The batch tags reached the coordinator, so a supervise call can wait on one group. Draining
        // reports both batches with the right worker counts.
        let drained = await store.superviseOrchestra(supervisor: parentID)
        let explore = drained.batches.first { $0.name == "explore" }
        let implement = drained.batches.first { $0.name == "implement" }
        #expect(explore != nil)
        #expect(implement != nil)
        #expect((explore?.running ?? 0) + (explore?.completed ?? 0) + (explore?.failed ?? 0) == 2)
        #expect((implement?.running ?? 0) + (implement?.completed ?? 0) + (implement?.failed ?? 0) == 1)

        await store.teardownOrchestra(supervisor: parentID)
    }

    /// A plan with no delegable task is refused outright rather than acknowledged, so the supervisor
    /// can't settle into a supervise loop waiting on workers that were never created.
    @Test func delegatePlanDeniesAPlanWithNoTasks() async throws {
        let repo = try await GitTestRepo()
        defer { repo.cleanup() }
        let store = makeStore(repo: repo) { _ in ScriptedBackend { _, _ in } }

        let project = try #require(
            await store.addProject(name: "demo", rootPath: repo.root, defaultBranch: "main"))
        let parentID = try await store.createSession(in: project, title: "supervisor", prompt: "")
        let parent = try #require(await store.liveSnapshot(parentID)).session

        let empty = await store.delegateOrchestraPlan(
            OrchestraPlanRequest(plan: "do it", batches: []), parent: parent, project: project)
        guard case .denied = empty else {
            Issue.record("expected .denied for an empty plan, got \(empty)")
            return
        }

        // A batch whose every task is unusable is likewise a denial, not a half-dispatch: the spawner
        // refuses a blank prompt, so nothing was created and the supervisor is told so.
        let blank = await store.delegateOrchestraPlan(
            OrchestraPlanRequest(
                plan: "do it",
                batches: [
                    OrchestraPlanBatch(
                        name: "explore", tasks: [OrchestraPlanTask(task: "noop", prompt: "   ")])
                ]),
            parent: parent, project: project)
        guard case .denied = blank else {
            Issue.record("expected .denied when no worker could spawn, got \(blank)")
            return
        }
        #expect(store.subagentSummaries(for: parentID).isEmpty)
    }

    // MARK: - The blocking subagent path (non-Orchestra, Control projects)

    /// The non-Orchestra managed subagent is **blocking**: `runBlockingSubagent` runs the worker to
    /// completion and returns its final reply inline (a drop-in for native Task), then archives the
    /// worker session (nothing resumes it) while keeping its `spawnedBy` link so it rolls up under the
    /// parent in the Subagents panel.
    @Test func blockingSubagentRunsToCompletionAndReturnsFinalReply() async throws {
        let repo = try await GitTestRepo()
        defer { repo.cleanup() }
        let store = makeStore(repo: repo) { _ in
            ScriptedBackend { e, _ in
                e.emit(.assistantText(
                    TextChunk(messageID: "m1", text: "Found the bug in Foo.swift.", isPartial: false)))
                e.emit(.turnCompleted(TurnCompleted(stopReason: "end_turn", usage: nil)))
                e.emit(.runFinished(RunFinished(outcome: .completed)))
            }
        }
        let project = try #require(
            await store.addProject(name: "demo", rootPath: repo.root, defaultBranch: "main"))
        let parentID = try await store.createSession(in: project, title: "chat", prompt: "")
        let parent = try #require(await store.liveSnapshot(parentID)).session

        let result = await store.runBlockingSubagent(
            OrchestraSubagentRequest(task: "hunt", prompt: "find the bug"),
            parent: parent, project: project)

        guard case .completed(let workerID, let title, _, let output) = result else {
            Issue.record("expected .completed, got \(result)")
            return
        }
        #expect(output == "Found the bug in Foo.swift.")  // the worker's report, returned inline
        #expect(title == "hunt")
        #expect(store.summaries.contains { $0.id == workerID })  // it ran as a real observed session
        #expect(store.summaries.first { $0.id == workerID }?.spawnedBySessionID == parentID)
        await waitFor { store.summaries.first { $0.id == workerID }?.archived == true }
        #expect(store.summaries.first { $0.id == workerID }?.archived == true)
    }

    /// An empty prompt is refused rather than spawning an empty worker.
    @Test func blockingSubagentDeniesEmptyPrompt() async throws {
        let repo = try await GitTestRepo()
        defer { repo.cleanup() }
        let store = makeStore(repo: repo) { _ in ScriptedBackend { _, _ in } }
        let project = try #require(
            await store.addProject(name: "demo", rootPath: repo.root, defaultBranch: "main"))
        let parentID = try await store.createSession(in: project, title: "chat", prompt: "")
        let parent = try #require(await store.liveSnapshot(parentID)).session

        let result = await store.runBlockingSubagent(
            OrchestraSubagentRequest(task: "noop", prompt: "   "),
            parent: parent, project: project)
        guard case .denied = result else {
            Issue.record("expected .denied, got \(result)")
            return
        }
    }

    /// Enforcement gate: on a Nucleic Control project, an ordinary (non-Orchestra) chat's run exposes
    /// the blocking subagent spawner — so native Task/Agent is denied toward `nucleic_subagent` — while
    /// the non-blocking Orchestra spawner stays absent (this isn't a supervisor).
    @Test func controlProjectExposesBlockingSubagentSpawnerNotOrchestra() async throws {
        let repo = try await GitTestRepo(controlled: true)
        defer { repo.cleanup() }
        let backends = LockedBox([SessionID: ScriptedBackend]())
        let store = makeStore(repo: repo) { session in
            let backend = ScriptedBackend { e, _ in
                e.emit(.runFinished(RunFinished(outcome: .completed)))
            }
            backends.set(backends.get().merging([session.id: backend]) { _, new in new })
            return backend
        }
        let project = try #require(
            await store.addProject(name: "demo", rootPath: repo.root, defaultBranch: "main"))
        #expect(project.isNucleicControlled)
        let id = try await store.createSession(in: project, title: "chat", prompt: "go")
        await waitFor { backends.get()[id]?.lastRun != nil }
        let run = try #require(backends.get()[id]?.lastRun)
        #expect(run.blockingSubagentSpawner != nil)
        #expect(run.orchestraSubagentSpawner == nil)
    }

    /// Off Nucleic Control, native Task/Agent stay untouched — no blocking spawner is exposed, so
    /// enforcement is scoped to Control projects.
    @Test func nonControlProjectDoesNotExposeBlockingSubagentSpawner() async throws {
        let repo = try await GitTestRepo()  // plain temp dir → not under the control base
        defer { repo.cleanup() }
        let backends = LockedBox([SessionID: ScriptedBackend]())
        let store = makeStore(repo: repo) { session in
            let backend = ScriptedBackend { e, _ in
                e.emit(.runFinished(RunFinished(outcome: .completed)))
            }
            backends.set(backends.get().merging([session.id: backend]) { _, new in new })
            return backend
        }
        let project = try #require(
            await store.addProject(name: "demo", rootPath: repo.root, defaultBranch: "main"))
        #expect(!project.isNucleicControlled)
        let id = try await store.createSession(in: project, title: "chat", prompt: "go")
        await waitFor { backends.get()[id]?.lastRun != nil }
        let run = try #require(backends.get()[id]?.lastRun)
        #expect(run.blockingSubagentSpawner == nil)
    }
}
