// ManagedGatewayTests.swift
// OSGKeyboard · Extension Tests
//
// Hermetic grant rotation and managed LLM tests using a private URLProtocol stub.

import Foundation
@testable import OSGKeyboardShared
import XCTest

final class ManagedGatewayTests: XCTestCase {
    private let baseURL = URL(string: "https://gateway.test")!

    override func setUp() {
        super.setUp()
        GatewayStub.shared.reset()
    }

    override func tearDown() {
        GatewayStub.shared.reset()
        super.tearDown()
    }

    func testFiveMinuteAccessBoundaryAndScope() {
        let start = Date(timeIntervalSince1970: 1_000)
        let value = credentials(accessExpiresAt: start.addingTimeInterval(3_600), receivedAt: start)
        XCTAssertEqual(value.effectiveAccessExpiresAt, start.addingTimeInterval(300))
        XCTAssertTrue(value.hasUsableAccessToken(for: .polish, at: start.addingTimeInterval(269)))
        XCTAssertFalse(value.hasUsableAccessToken(for: .polish, at: start.addingTimeInterval(270)))
        XCTAssertFalse(value.hasUsableAccessToken(for: .asr, at: start))
    }

    func testScopePolicyExcludesASRForLocalEngineAndIncludesItForCloud() {
        let localScopes = ManagedGatewayScopePolicy.scopes(engineMode: "local")
        XCTAssertEqual(localScopes, [.polish, .assistant, .agent])
        XCTAssertFalse(localScopes.contains(.asr))

        let cloudScopes = ManagedGatewayScopePolicy.scopes(engineMode: "cloud")
        XCTAssertEqual(cloudScopes, [.polish, .assistant, .agent, .asr])
    }

    func testOOBEFeatureConflictIsNotMappedToNetworkFailure() {
        let error = ManagedGatewayHTTP.error(
            data: errorJSON("oobe_feature_already_used"),
            status: 409,
            requestId: "oobe-conflict"
        )

        XCTAssertEqual(error, .oobeFeatureAlreadyUsed)
    }

    func testCreateGrantSendsAccountTokenButStoresOnlyGrant() async throws {
        GatewayStub.shared.enqueue(
            201,
            grantJSON(
                access: "grant-access",
                refresh: "grant-refresh",
                scopes: ["polish", "ai"]
            )
        )
        let store = MemoryGrantStore()
        let coordinator = makeCoordinator(store: store)

        _ = try await coordinator.createGrant(
            accountAccessToken: "account-secret",
            scopes: [.polish, .assistant],
            lifetimeSeconds: 3_600,
            idempotencyKey: "create-idempotency"
        )

        let stored = await store.value()
        XCTAssertEqual(stored?.refreshToken, "grant-refresh")
        let request = try XCTUnwrap(GatewayStub.shared.requests().first)
        XCTAssertEqual(request.url?.path, "/v1/gateway/grants")
        XCTAssertEqual(request.value(forHTTPHeaderField: "Authorization"), "Bearer account-secret")
        XCTAssertEqual(request.value(forHTTPHeaderField: "Idempotency-Key"), "create-idempotency")
        let body = try jsonBody(request)
        XCTAssertEqual(Set(body["scopes"] as? [String] ?? []), Set(["polish", "ai"]))
        XCTAssertEqual(body["lifetimeSeconds"] as? Int, 3_600)
        let persisted = try JSONEncoder().encode(stored)
        XCTAssertFalse(
            (String(bytes: persisted, encoding: .utf8) ?? "").contains("account-secret")
        )
    }

    func testCreateGrantRejectsBroaderReturnedScopes() async throws {
        GatewayStub.shared.enqueue(
            201,
            grantJSON(
                access: "grant-access",
                refresh: "grant-refresh",
                scopes: ["polish", "ai", "agent"]
            )
        )
        let store = MemoryGrantStore()
        let coordinator = makeCoordinator(store: store)

        await XCTAssertThrowsManaged(.invalidGrant) {
            _ = try await coordinator.createGrant(
                accountAccessToken: "account-secret",
                scopes: [.polish],
                idempotencyKey: "create-scope-boundary"
            )
        }
        let stored = await store.value()
        XCTAssertNil(stored)
    }

