import Foundation
import Network

/// Measures round-trip time to an endpoint with an unprivileged TCP connect
/// (time from start to `.ready` ≈ one handshake RTT). Results are cached.
///
/// We probe a port we've already seen traffic on, so we're not knocking on
/// closed doors — and the connection carries no data, just the handshake.
public actor LatencyProber {
    private var cache: [String: Double] = [:]

    public init() {}

    public func cached(_ host: String) -> Double? { cache[host] }

    /// Measure (or refresh) RTT in milliseconds. Returns nil if unreachable.
    public func measure(host: String, port: UInt16) async -> Double? {
        let rtt = await Self.tcpHandshakeRTT(host: host, port: port)
        if let rtt { cache[host] = rtt }
        return rtt
    }

    private static func tcpHandshakeRTT(host: String, port: UInt16) async -> Double? {
        guard let nwPort = NWEndpoint.Port(rawValue: port) else { return nil }

        return await withCheckedContinuation { (cont: CheckedContinuation<Double?, Never>) in
            let conn = NWConnection(host: NWEndpoint.Host(host), port: nwPort, using: .tcp)
            let queue = DispatchQueue(label: "latency.probe")
            let start = DispatchTime.now()
            let state = ProbeState(conn: conn, cont: cont)

            conn.stateUpdateHandler = { st in
                switch st {
                case .ready:
                    let ns = DispatchTime.now().uptimeNanoseconds - start.uptimeNanoseconds
                    state.finish(Double(ns) / 1_000_000)
                case .failed, .cancelled:
                    state.finish(nil)
                default:
                    break
                }
            }
            conn.start(queue: queue)
            queue.asyncAfter(deadline: .now() + 2) { state.finish(nil) }   // timeout
        }
    }

    /// Resumes the continuation exactly once, whichever callback fires first.
    private final class ProbeState: @unchecked Sendable {
        private let lock = NSLock()
        private var done = false
        private let conn: NWConnection
        private let cont: CheckedContinuation<Double?, Never>

        init(conn: NWConnection, cont: CheckedContinuation<Double?, Never>) {
            self.conn = conn
            self.cont = cont
        }

        func finish(_ value: Double?) {
            lock.lock()
            if done { lock.unlock(); return }
            done = true
            lock.unlock()
            conn.cancel()
            cont.resume(returning: value)
        }
    }
}