import Foundation
import Testing

@testable import NucleicCore

@Suite("Container network recovery")
struct ContainerNetworkRecoveryTests {
    @Test func waitsForEveryActiveContainerBeforeResettingSharedNetwork() async throws {
        let engine = NetworkRecoveryTestEngine(dnsProbe: .lookupFailed("EAI_AGAIN"))
        let manager = ContainerManager(engine: engine)
        let first = Self.spec("network-first")
        let second = Self.spec("network-second")
        _ = try await manager.ensureRunning(first)
        _ = try await manager.ensureRunning(second)

        let recovery = Task {
            await manager.requestNetworkRecovery(name: first.name, hostname: "chatgpt.com")
        }
        // Let recovery enter its idle wait, then drain only the affected container. Networking is
        // shared, so the unrelated active turn must still prevent the reset.
        try? await Task.sleep(nanoseconds: 250_000_000)
        await manager.finished(name: first.name)
        try? await Task.sleep(nanoseconds: 250_000_000)
        let resetWhileSecondWasActive = await engine.networkResetCount()
        #expect(resetWhileSecondWasActive == 0)

        await manager.finished(name: second.name)
        await recovery.value
        let finalResetCount = await engine.networkResetCount()
        let runningAfterReset = await engine.runningCount()
        #expect(finalResetCount == 1)
        #expect(runningAfterReset == 0)
    }

    @Test func responsiveReprobeCancelsNetworkReset() async throws {
        let engine = NetworkRecoveryTestEngine(dnsProbe: .responsive)
        let manager = ContainerManager(engine: engine)
        let spec = Self.spec("network-healed")
        _ = try await manager.ensureRunning(spec)
        await manager.finished(name: spec.name)

        await manager.requestNetworkRecovery(name: spec.name, hostname: "chatgpt.com")

        let resetCount = await engine.networkResetCount()
        let probedHostnames = await engine.probedHostnames()
        #expect(resetCount == 0)
        #expect(probedHostnames == ["chatgpt.com"])
    }

    private static func spec(_ name: String) -> ContainerSpec {
        ContainerSpec(
            name: name, image: "test", mounts: [], workdir: "/work", env: [:],
            idleTimeout: 60, claudeHomeStaging: "/staging", claudeHomeWritable: "/home")
    }
}

/// Minimal mechanism fake for exercising `ContainerManager`'s global lifecycle policy. It keeps
/// enough state to prove the reset is delayed until all activity drains and that a healthy
/// last-chance DNS lookup avoids needless churn.
private actor NetworkRecoveryTestEngine: SandboxEngine {
    private enum Unexpected: Error { case call }

    private var running: Set<String> = []
    private var resetCount = 0
    private var hostnames: [String] = []
    private let dnsProbe: SandboxDNSProbe

    init(dnsProbe: SandboxDNSProbe) {
        self.dnsProbe = dnsProbe
    }

    func ensureRunning(_ spec: ContainerSpec) async throws
        -> (name: String, hostGateway: String)
    {
        running.insert(spec.name)
        return (spec.name, "192.168.64.1")
    }

    func exec(
        name: String, workdir: String, env: [String: String], argv: [String],
        uid: Int?, gid: Int?, memoryLimitBytes: UInt64?
    ) async throws -> any ProcessHandle {
        throw Unexpected.call
    }

    func runCapturing(
        name: String, argv: [String], env: [String: String], uid: Int, gid: Int
    ) async -> (exitCode: Int32, output: String)? {
        nil
    }

    func stop(name: String) async {
        running.remove(name)
    }

    func restart(name: String) async -> Bool {
        running.insert(name)
        return true
    }

    func remove(name: String) async -> Bool {
        running.remove(name)
        return true
    }

    func isRunning(_ name: String) async -> Bool {
        running.contains(name)
    }

    func containerExists(_ name: String) async -> Bool {
        running.contains(name)
    }

    func list() async -> [String] {
        Array(running)
    }

    func sampleResourceUsage(name: String) async -> ContainerResourceSample? { nil }
    func currentDownloadProgress() async -> ContainerDownloadProgress? { nil }

    func probeControlPlane(name: String) async -> ContainerEngine.ControlPlaneProbe {
        .responsive
    }

    func probeDNS(name: String, hostname: String) async -> SandboxDNSProbe {
        hostnames.append(hostname)
        return dnsProbe
    }

    func resetNetwork() async {
        resetCount += 1
        running.removeAll()
    }

    func diagnoseKill(name: String) async -> ContainerKillDiagnosis {
        ContainerKillDiagnosis(
            containerRunning: running.contains(name), oomKills: nil, sample: nil)
    }

    func reconcileDisk(keepNames: Set<String>) async {}
    func removeDefaultRootfs() async {}
    func checkDefaultImageUpdate() async -> ImageUpdateStatus { .idle }
    func reclaimMemoryNow() async {}

    func networkResetCount() -> Int { resetCount }
    func runningCount() -> Int { running.count }
    func probedHostnames() -> [String] { hostnames }
}