    func testConcurrentRefreshIsMergedAndHashesOldTokenForIdempotency() async throws {
        let now = Date(timeIntervalSince1970: 2_000)
        let oldRefresh = "old-refresh-token-that-is-at-least-32-characters"
        let store = MemoryGrantStore(
            credentials(
                accessToken: "expired",
                refreshToken: oldRefresh,
                accessExpiresAt: now.addingTimeInterval(-1),
                receivedAt: now.addingTimeInterval(-600)
            )
        )
        GatewayStub.shared.enqueue(
            200,
            grantJSON(access: "new-access", refresh: "new-refresh"),
            delay: 0.05
        )
        let coordinator = makeCoordinator(store: store, now: { now })

        async let first = coordinator.accessToken(for: .assistant)
        async let second = coordinator.accessToken(for: .assistant)
        let values = try await (first, second)

        XCTAssertEqual(values.0, "new-access")
        XCTAssertEqual(values.1, "new-access")
        let requests = GatewayStub.shared.requests()
        XCTAssertEqual(requests.count, 1)
        let request = try XCTUnwrap(requests.first)
        let key = try XCTUnwrap(request.value(forHTTPHeaderField: "Idempotency-Key"))
        XCTAssertEqual(key, GatewayGrantCoordinator.refreshTokenIdempotencyKey(oldRefresh))
        XCTAssertFalse(key.contains(oldRefresh))
        XCTAssertEqual(try jsonBody(request)["refreshToken"] as? String, oldRefresh)
    }

    func testSignedOutPolicyRejectsCachedAccountGrant() async throws {
        let suiteName = "ManagedGatewayTests.signedOut.\(UUID().uuidString)"
        let defaults = try XCTUnwrap(UserDefaults(suiteName: suiteName))
        defer { defaults.removePersistentDomain(forName: suiteName) }
        let now = Date(timeIntervalSince1970: 3_000)
        let store = MemoryGrantStore(
            credentials(
                accessToken: "cached-account-access",
                accessExpiresAt: now.addingTimeInterval(240),
                refreshExpiresAt: now.addingTimeInterval(3_600),
                receivedAt: now
            )
        )
        let accountState = AppGroupStore(defaults: defaults)
        let coordinator = GatewayGrantCoordinator(
            baseURL: baseURL,
            store: store,
            session: stubSession(),
            now: { now },
            accountAccessPolicy: AppGroupManagedGatewayAccountAccessPolicy(
                defaults: defaults
            )
        )

        await XCTAssertThrowsManaged(.missingGrant) {
            _ = try await coordinator.accessToken(for: .polish)
        }
        XCTAssertTrue(GatewayStub.shared.requests().isEmpty)

        accountState.setManagedGatewayAccountSessionAvailable(true)
        let token = try await coordinator.accessToken(for: .polish)
        XCTAssertEqual(token, "cached-account-access")
    }

