import Foundation

#if canImport(Darwin)
import Darwin
#elseif canImport(Musl)
import Musl
#elseif canImport(Glibc)
import Glibc
#elseif os(Windows)
import WinSDK
#endif

/// A peer's UDP endpoint as the punch layer sees it — one IP literal + port, family-tagged.
/// String-backed (not `sockaddr` structs) because these travel inside Noise as candidate
/// lists (`DirectCandidate`) and get compared/logged; conversion to raw sockaddr happens
/// only at the send/recv boundary in `UDPSocket`.
public enum UDPEndpoint: Sendable, Hashable, Codable {
    case v4(ip: String, port: UInt16)
    case v6(ip: String, port: UInt16)

    public var ip: String {
        switch self {
        case .v4(let ip, _), .v6(let ip, _): ip
        }
    }

    public var port: UInt16 {
        switch self {
        case .v4(_, let port), .v6(_, let port): port
        }
    }

    public var isIPv6: Bool {
        if case .v6 = self { return true }
        return false
    }

    public var description: String {
        switch self {
        case .v4(let ip, let port): "\(ip):\(port)"
        case .v6(let ip, let port): "[\(ip)]:\(port)"
        }
    }
}

/// A plain BSD-sockets UDP socket — deliberately *not* Network.framework, for two reasons: the
/// punch phase needs precise control of the local port (the QUIC leg later rebinds to it,
/// see `SYNC_PROTOCOL` §3.3), and the whole pre-QUIC core must build on Linux so the NAT
/// simulation (`scripts/nat-sim`) exercises the real STUN/punch code, not a re-implementation.
///
/// Winsock differs from POSIX in handle type, error reporting and option encoding; every one
/// of those differences lives in ``DirectSock`` so this file reads the same on all platforms.
///
/// One socket per address family: the v4 socket carries STUN + v4 punching, an optional v6
/// sibling carries plain v6 reachability probes. Sockets are non-blocking; `receive` polls in
/// short slices so a cancelled Task exits promptly instead of hanging in `recvfrom`.
public final class UDPSocket: @unchecked Sendable {
    public enum SocketError: Error, Equatable {
        case createFailed(Int32)
        case bindFailed(Int32)
        case sendFailed(Int32)
        case closed
        case badAddress(String)
    }

    private let fd: DirectSocketHandle
    private let lock = NSLock()
    private var isClosed = false
    public let isIPv6: Bool
    /// The OS-assigned (or requested) local port — the value that becomes the candidate port
    /// and, after the punch, the QUIC leg's `requiredLocalEndpoint`.
    public let localPort: UInt16

    /// `bindAddress` binds to a specific local IP instead of the wildcard — used by the NAT
    /// sim's multi-homed STUN server so its replies leave from the queried address (a real
    /// STUN client drops a reply whose source doesn't match where it sent the request, and a
    /// NAT's conntrack drops it outright).
    public init(port: UInt16 = 0, ipv6: Bool = false, bindAddress: String? = nil) throws {
        self.isIPv6 = ipv6
        let family: Int32 = ipv6 ? AF_INET6 : AF_INET
        let fd = DirectSock.makeUDP(family: family)
        guard DirectSock.isValid(fd) else { throw SocketError.createFailed(DirectSock.lastError()) }

        // Non-blocking: receive() polls in slices; a blocking recvfrom would pin a
        // concurrency thread for the whole punch budget.
        DirectSock.setNonBlocking(fd)
        DirectSock.setIntOption(fd, level: SOL_SOCKET, name: SO_REUSEADDR, value: 1)

        var bindResult: Int32 = -1
        if ipv6 {
            var addr = sockaddr_in6()
            addr.sin6_family = sa_family_t(AF_INET6)
            addr.sin6_port = port.bigEndian
            addr.sin6_addr = in6addr_any
            // v6-only: the v4 socket handles v4 traffic, so a dual-stack v6 socket would
            // just double-deliver punches on mapped addresses.
            DirectSock.setIntOption(
                fd, level: DirectSock.ipv6OptionLevel, name: IPV6_V6ONLY, value: 1)
            bindResult = withUnsafePointer(to: &addr) {
                $0.withMemoryRebound(to: sockaddr.self, capacity: 1) {
                    DirectSock.bindSocket(fd, $0, MemoryLayout<sockaddr_in6>.size)
                }
            }
        } else {
            var addr = sockaddr_in()
            addr.sin_family = sa_family_t(AF_INET)
            addr.sin_port = port.bigEndian
            if let bindAddress, DirectSock.pton(family: AF_INET, bindAddress, into: &addr.sin_addr) {
                // bound to the given IP
            } else {
                addr.sin_addr = DirectSock.anyIPv4Address()
            }
            bindResult = withUnsafePointer(to: &addr) {
                $0.withMemoryRebound(to: sockaddr.self, capacity: 1) {
                    DirectSock.bindSocket(fd, $0, MemoryLayout<sockaddr_in>.size)
                }
            }
        }
        guard bindResult == 0 else {
            let err = DirectSock.lastError()
            DirectSock.closeHandle(fd)
            throw SocketError.bindFailed(err)
        }

        // Read back the OS-assigned port when we bound to 0.
        var storage = sockaddr_storage()
        var len = MemoryLayout<sockaddr_storage>.size
        let boundPort: UInt16 = withUnsafeMutablePointer(to: &storage) {
            $0.withMemoryRebound(to: sockaddr.self, capacity: 1) { sa -> UInt16 in
                guard DirectSock.localName(fd, sa, &len) == 0 else { return port }
                if ipv6 {
                    return sa.withMemoryRebound(to: sockaddr_in6.self, capacity: 1) {
                        UInt16(bigEndian: $0.pointee.sin6_port)
                    }
                }
                return sa.withMemoryRebound(to: sockaddr_in.self, capacity: 1) {
                    UInt16(bigEndian: $0.pointee.sin_port)
                }
            }
        }

        self.fd = fd
        self.localPort = boundPort
    }

