import Foundation
import NucleicProtocol
import Testing

@testable import NucleicCore

/// Scheduling contracts of the mesh-wide AFM queue (docs/COVALENCE_RUNNER.md §5, item 6):
/// priority→device-class placement, the Apple Silicon power bias, cross-device parallelism,
/// deadlines, disconnect re-dispatch, and result correlation.
@Suite struct MeshIntelligenceQueueTests {
    /// A scripted worker: records what it's handed and, when `reply` is set, answers through
    /// the queue like a real device would.
    private actor Worker {
        private(set) var received: [WireIntelligenceRequest] = []
        private let reply: (@Sendable (WireIntelligenceRequest) -> WireIntelligenceResult?)?
        private let queue: MeshIntelligenceQueue
        private let deviceID: String

        init(
            deviceID: String, queue: MeshIntelligenceQueue,
            reply: (@Sendable (WireIntelligenceRequest) -> WireIntelligenceResult?)? = nil
        ) {
            self.deviceID = deviceID
            self.queue = queue
            self.reply = reply
        }

        func handle(_ request: WireIntelligenceRequest) async {
            received.append(request)
            if let result = reply?(request) {
                await queue.receiveResult(result, from: deviceID)
            }
        }
        func requests() -> [WireIntelligenceRequest] { received }
    }

    private func register(
        _ worker: Worker, as deviceID: String, on queue: MeshIntelligenceQueue,
        connectionID: String = "c1", deviceClass: IntelligenceDeviceClass, chip: String? = nil
    ) async {
        await queue.registerWorker(
            deviceID: deviceID, connectionID: connectionID,
            profile: IntelligenceWorkerProfile(deviceClass: deviceClass, chip: chip),
            send: { request in await worker.handle(request) })
    }

    private func request(
        _ id: String, priority: IntelligencePriority, deadline: Double = 3
    ) -> WireIntelligenceRequest {
        WireIntelligenceRequest(
            id: id, kind: .summarizeBashCommand, inputs: ["swift test"],
            deadlineSeconds: deadline, priority: priority)
    }

    /// An echo reply: answers every request with its own id as the output.
    private static let echo: @Sendable (WireIntelligenceRequest) -> WireIntelligenceResult? = {
        WireIntelligenceResult(id: $0.id, outputs: ["echo:\($0.id)"])
    }

    @Test func interactiveNeverRunsOnMobile() async {
        // Only a phone is connected: a top-tier job must fail fast (nil → instant heuristic
        // fallback), and the phone must never even see it.
        let queue = MeshIntelligenceQueue()
        let phone = Worker(deviceID: "phone", queue: queue, reply: Self.echo)
        await register(phone, as: "phone", on: queue, deviceClass: .mobile, chip: "Apple A18 Pro")

        let result = await queue.submit(request("job-1", priority: .interactive))
        #expect(result == nil)
        #expect(await phone.requests().isEmpty)
        // Same for the bash tier — both top tiers are desktop-only.
        #expect(await queue.submit(request("job-2", priority: .bashSummary)) == nil)
        #expect(await phone.requests().isEmpty)
    }

    @Test func lowTiersRunOnMobile() async {
        let queue = MeshIntelligenceQueue()
        let phone = Worker(deviceID: "phone", queue: queue, reply: Self.echo)
        await register(phone, as: "phone", on: queue, deviceClass: .mobile, chip: "Apple A18 Pro")

        let completion = await queue.submit(request("job-1", priority: .completion))
        #expect(completion?.outputs == ["echo:job-1"])
        let background = await queue.submit(request("job-2", priority: .background))
        #expect(background?.outputs == ["echo:job-2"])
    }

    @Test func highestPriorityLandsOnMostPowerfulChip() async {
        // Two idle desktop-class workers: the interactive job must land on the stronger chip.
        let queue = MeshIntelligenceQueue()
        let base = Worker(deviceID: "m1", queue: queue, reply: Self.echo)
        let max = Worker(deviceID: "m4max", queue: queue, reply: Self.echo)
        await register(base, as: "m1", on: queue, deviceClass: .desktop, chip: "Apple M1")
        await register(max, as: "m4max", on: queue, deviceClass: .desktop, chip: "Apple M4 Max")

        let result = await queue.submit(request("hot", priority: .interactive))
        #expect(result?.outputs == ["echo:hot"])
        #expect(await max.requests().map(\.id) == ["hot"])
        #expect(await base.requests().isEmpty)
    }