    func testManagedClientMapsAllCapabilitiesHeadersBodyAndRequestIds() async throws {
        let now = Date(timeIntervalSince1970: 3_000)
        let store = MemoryGrantStore(credentials(accessToken: "access", receivedAt: now))
        let coordinator = makeCoordinator(store: store, now: { now })
        let ids = LockedValues(["request-0001", "request-0002", "request-0003"])

        let cases: [(ManagedLLMClient.Capability, ManagedGatewayTaskKind)] = [
            (.polish, .dictationPolish),
            (.assistant, .aiQuestion),
            (.agent, .agentPlanning)
        ]
        for (capability, _) in cases {
            GatewayStub.shared.enqueue(200, Data(#"{"choices":[{"message":{"content":"ok"}}]}"#.utf8))
            let client = ManagedLLMClient(
                capability: capability,
                grants: coordinator,
                baseURL: baseURL,
                session: stubSession(),
                requestId: { ids.next() }
            )
            let result = try await client.polish(
                "input",
                systemPrompt: "context",
                timeout: 7,
                options: LLMGenerationOptions(temperature: 0.4, maxTokens: 321)
            )
            XCTAssertEqual(result, "ok")
        }

        let requests = GatewayStub.shared.requests()
        XCTAssertEqual(
            requests.compactMap(\.url?.path),
            [
                "/v1/gateway/llm/polish",
                "/v1/gateway/llm/ai",
                "/v1/gateway/llm/agent"
            ]
        )
        XCTAssertEqual(Set(requests.compactMap { $0.value(forHTTPHeaderField: "X-Request-ID") }).count, 3)
        for (request, expectedCase) in zip(requests, cases) {
            XCTAssertEqual(request.value(forHTTPHeaderField: "Authorization"), "Bearer access")
            let body = try jsonBody(request)
            XCTAssertEqual(body["input"] as? String, "input")
            XCTAssertEqual(body["context"] as? String, "context")
            XCTAssertEqual(body["maxOutputTokens"] as? Int, 321)
            XCTAssertEqual(body["temperature"] as? Double, 0.4)
            XCTAssertEqual(body["stream"] as? Bool, false)
            XCTAssertEqual(body["taskKind"] as? String, expectedCase.1.rawValue)
            XCTAssertEqual(request.timeoutInterval, 7)
        }
    }

    func testManagedClientSerializesFineGrainedTaskKinds() async throws {
        let now = Date(timeIntervalSince1970: 3_500)
        let store = MemoryGrantStore(credentials(accessToken: "access", receivedAt: now))
        let coordinator = makeCoordinator(store: store, now: { now })
        let cases: [(ManagedLLMClient.Capability, ManagedGatewayTaskKind)] = [
            (.polish, .translation),
            (.polish, .editLastInput),
            (.assistant, .currentInformationQuestion),
            (.assistant, .clipboardTransform),
            (.assistant, .customSkill)
        ]

        for (capability, taskKind) in cases {
            GatewayStub.shared.enqueue(200, Data(#"{"output_text":"ok"}"#.utf8))
            let client = ManagedLLMClient(
                capability: capability,
                taskKind: taskKind,
                grants: coordinator,
                baseURL: baseURL,
                session: stubSession()
            )
            _ = try await client.polish("input", systemPrompt: "context")
        }

        let taskKinds = try GatewayStub.shared.requests().map {
            try XCTUnwrap(jsonBody($0)["taskKind"] as? String)
        }
        XCTAssertEqual(taskKinds, cases.map { $0.1.rawValue })
    }

    func testManagedTaskKindsResolveToMatchingCapabilities() {
        XCTAssertEqual(
            ManagedLLMClient.Capability.resolve(taskKind: .dictationPolish),
            .polish
        )
        XCTAssertEqual(
            ManagedLLMClient.Capability.resolve(taskKind: .customSkill),
            .assistant
        )
        XCTAssertEqual(
            ManagedLLMClient.Capability.resolve(taskKind: .agentPlanning),
            .agent
        )
    }

    func testManagedHotwordSerializesSourceAndCurrentInformationIntent() async throws {
        let now = Date(timeIntervalSince1970: 3_600)
        let store = MemoryGrantStore(credentials(accessToken: "access", receivedAt: now))
        GatewayStub.shared.enqueue(200, Data(#"{"output_text":"ok"}"#.utf8))
        let client = ManagedLLMClient(
            capability: .assistant,
            taskKind: .currentInformationQuestion,
            requestSource: .hotword,
            grants: makeCoordinator(store: store, now: { now }),
            baseURL: baseURL,
            session: stubSession()
        )

        _ = try await client.polish("today's news", systemPrompt: "context")

        let request = try XCTUnwrap(GatewayStub.shared.requests().first)
        let body = try jsonBody(request)
        XCTAssertEqual(body["taskKind"] as? String, "current_information_question")
        XCTAssertEqual(body["requestSource"] as? String, "hotword")
    }

    func testManagedClientSerializesOOBERequestPurpose() async throws {
        let now = Date(timeIntervalSince1970: 3_750)
        let store = MemoryGrantStore(credentials(accessToken: "account-access", receivedAt: now))
        let oobeStore = MemoryGrantStore(
            credentials(
                accessToken: "oobe-access",
                receivedAt: now,
                scopes: [.polish, .assistant]
            )
        )
        let suite = "managed.oobe.\(UUID().uuidString)"
        let defaults = try XCTUnwrap(UserDefaults(suiteName: suite))
        defer { defaults.removePersistentDomain(forName: suite) }
        let practice = try XCTUnwrap(
            KeyboardSetupBridge.beginOOBEPracticeSession(
                expectedFeature: .voiceInput,
                defaults: defaults,
                now: now
            )
        )
        XCTAssertEqual(practice.expectedFeature, .voiceInput)
        GatewayStub.shared.enqueue(200, Data(#"{"output_text":"ok"}"#.utf8))
        let client = ManagedLLMClient(
            capability: .polish,
            taskKind: .dictationPolish,
            requestPurpose: .oobe,
            oobeFeature: .voiceInput,
            grants: makeCoordinator(store: store, now: { now }),
            oobeGrants: OOBEGatewayGrantCoordinator(
                baseURL: baseURL,
                store: oobeStore,
                session: stubSession(),
                now: { now },
                practiceSession: { practice }
            ),
            baseURL: baseURL,
            session: stubSession()
        )

        _ = try await client.polish("input", systemPrompt: "context")

        let request = try XCTUnwrap(GatewayStub.shared.requests().last)
        let body = try jsonBody(request)
        XCTAssertEqual(body["taskKind"] as? String, "dictation_polish")
        XCTAssertEqual(body["requestPurpose"] as? String, "oobe")
        XCTAssertEqual(body["oobeFeature"] as? String, "voice_input")
        XCTAssertEqual(request.value(forHTTPHeaderField: "Authorization"), "Bearer oobe-access")
        let accountCredentials = await store.value()
        XCTAssertEqual(accountCredentials?.accessToken, "account-access")
    }

    func testOOBEGrantRejectsAgentOrASRScopes() async {
        let now = Date(timeIntervalSince1970: 3_775)
        let coordinator = OOBEGatewayGrantCoordinator(
            baseURL: baseURL,
            store: MemoryGrantStore(),
            session: stubSession(),
            now: { now }
        )

        await XCTAssertThrowsManaged(.invalidGrant) {
            try await coordinator.install(
                self.credentials(
                    receivedAt: now,
                    scopes: [.polish, .assistant, .agent]
                )
            )
        }
        await XCTAssertThrowsManaged(.invalidGrant) {
            try await coordinator.install(
                self.credentials(
                    receivedAt: now,
                    scopes: [.polish, .assistant, .asr]
                )
            )
        }
    }

    func testOOBERefreshUsesDedicatedEndpoint() async throws {
        let now = Date(timeIntervalSince1970: 3_790)
        let practice = OOBEPracticeSession(
            sessionID: UUID(),
            expectedFeature: .voiceInput,
            startedAt: now.addingTimeInterval(-10),
            expiresAt: now.addingTimeInterval(300)
        )
        let store = MemoryGrantStore(
            credentials(
                accessToken: "expired",
                accessExpiresAt: now.addingTimeInterval(-1),
                receivedAt: now.addingTimeInterval(-600),
                scopes: [.polish, .assistant]
            )
        )
        GatewayStub.shared.enqueue(
            200,
            grantJSON(
                access: "refreshed",
                refresh: "rotated",
                scopes: ["polish", "ai"]
            )
        )
        let coordinator = OOBEGatewayGrantCoordinator(
            baseURL: baseURL,
            store: store,
            session: stubSession(),
            now: { now },
            practiceSession: { practice }
        )

        let token = try await coordinator.accessToken(
            for: .polish,
            feature: .voiceInput
        )

        XCTAssertEqual(token, "refreshed")
        XCTAssertEqual(
            GatewayStub.shared.requests().first?.url?.path,
            "/v1/oobe/grants/refresh"
        )
    }

    func testOOBEClientRejectsMissingFeatureWithoutUsingAccountGrant() async {
        let now = Date(timeIntervalSince1970: 3_800)
        let store = MemoryGrantStore(credentials(accessToken: "account-access", receivedAt: now))
        let client = ManagedLLMClient(
            capability: .polish,
            requestPurpose: .oobe,
            grants: makeCoordinator(store: store, now: { now }),
            baseURL: baseURL,
            session: stubSession()
        )

        await XCTAssertThrowsManaged(
            .server(code: "missing_oobe_feature", status: 400, requestId: "fixed")
        ) {
            _ = try await ManagedLLMClient(
                capability: client.capability,
                requestPurpose: client.requestPurpose,
                grants: makeCoordinator(store: store, now: { now }),
                baseURL: self.baseURL,
                session: self.stubSession(),
                requestId: { "fixed" }
            ).polish("input", systemPrompt: "context")
        }
        XCTAssertTrue(GatewayStub.shared.requests().isEmpty)
    }

    func testManagedTaskKindWireValuesMatchServerContract() {
        XCTAssertEqual(
            ManagedGatewayTaskKind.allCases.map(\.rawValue),
            [
                "dictation_polish",
                "translation",
                "edit_last_input",
                "ai_question",
                "current_information_question",
                "clipboard_transform",
                "custom_skill",
                "agent_planning"
            ]
        )
    }

    func testConversationMappingAndRawTextFallback() async throws {
        let now = Date(timeIntervalSince1970: 4_000)
        let store = MemoryGrantStore(credentials(accessToken: "access", receivedAt: now))
        GatewayStub.shared.enqueue(200, Data("raw answer".utf8), contentType: "text/plain")
        let client = ManagedLLMClient(
            capability: .assistant,
            grants: makeCoordinator(store: store, now: { now }),
            baseURL: baseURL,
            session: stubSession()
        )

        let answer = try await client.complete(
            messages: [.system("system"), .user("old"), .assistant("reply"), .user("latest")],
            timeout: nil,
            options: .polishDefault
        )

        XCTAssertEqual(answer, "raw answer")
        let request = try XCTUnwrap(GatewayStub.shared.requests().first)
        let body = try jsonBody(request)
        XCTAssertEqual(body["input"] as? String, "latest")
        XCTAssertEqual(body["context"] as? String, "system:\nsystem\n\nuser:\nold\n\nassistant:\nreply")
    }

    func testUnauthorizedRequestRefreshesOnceWithSameRequestId() async throws {
        let now = Date(timeIntervalSince1970: 5_000)
        let oldRefresh = "old-refresh-value-123456789012345678901"
        let store = MemoryGrantStore(
            credentials(accessToken: "old-access", refreshToken: oldRefresh, receivedAt: now)
        )
        GatewayStub.shared.enqueue(401, errorJSON("unauthorized"))
        GatewayStub.shared.enqueue(200, grantJSON(access: "new-access", refresh: "new-refresh"))
        GatewayStub.shared.enqueue(200, Data(#"{"output_text":"done"}"#.utf8))
        let client = ManagedLLMClient(
            capability: .polish,
            grants: makeCoordinator(store: store, now: { now }),
            baseURL: baseURL,
            session: stubSession(),
            requestId: { "logical-request" }
        )

        let result = try await client.polish("input", systemPrompt: "prompt")
        XCTAssertEqual(result, "done")
        let requests = GatewayStub.shared.requests()
        XCTAssertEqual(requests.count, 3)
        XCTAssertEqual(requests[0].value(forHTTPHeaderField: "Authorization"), "Bearer old-access")
        XCTAssertEqual(requests[1].url?.path, "/v1/gateway/grants/refresh")
        XCTAssertEqual(requests[2].value(forHTTPHeaderField: "Authorization"), "Bearer new-access")
        XCTAssertEqual(
            requests[0].value(forHTTPHeaderField: "X-Request-ID"),
            requests[2].value(forHTTPHeaderField: "X-Request-ID")
        )
        XCTAssertEqual(
            requests[1].value(forHTTPHeaderField: "Idempotency-Key"),
            GatewayGrantCoordinator.refreshTokenIdempotencyKey(oldRefresh)
        )
    }

    func testCreditAndInvalidRefreshErrorsAreMapped() async throws {
        let now = Date(timeIntervalSince1970: 6_000)
        let creditStore = MemoryGrantStore(credentials(accessToken: "access", receivedAt: now))
        GatewayStub.shared.enqueue(402, errorJSON("insufficient_credits"))
        let creditClient = ManagedLLMClient(
            capability: .assistant,
            grants: makeCoordinator(store: creditStore, now: { now }),
            baseURL: baseURL,
            session: stubSession()
        )
        await XCTAssertThrowsManaged(.insufficientCredits) {
            _ = try await creditClient.complete(
                messages: [.user("question")],
                timeout: nil,
                options: .polishDefault
            )
        }

        let expiredStore = MemoryGrantStore(
            credentials(
                accessToken: "expired",
                accessExpiresAt: now.addingTimeInterval(-1),
                receivedAt: now.addingTimeInterval(-600)
            )
        )
        GatewayStub.shared.enqueue(401, errorJSON("invalid_gateway_refresh"))
        let coordinator = makeCoordinator(store: expiredStore, now: { now })
        await XCTAssertThrowsManaged(.invalidGrant) {
            _ = try await coordinator.accessToken(for: .assistant)
        }
        let remaining = await expiredStore.value()
        XCTAssertNil(remaining)
    }

    func testProviderFailuresAreMappedToDistinctManagedErrors() {
        XCTAssertEqual(
            ManagedGatewayHTTP.error(
                data: errorJSON("provider_unavailable"),
                status: 503,
                requestId: "provider-unavailable"
            ),
            .providerUnavailable(requestId: "request-id")
        )
        XCTAssertEqual(
            ManagedGatewayHTTP.error(
                data: errorJSON("provider_rate_limited"),
                status: 503,
                requestId: "provider-rate-limited"
            ),
            .providerRateLimited(requestId: "request-id")
        )
        XCTAssertEqual(
            ManagedGatewayHTTP.error(
                data: errorJSON("provider_timeout"),
                status: 504,
                requestId: "provider-timeout"
            ),
            .providerTimeout(requestId: "request-id")
        )
        XCTAssertEqual(
            ManagedGatewayHTTP.error(
                data: errorJSON("provider_invalid_response"),
                status: 502,
                requestId: "provider-invalid"
            ),
            .providerFailure(requestId: "request-id")
        )
        XCTAssertEqual(
            ManagedGatewayHTTP.error(
                data: errorJSON("internal_failure"),
                status: 500,
                requestId: "internal-failure"
            ),
            .internalFailure(requestId: "request-id")
        )
    }

    func testTransportTimeoutMapsToManagedTimeout() async {
        let now = Date(timeIntervalSince1970: 6_500)
        let store = MemoryGrantStore(credentials(accessToken: "access", receivedAt: now))
        GatewayStub.shared.enqueue(
            0,
            Data(),
            transportError: .timedOut
        )
        let client = ManagedLLMClient(
            capability: .polish,
            grants: makeCoordinator(store: store, now: { now }),
            baseURL: baseURL,
            session: stubSession()
        )

        await XCTAssertThrowsManaged(.timeout) {
            _ = try await client.polish("input", systemPrompt: "prompt", timeout: 0.1)
        }
    }

    func testCancellationAndStreaming() async throws {
        let now = Date(timeIntervalSince1970: 7_000)
        let store = MemoryGrantStore(credentials(accessToken: "access", receivedAt: now))
        let coordinator = makeCoordinator(store: store, now: { now })
        GatewayStub.shared.enqueue(200, Data(#"{"output_text":"late"}"#.utf8), delay: 5)
        let cancelledClient = ManagedLLMClient(
            capability: .polish,
            grants: coordinator,
            baseURL: baseURL,
            session: stubSession()
        )
        let task = Task { try await cancelledClient.polish("input", systemPrompt: "prompt") }
        try await Task.sleep(for: .milliseconds(50))
        task.cancel()
        do {
            _ = try await task.value
            XCTFail("expected cancellation")
        } catch let error as LLMError {
            XCTAssertEqual(error, .cancelled)
        } catch is CancellationError {
            // Foundation may surface native cancellation before client mapping.
        }

        let sse = """
        data: {"choices":[{"delta":{"reasoning_content":"hidden","content":"Hello"}}]}

        data: {"type":"response.output_text.delta","delta":" world"}

        data: [DONE]

        """
        GatewayStub.shared.enqueue(200, Data(sse.utf8), contentType: "text/event-stream")
        let streamClient = ManagedLLMClient(
            capability: .assistant,
            grants: coordinator,
            baseURL: baseURL,
            session: stubSession()
        )
        var answer = ""
        for try await event in streamClient.completeStreaming(
            messages: [.user("question")],
            timeout: 2,
            options: .polishDefault
        ) {
            if case .delta(let value) = event { answer += value }
        }
        XCTAssertEqual(answer, "Hello world")
        let streamRequest = try XCTUnwrap(GatewayStub.shared.requests().last)
        XCTAssertEqual(try jsonBody(streamRequest)["stream"] as? Bool, true)
    }

    private func makeCoordinator(
        store: MemoryGrantStore,
        now: @escaping @Sendable () -> Date = Date.init
    ) -> GatewayGrantCoordinator {
        GatewayGrantCoordinator(
            baseURL: baseURL,
            store: store,
            session: stubSession(),
            now: now,
            accountAccessPolicy: UnrestrictedManagedGatewayAccountAccessPolicy()
        )
    }

    private func stubSession() -> URLSession {
        let config = URLSessionConfiguration.ephemeral
        config.protocolClasses = [GatewayURLProtocol.self]
        return URLSession(configuration: config)
    }

    private func credentials(
        accessToken: String = "access-token",
        refreshToken: String = "refresh-token-value-12345678901234567890",
        accessExpiresAt: Date = Date(timeIntervalSince1970: 20_000),
        refreshExpiresAt: Date = Date(timeIntervalSince1970: 40_000),
        receivedAt: Date = Date(timeIntervalSince1970: 10_000),
        scopes: Set<ManagedGatewayCapability> = [.polish, .assistant, .agent]
    ) -> ManagedGatewayGrantCredentials {
        ManagedGatewayGrantCredentials(
            grantId: "11111111-1111-1111-1111-111111111111",
            scopes: scopes,
            accessToken: accessToken,
            accessExpiresAt: accessExpiresAt,
            refreshToken: refreshToken,
            refreshExpiresAt: refreshExpiresAt,
            receivedAt: receivedAt
        )
    }

    private func grantJSON(
        access: String,
        refresh: String,
        scopes: [String] = ["polish", "ai", "agent"]
    ) -> Data {
        let encodedScopes = scopes.map { "\"\($0)\"" }.joined(separator: ",")
        return Data(
            """
            {"grantId":"11111111-1111-1111-1111-111111111111",
             "scopes":[\(encodedScopes)],"accessToken":"\(access)",
             "accessExpiresAt":"2030-01-01T00:05:00Z","refreshToken":"\(refresh)",
             "refreshExpiresAt":"2030-01-02T00:00:00Z"}
            """.utf8
        )
    }

    private func errorJSON(_ code: String) -> Data {
        Data(#"{"code":"\#(code)","message":"error","requestId":"request-id"}"#.utf8)
    }

    private func jsonBody(_ request: URLRequest) throws -> [String: Any] {
        try XCTUnwrap(
            JSONSerialization.jsonObject(with: XCTUnwrap(request.httpBody)) as? [String: Any]
        )
    }

    private func XCTAssertThrowsManaged(
        _ expected: ManagedGatewayError,
        operation: () async throws -> Void
    ) async {
        do {
            try await operation()
            XCTFail("expected \(expected)")
        } catch let error as ManagedGatewayError {
            XCTAssertEqual(error, expected)
        } catch {
            XCTFail("unexpected error: \(error)")
        }
    }
}

private actor MemoryGrantStore: GatewayGrantCredentialStore {
    private var credentials: ManagedGatewayGrantCredentials?

    init(_ credentials: ManagedGatewayGrantCredentials? = nil) {
        self.credentials = credentials
    }

    func load() async throws -> ManagedGatewayGrantCredentials? { credentials }
    func save(_ credentials: ManagedGatewayGrantCredentials) async throws { self.credentials = credentials }
    func delete() async throws { credentials = nil }
    func value() -> ManagedGatewayGrantCredentials? { credentials }
}

private final class LockedValues: @unchecked Sendable {
    private let lock = NSLock()
    private var values: [String]

    init(_ values: [String]) { self.values = values }

    func next() -> String {
        lock.withLock { values.isEmpty ? UUID().uuidString : values.removeFirst() }
    }
}

private final class GatewayStub: @unchecked Sendable {
    struct Response {
        let status: Int
        let body: Data
        let contentType: String
        let delay: TimeInterval
        let transportError: URLError.Code?
    }

    static let shared = GatewayStub()
    private let lock = NSLock()
    private var responses: [Response] = []
    private var captured: [URLRequest] = []

    func enqueue(
        _ status: Int,
        _ body: Data,
        contentType: String = "application/json",
        delay: TimeInterval = 0,
        transportError: URLError.Code? = nil
    ) {
        lock.withLock {
            responses.append(
                Response(
                    status: status,
                    body: body,
                    contentType: contentType,
                    delay: delay,
                    transportError: transportError
                )
            )
        }
    }

    func take(_ request: URLRequest) -> Response {
        lock.withLock {
            captured.append(Self.materialize(request))
            return responses.isEmpty
                ? Response(
                    status: 500,
                    body: Data(),
                    contentType: "application/json",
                    delay: 0,
                    transportError: nil
                )
                : responses.removeFirst()
        }
    }

    func requests() -> [URLRequest] { lock.withLock { captured } }

    func reset() {
        lock.withLock {
            responses.removeAll()
            captured.removeAll()
        }
    }

    private static func materialize(_ source: URLRequest) -> URLRequest {
        var request = source
        guard request.httpBody == nil, let stream = request.httpBodyStream else { return request }
        stream.open()
        defer { stream.close() }
        var data = Data()
        let buffer = UnsafeMutablePointer<UInt8>.allocate(capacity: 4_096)
        defer { buffer.deallocate() }
        while stream.hasBytesAvailable {
            let count = stream.read(buffer, maxLength: 4_096)
            if count <= 0 { break }
            data.append(buffer, count: count)
        }
        request.httpBody = data
        return request
    }
}

private final class GatewayURLProtocol: URLProtocol, @unchecked Sendable {
    private let lock = NSLock()
    private var stopped = false

    override static func canInit(with request: URLRequest) -> Bool { true }
    override static func canonicalRequest(for request: URLRequest) -> URLRequest { request }

    override func startLoading() {
        let response = GatewayStub.shared.take(request)
        DispatchQueue.global().asyncAfter(deadline: .now() + response.delay) { [weak self] in
            guard let self, !self.lock.withLock({ self.stopped }) else { return }
            if let errorCode = response.transportError {
                self.client?.urlProtocol(self, didFailWithError: URLError(errorCode))
                return
            }
            let http = HTTPURLResponse(
                url: self.request.url!,
                statusCode: response.status,
                httpVersion: "HTTP/1.1",
                headerFields: ["Content-Type": response.contentType]
            )!
            self.client?.urlProtocol(self, didReceive: http, cacheStoragePolicy: .notAllowed)
            self.client?.urlProtocol(self, didLoad: response.body)
            self.client?.urlProtocolDidFinishLoading(self)
        }
    }

    override func stopLoading() {
        lock.withLock { stopped = true }
    }
}