    deinit { closeSocket() }

    public func closeSocket() {
        lock.lock()
        defer { lock.unlock() }
        guard !isClosed else { return }
        isClosed = true
        DirectSock.closeHandle(fd)
    }

    /// Sync by design: `NSLock` is unavailable from async contexts, so the async `receive`
    /// loop reads the flag through this hop.
    private var closed: Bool {
        lock.lock()
        defer { lock.unlock() }
        return isClosed
    }

    public func send(_ data: Data, to endpoint: UDPEndpoint) throws {
        guard !closed else { throw SocketError.closed }
        guard endpoint.isIPv6 == isIPv6 else { throw SocketError.badAddress(endpoint.description) }

        var storage = sockaddr_storage()
        let saLen = try Self.fill(&storage, from: endpoint)
        let sent = data.withUnsafeBytes { buf in
            withUnsafePointer(to: &storage) {
                $0.withMemoryRebound(to: sockaddr.self, capacity: 1) { sa in
                    DirectSock.sendPacket(fd, buf.baseAddress, data.count, to: sa, saLen)
                }
            }
        }
        // EWOULDBLOCK on a UDP socket means a transient full buffer; punch/STUN datagrams
        // are retransmitted by their own cadence, so dropping one silently is correct.
        if sent < 0 {
            let err = DirectSock.lastError()
            if !DirectSock.isWouldBlock(err) { throw SocketError.sendFailed(err) }
        }
    }

    /// One datagram, or nil at the deadline. Polls in ≤50 ms slices so Task cancellation
    /// and short punch cadences stay responsive without a blocking recvfrom.
    public func receive(
        deadline: ContinuousClock.Instant,
        clock: ContinuousClock = ContinuousClock()
    ) async throws -> (data: Data, from: UDPEndpoint)? {
        while true {
            try Task.checkCancellation()
            guard !closed else { throw SocketError.closed }

            if let packet = tryReceive() { return packet }

            let now = clock.now
            guard now < deadline else { return nil }
            let slice = min(Duration.milliseconds(50), deadline - now)
            try await clock.sleep(for: slice)
        }
    }

    /// Non-blocking read of one datagram, nil when nothing is queued.
    public func tryReceive() -> (data: Data, from: UDPEndpoint)? {
        var buffer = [UInt8](repeating: 0, count: 2048)
        var storage = sockaddr_storage()
        var len = MemoryLayout<sockaddr_storage>.size
        let count = buffer.withUnsafeMutableBytes { buf in
            withUnsafeMutablePointer(to: &storage) {
                $0.withMemoryRebound(to: sockaddr.self, capacity: 1) { sa in
                    DirectSock.receivePacket(fd, buf.baseAddress, buf.count, from: sa, &len)
                }
            }
        }
        guard count > 0, let from = Self.endpoint(from: &storage) else { return nil }
        return (Data(buffer[0..<count]), from)
    }

    // MARK: sockaddr conversion

    /// Fills `storage` and returns the significant length, for the send path.
    static func fill(_ storage: inout sockaddr_storage, from endpoint: UDPEndpoint) throws -> Int {
        switch endpoint {
        case .v4(let ip, let port):
            var addr = sockaddr_in()
            addr.sin_family = sa_family_t(AF_INET)
            addr.sin_port = port.bigEndian
            guard DirectSock.pton(family: AF_INET, ip, into: &addr.sin_addr) else {
                throw SocketError.badAddress(ip)
            }
            withUnsafeMutablePointer(to: &storage) {
                $0.withMemoryRebound(to: sockaddr_in.self, capacity: 1) { $0.pointee = addr }
            }
            return MemoryLayout<sockaddr_in>.size
        case .v6(let ip, let port):
            var addr = sockaddr_in6()
            addr.sin6_family = sa_family_t(AF_INET6)
            addr.sin6_port = port.bigEndian
            guard DirectSock.pton(family: AF_INET6, ip, into: &addr.sin6_addr) else {
                throw SocketError.badAddress(ip)
            }
            withUnsafeMutablePointer(to: &storage) {
                $0.withMemoryRebound(to: sockaddr_in6.self, capacity: 1) { $0.pointee = addr }
            }
            return MemoryLayout<sockaddr_in6>.size
        }
    }

    static func endpoint(from storage: inout sockaddr_storage) -> UDPEndpoint? {
        let family = Int32(storage.ss_family)
        if family == AF_INET {
            return withUnsafePointer(to: &storage) {
                $0.withMemoryRebound(to: sockaddr_in.self, capacity: 1) { sa -> UDPEndpoint? in
                    guard
                        let ip = withUnsafePointer(to: sa.pointee.sin_addr, {
                            DirectSock.ntop(family: AF_INET, $0)
                        })
                    else { return nil }
                    return .v4(ip: ip, port: UInt16(bigEndian: sa.pointee.sin_port))
                }
            }
        }
        if family == AF_INET6 {
            return withUnsafePointer(to: &storage) {
                $0.withMemoryRebound(to: sockaddr_in6.self, capacity: 1) { sa -> UDPEndpoint? in
                    guard
                        let ip = withUnsafePointer(to: sa.pointee.sin6_addr, {
                            DirectSock.ntop(family: AF_INET6, $0)
                        })
                    else { return nil }
                    return .v6(ip: ip, port: UInt16(bigEndian: sa.pointee.sin6_port))
                }
            }
        }
        return nil
    }
}