    @Test func jobsRunInParallelAcrossWorkers() async {
        // Two workers, two concurrent jobs: each device gets one — the queue fans out rather
        // than serializing everything behind one device.
        let queue = MeshIntelligenceQueue()
        let a = Worker(deviceID: "a", queue: queue)  // records, never answers
        let b = Worker(deviceID: "b", queue: queue)
        await register(a, as: "a", on: queue, deviceClass: .desktop, chip: "Apple M4")
        await register(b, as: "b", on: queue, deviceClass: .desktop, chip: "Apple M1")

        async let first = queue.submit(request("j1", priority: .background, deadline: 0.5))
        async let second = queue.submit(request("j2", priority: .background, deadline: 0.5))
        let dispatched = await poll {
            await a.requests().count + b.requests().count == 2 ? [true] : []
        }
        #expect(dispatched == [true])
        #expect(await a.requests().count == 1)
        #expect(await b.requests().count == 1)
        // Nobody answers — both resolve nil at their deadline.
        #expect(await first == nil)
        #expect(await second == nil)
    }

    @Test func deadlineResolvesNil() async {
        let queue = MeshIntelligenceQueue()
        let silent = Worker(deviceID: "mac", queue: queue)  // takes the job, never answers
        await register(silent, as: "mac", on: queue, deviceClass: .desktop, chip: "Apple M2")

        let started = ContinuousClock.now
        let result = await queue.submit(request("slow", priority: .completion, deadline: 0.2))
        #expect(result == nil)
        #expect(ContinuousClock.now - started < .seconds(2))
        #expect(await silent.requests().map(\.id) == ["slow"])
    }

    @Test func noWorkersFailsFast() async {
        let queue = MeshIntelligenceQueue()
        let started = ContinuousClock.now
        #expect(await queue.submit(request("j", priority: .background, deadline: 5)) == nil)
        // Fail-fast: nowhere near the 5 s deadline.
        #expect(ContinuousClock.now - started < .seconds(1))
    }

    @Test func workerDisconnectRedispatchesToAnother() async {
        let queue = MeshIntelligenceQueue()
        let flaky = Worker(deviceID: "flaky", queue: queue)  // strongest chip; never answers
        await register(flaky, as: "flaky", on: queue, deviceClass: .desktop, chip: "Apple M4 Ultra")

        async let pending = queue.submit(request("j", priority: .completion, deadline: 3))
        _ = await poll { await flaky.requests() }
        #expect(await flaky.requests().map(\.id) == ["j"])

        // The flaky device drops; a healthy one is (already/newly) around to pick the job up.
        let healthy = Worker(deviceID: "healthy", queue: queue, reply: Self.echo)
        await register(healthy, as: "healthy", on: queue, deviceClass: .desktop, chip: "Apple M1")
        await queue.unregisterWorker(deviceID: "flaky", connectionID: "c1")

        let result = await pending
        #expect(result?.outputs == ["echo:j"])
    }

    @Test func staleConnectionTeardownKeepsReplacementRegistered() async {
        // A reconnect replaces the worker (connection c2); the deduped old socket's teardown
        // (c1) must not unregister it — the connectionID guard.
        let queue = MeshIntelligenceQueue()
        let worker = Worker(deviceID: "mac", queue: queue, reply: Self.echo)
        await register(worker, as: "mac", on: queue, connectionID: "c1", deviceClass: .desktop)
        await register(worker, as: "mac", on: queue, connectionID: "c2", deviceClass: .desktop)
        await queue.unregisterWorker(deviceID: "mac", connectionID: "c1")

        let result = await queue.submit(request("j", priority: .background))
        #expect(result?.outputs == ["echo:j"])
    }

    @Test func resultsCorrelateAndStaleOnesDrop() async {
        let queue = MeshIntelligenceQueue()
        let mac = Worker(deviceID: "mac", queue: queue)
        await register(mac, as: "mac", on: queue, deviceClass: .desktop, chip: "Apple M3")

        async let pending = queue.submit(request("j", priority: .background, deadline: 3))
        _ = await poll { await mac.requests() }

        // An answer from a device the job wasn't dispatched to is dropped…
        await queue.receiveResult(
            WireIntelligenceResult(id: "j", outputs: ["forged"]), from: "someone-else")
        // …the real worker's answer resolves the submitter.
        await queue.receiveResult(
            WireIntelligenceResult(id: "j", outputs: ["real"]), from: "mac")
        #expect(await pending?.outputs == ["real"])
    }

    @Test func executorErrorNormalizesToNil() async {
        // A device that can't serve (model off, unknown kind) answers with `error` set — the
        // submitter reads nil and falls back to heuristics.
        let queue = MeshIntelligenceQueue()
        let mac = Worker(
            deviceID: "mac", queue: queue,
            reply: { WireIntelligenceResult(id: $0.id, error: "model unavailable") })
        await register(mac, as: "mac", on: queue, deviceClass: .desktop)
        #expect(await queue.submit(request("j", priority: .background)) == nil)
    }
}
