import Foundation
import System

extension FileDescriptor {

    func dup() throws -> FileDescriptor {
        let result = try errnoQ(retryOnInterrupt: false) {
            Foundation.dup(self.rawValue)
        }
        return FileDescriptor(rawValue: result)
    }

    func setNonBlocking() throws {
        let flags = try errnoQ(retryOnInterrupt: false) {
            fcntl(self.rawValue, F_GETFL)
        }
        try errnoQ(retryOnInterrupt: false) {
            fcntl(self.rawValue, F_SETFL, flags | O_NONBLOCK)
        }
    }

    static func socket(_ domain: Int32, _ type: Int32, _ proto: Int32) throws -> FileDescriptor {
        let socket = try errnoQ(retryOnInterrupt: false) {
            Foundation.socket(domain, type, proto)
        }
        return FileDescriptor(rawValue: socket)
    }
    
    static func socketpair(_ domain: Int32, _ type: Int32, _ proto: Int32) throws -> (FileDescriptor, FileDescriptor) {
        var sockets = [Int32](repeating: -1, count: 2)
        try errnoQ(retryOnInterrupt: false) {
            Foundation.socketpair(domain, type, proto, &sockets)
        }
        return (FileDescriptor(rawValue: sockets[0]), FileDescriptor(rawValue: sockets[1]))
    }
    
    func setSocketOption<Value>(_ level: Int32, _ option: Int32, _ value: Value) throws
        where Value: SocketOptionSettable
    {
        try withUnsafeBytes(of: value) { buf in
            _ = try errnoQ(retryOnInterrupt: false) {
                Foundation.setsockopt(self.rawValue, level, option, buf.baseAddress, socklen_t(buf.count))
            }
        }
    }
    
    func getSocketOption<Value>(_ level: Int32, _ option: Int32) throws -> Value
        where Value: SocketOptionGettable
    {
        var result = Value()
        try withUnsafeMutableBytes(of: &result) { buf in
            var len = socklen_t(buf.count)
            _ = try errnoQ(retryOnInterrupt: false) {
                Foundation.getsockopt(self.rawValue, level, option, buf.baseAddress!, &len)
            }
            guard len == buf.count else {
                throw Errno.invalidArgument
            }
        }
        return result
    }

    func bind(_ addr: SockAddr) throws {
        try addr.withSockAddr { sa, saLen in
            _ = try errnoQ(retryOnInterrupt: false) {
                Foundation.bind(self.rawValue, sa, saLen)
            }
        }
    }

    func connect(_ addr: SockAddr, retryOnInterrupt: Bool = true) throws {
        try addr.withSockAddr { sa, saLen in
            _ = try errnoQ(retryOnInterrupt: retryOnInterrupt) {
                Foundation.connect(self.rawValue, sa, saLen)
            }
        }
    }

    func getSocketName() throws -> SockAddr {
        try SockAddr.fromMutableSockAddr { sa, saLenPtr in
            _ = try errnoQ(retryOnInterrupt: false) {
                Foundation.getsockname(self.rawValue, sa, saLenPtr)
            }
        }.1
    }

    func getPeerName() throws -> SockAddr {
        try SockAddr.fromMutableSockAddr { sa, saLenPtr in
            _ = try errnoQ(retryOnInterrupt: false) {
                Foundation.getpeername(self.rawValue, sa, saLenPtr)
            }
        }.1
    }

    func listen(_ backlog: Int32) throws {
        try errnoQ(retryOnInterrupt: false) {
            Foundation.listen(self.rawValue, backlog)
        }
    }

    func accept(retryOnInterrupt: Bool = true) throws -> (newSocket: FileDescriptor, addr: SockAddr) {
        let (newSocket, addr) = try SockAddr.fromMutableSockAddr { sa, saLenPtr in
            try errnoQ(retryOnInterrupt: retryOnInterrupt) {
                Foundation.accept(self.rawValue, sa, saLenPtr)
            }
        }
        return (FileDescriptor(rawValue: newSocket), addr)
    }
    
    func shutdown(_ how: ShutdownHow) throws {
        try errnoQ(retryOnInterrupt: false) {
            Foundation.shutdown(self.rawValue, how.rawValue)
        }
    }
    
    struct ShutdownHow: OptionSet {
        var rawValue: Int32
        static let read = ShutdownHow(rawValue: SHUT_RD)
        static let write = ShutdownHow(rawValue: SHUT_WR)
    }
    
    func send(_ buffer: UnsafeRawBufferPointer, flags: MessageFlags = [], to addr: SockAddr, retryOnInterrupt: Bool = true) throws -> Int {
        try addr.withSockAddr { sa, saLen in
            try errnoQ(retryOnInterrupt: retryOnInterrupt) {
                Foundation.sendto(self.rawValue, buffer.baseAddress, buffer.count, flags.rawValue, sa, saLen)
            }
        }
    }
    
    func receive(_ buffer: UnsafeMutableRawBufferPointer, flags: MessageFlags = [], from addr: inout SockAddr, retryOnInterrupt: Bool = true) throws -> Int {
        let result = try SockAddr.fromMutableSockAddr { sa, saLenPtr in
            try errnoQ(retryOnInterrupt: retryOnInterrupt) {
                Foundation.recvfrom(self.rawValue, buffer.baseAddress, buffer.count, flags.rawValue, sa, saLenPtr)
            }
        }
        addr = result.1
        return result.0
    }
    
    struct MessageFlags: OptionSet {
        var rawValue: Int32
        
        static let outOfBand = MessageFlags(rawValue: MSG_OOB)
        static let peek = MessageFlags(rawValue: MSG_PEEK)
        static let dontRoute = MessageFlags(rawValue: MSG_DONTROUTE)
        static let waitAll = MessageFlags(rawValue: MSG_WAITALL)
        // The flags here are those documented in the man page `sendto` and
        // `recvfrom`.  There are lots of others; if you need ’em, define ’em!
    }
}

protocol SocketOptionSettable {
}

protocol SocketOptionGettable {
    init()
}

extension UInt8: SocketOptionSettable { }
extension UInt8: SocketOptionGettable { }
extension Int32: SocketOptionSettable { }
extension Int32: SocketOptionGettable { }
extension UInt32: SocketOptionSettable { }
extension UInt32: SocketOptionGettable { }
extension timeval: SocketOptionSettable { }
extension timeval: SocketOptionGettable { }
extension ip_mreq: SocketOptionSettable { }
extension ip_mreq: SocketOptionGettable { }
extension ipv6_mreq: SocketOptionSettable { }
extension ipv6_mreq: SocketOptionGettable { }

// MARK: Address

struct SockAddr {

    private var addr: Data
    
    init(addr: Data) {
        // Must have at least the `sa_len` and `sa_family` fields.
        precondition(addr.count >= MemoryLayout.offset(of: \sockaddr.sa_data)!)
        // `SOCK_MAXADDRLEN` is the hard limit on socket address lengths.
        precondition(addr.count <= SOCK_MAXADDRLEN)
        // The length must fit into a `UInt8`, which is the type of `sa_len`.
        precondition(UInt8(exactly: addr.count) != nil)
        self.addr = addr
        // Run this check after setting up `addr` so that we can use
        // our helper.
        //
        // We are expecting `sa_len` to match the address length.  It’s
        // conceivable that this assert might trip because many folks
        // fail to set up `sa_len` properly.  However, I’m leaving it
        // in place to help flush out those problems.
        self.withSockAddr { sa, _ in
            precondition(sa.pointee.sa_len == addr.count)
        }
    }
    
    init?(saPtr: UnsafePointer<sockaddr>) {
        guard saPtr.pointee.sa_len > 0 && saPtr.pointee.sa_family != AF_UNSPEC else { return nil }
        let buf = UnsafeRawBufferPointer.init(start: saPtr, count: Int(saPtr.pointee.sa_len))
        self.init(addr: Data(buf))
    }
    
    var saLength: UInt8 {
        self.withSockAddr { sa, _ in
            sa.pointee.sa_len
        }
    }
    
    var saFamily: sa_family_t {
        self.withSockAddr { sa, _ in
            sa.pointee.sa_family
        }
    }
    
    static let empty: SockAddr = {
        var sa = sockaddr()
        sa.sa_len = UInt8(sockAddrMinSize)
        let saData = withUnsafeBytes(of: &sa) { buf in
            Data(buf.prefix(sockAddrMinSize))
        }
        return SockAddr(addr: saData)
    }()
}

// MARK: Address 'Base Class' Conversion

extension SockAddr {

    static func fromMutableSockAddr<Result>(_ body: (_ sa: UnsafeMutablePointer<sockaddr>, _ saLenPtr: UnsafeMutablePointer<socklen_t>) throws -> Result) rethrows -> (Result, SockAddr) {
        var addr = Data(repeating: 0, count: Int(SOCK_MAXADDRLEN))
        var saLen = socklen_t(addr.count)
        let result = try addr.withUnsafeMutableBytes { buf -> Result in
            let sa = buf.bindMemory(to: sockaddr.self).baseAddress!
            return try body(sa, &saLen)
        }
        assert(saLen <= addr.count)
        let sockAddr = SockAddr(addr: addr.prefix(Int(saLen)))
        return (result, sockAddr)
    }

    func withSockAddr<Result>(_ body: (_ sa: UnsafePointer<sockaddr>, _ saLen: socklen_t) throws -> Result) rethrows -> Result {
        try self.addr.withUnsafeBytes { buf in
            let sa = buf.bindMemory(to: sockaddr.self).baseAddress!
            let saLen = socklen_t(buf.count)
            return try body(sa, saLen)
        }
    }
}

// MARK: Address 'Subclass' Conversion

extension SockAddr {

    static let sockAddrMinSize = MemoryLayout.offset(of: \sockaddr.sa_data)!

    init<SockAddrType>(_ sax: SockAddrType)
        where SockAddrType: SockAddrCompatible
    {
        precondition(MemoryLayout.size(ofValue: sax) >= Self.sockAddrMinSize)
        var tmp = sax
        let addr = withUnsafeBytes(of: &tmp) { buf in
            Data(buf)
        }
        self.init(addr: addr)
    }

    // This is for the common case where you just need to access a field in a
    // fixed-size address, like `sockaddr_in`.
    
    func withSockAddrType<SockAddrType, Return>(_ body: (_ sax: SockAddrType) throws -> Return) rethrows -> Return
        where SockAddrType: SockAddrCompatible
    {
        precondition(MemoryLayout<SockAddrType>.size >= Self.sockAddrMinSize)
        return try self.addr.withUnsafeBytes { buf in
            let sax = buf.bindMemory(to: SockAddrType.self).baseAddress!
            return try body(sax.pointee)
        }
    }

    // This is for the weird case, for example, dealing with `sockaddr_un`,
    // where the actual bytes might extend beyond the fixed size of the
    // `sockaddr` ‘subclass’.

    func withSockAddrTypePtr<SockAddrType, Return>(_ body: (_ sax: UnsafePointer<SockAddrType>) throws -> Return) rethrows -> Return
        where SockAddrType: SockAddrCompatible
    {
        precondition(MemoryLayout<SockAddrType>.size >= Self.sockAddrMinSize)
        var tmp = addr
        return try tmp.withUnsafeMutableBytes { buf in
            let sax = buf.bindMemory(to: SockAddrType.self).baseAddress!
            return try body(sax)
        }
    }
}

protocol SockAddrCompatible {
}

extension sockaddr_storage: SockAddrCompatible {
}

extension sockaddr_in: SockAddrCompatible {
}

extension sockaddr_in6: SockAddrCompatible {
}

extension sockaddr_un: SockAddrCompatible {
}

extension sockaddr_dl: SockAddrCompatible {
}

//extension sockaddr_ctl: SockAddrCompatible {
//}
//
//extension sockaddr_sys: SockAddrCompatible {
//}

// MARK: Adress Text Representation Conversion

extension SockAddr {

    init?(textRepresentation: (address: String, service: String)) {
        var hints = addrinfo()
        hints.ai_flags = AI_NUMERICHOST | AI_NUMERICSERV
        var addrList: UnsafeMutablePointer<addrinfo>? = nil
        guard getaddrinfo(textRepresentation.address, textRepresentation.service, &hints, &addrList) == 0 else {
            return nil
        }
        defer { freeaddrinfo(addrList) }
        guard
            let sa = addrList?.pointee.ai_addr,
            let saLen = addrList?.pointee.ai_addrlen,
            saLen != 0      // Just being paranoid.
        else { return nil }
        let addr = Data(bytes: sa, count: Int(saLen))
        self.init(addr: addr)
    }

    var textRepresentation: (address: String, service: String) {
        self.withSockAddr { sa, saLen -> (String, String) in
            var host = [CChar](repeating: 0, count: Int(NI_MAXHOST))
            var service = [CChar](repeating: 0, count: Int(NI_MAXSERV))
            guard getnameinfo(sa, saLen, &host, socklen_t(host.count), &service, socklen_t(service.count), NI_NUMERICHOST | NI_NUMERICSERV) == 0 else {
                return ("?", "?")
            }
            return (String(cString: host), String(cString: service))
        }
    }
}

// MARK: IPv4 and IPv6

extension SockAddr {

    private var portOffset: Int? {
        switch Int32(self.saFamily) {
        case AF_INET:  return MemoryLayout.offset(of: \sockaddr_in.sin_port)!
        case AF_INET6: return MemoryLayout.offset(of: \sockaddr_in6.sin6_port)!
        default: return nil
        }
    }

    var port: UInt16? {
        guard let offset = self.portOffset else { return nil }
        return self.addr.withUnsafeBytes { buf in
            UInt16(bigEndian: buf.baseAddress!.load(fromByteOffset: offset, as: UInt16.self))
        }
    }
    
    mutating func setPort(_ port: UInt16) -> Bool {
        guard let offset = self.portOffset else { return false }
        self.addr.withUnsafeMutableBytes { buf in
            buf.baseAddress!.storeBytes(of: port.bigEndian, toByteOffset: offset, as: UInt16.self)
        }
        return true
    }
}

private func inaddrXFrom<SockAddrType, AddrType>(textRepresentation: String, family: Int32, sockAddrKeyPath: KeyPath<SockAddrType, AddrType>) -> AddrType?
    where SockAddrType: SockAddrCompatible
{
    guard
        let addr = SockAddr(textRepresentation: (textRepresentation, "0")),
        Int32(addr.saFamily) == family
    else {
        return nil
    }
    return addr.withSockAddrType { (sax: SockAddrType) -> AddrType in
        return sax[keyPath: sockAddrKeyPath]
    }
}

// MARK: IPv4

extension SockAddr {

    init(inAddr: in_addr, port: UInt16) {
        var sin = sockaddr_in()
        sin.sin_len = UInt8(MemoryLayout.size(ofValue: sin))
        sin.sin_family = sa_family_t(AF_INET)
        sin.sin_addr = inAddr
        sin.sin_port = port.bigEndian
        self.init(sin)
    }
}

extension in_addr {

    static let any = in_addr(s_addr: INADDR_ANY.bigEndian)
    static let broadcast = in_addr(s_addr: INADDR_BROADCAST.bigEndian)
    static let loopback = in_addr(s_addr: INADDR_LOOPBACK.bigEndian)
}

extension in_addr {

    init?(textRepresentation: String) {
        guard let result = inaddrXFrom(textRepresentation: textRepresentation, family: AF_INET, sockAddrKeyPath: \sockaddr_in.sin_addr) else {
            return nil
        }
        self = result
    }
    
    var textRepresentation: String {
        SockAddr(inAddr: self, port: 0).textRepresentation.address
    }
}

extension in_addr: ExpressibleByStringLiteral {

    public init(stringLiteral: String) {
        self.init(textRepresentation: stringLiteral)!
    }
}

// MARK: IPv6

extension SockAddr {

    init(in6Addr: in6_addr, port: UInt16) {
        var sin6 = sockaddr_in6()
        sin6.sin6_len = UInt8(MemoryLayout.size(ofValue: sin6))
        sin6.sin6_family = sa_family_t(AF_INET6)
        sin6.sin6_addr = in6Addr
        sin6.sin6_port = port.bigEndian
        self.init(sin6)
    }
}

extension in6_addr {

    static let any = in6addr_any
    // No `.broadcast` because IPv6 doesn’t support it.
    static let loopback = in6addr_loopback
}

extension in6_addr {

    init?(textRepresentation: String) {
        guard let result = inaddrXFrom(textRepresentation: textRepresentation, family: AF_INET6, sockAddrKeyPath: \sockaddr_in6.sin6_addr) else {
            return nil
        }
        self = result
    }
    
    var textRepresentation: String {
        SockAddr(in6Addr: self, port: 0).textRepresentation.address
    }
}

extension in6_addr: ExpressibleByStringLiteral {

    public init(stringLiteral: String) {
        self.init(textRepresentation: stringLiteral)!
    }
}

// MARK: Unix Domain

extension SockAddr {

    private static let sunHeaderLength = MemoryLayout.offset(of: \sockaddr_un.sun_path)!

    init?(path: FilePath) {
        let addrQ = path.withCString { pathPtr -> Data? in
            // Note that `pathLength` does not include the trailing nul because
            // the final `sockaddr` does not have it.
            let pathLength = strlen(pathPtr)
            guard
                let sunLength = UInt8(exactly: Self.sunHeaderLength + pathLength),
                sunLength <= SOCK_MAXADDRLEN
            else {
                return nil
            }
            var addr = Data(count: Int(sunLength))
            addr.withUnsafeMutableBytes { addrBuf in
                let sun = addrBuf.baseAddress!.bindMemory(to: sockaddr_un.self, capacity: 1)
                sun.pointee.sun_family = sa_family_t(AF_UNIX)
                sun.pointee.sun_len = sunLength
                memcpy(sun[\.sun_path.0]!, pathPtr, pathLength)
            }
            return addr
        }
        guard let addr = addrQ else { return nil }
        self.init(addr: addr)
    }
    
    var path: String? {
        guard self.saFamily == AF_UNIX else { return nil }
        return self.withSockAddrTypePtr { (sun: UnsafePointer<sockaddr_un>) -> String? in
            // We need to initialise `pathPtr` using the subscript because we
            // can’t afford for the compiler to make a copy of `sun_path.0` or,
            // worse yet, the whole `sun_path` tuple.  Remember that the path
            // can be longer than the 104 bytes allocated to it in `sockaddr_un`.
            //
            // We can’t use the `String(cString:)` initialiser because the path
            // is _not_ nul terminated.  Instead we use `String(bytes:encoding:)`.
            // Annoyingly, that takes a sequence of `UInt8` while `pathPtr` is a
            // point to `CChar`.  We use `withMemoryRebound(…)` to fix that.
            //
            // Unlikely `String(cString:)`, `String(bytes:encoding:)` might fail
            // if you give it bogus UTF-8.  Fortunately that doesn’t matter here
            // because a) it’s unlikely to happen in practice, and b) we are
            // allowed to return `nil`.
            //
            // Wow, more comments than code.  Achievement unlocked!
            let pathPtr = sun[\.sun_path.0]!
            let pathLength = Int(sun.pointee.sun_len) - Self.sunHeaderLength
            return pathPtr.withMemoryRebound(to: UInt8.self, capacity: pathLength) { pathPtr in
                let pathBuffer = UnsafeBufferPointer(start: pathPtr, count: pathLength)
                return String(bytes: pathBuffer, encoding: .utf8)
            }
        }
    }
}

// These were stolen, more or less, from [SR-11156][SR-11156]
//
// [SR-11156]: <https://bugs.swift.org/browse/SR-11156>

extension UnsafePointer {

    fileprivate subscript<Field>(field: KeyPath<Pointee, Field>) -> UnsafePointer<Field>? {
        guard let offset = MemoryLayout<Pointee>.offset(of: field) else { return nil }
        return (UnsafeRawPointer(self) + offset).assumingMemoryBound(to: Field.self)
    }
}

extension UnsafeMutablePointer {

    subscript<Field>(field: KeyPath<Pointee, Field>) -> UnsafePointer<Field>? {
        guard let offset = MemoryLayout<Pointee>.offset(of: field) else { return nil }
        return (UnsafeRawPointer(self) + offset).assumingMemoryBound(to: Field.self)
    }

    subscript<Field>(field: WritableKeyPath<Pointee, Field>) -> UnsafeMutablePointer<Field>? {
        guard let offset = MemoryLayout<Pointee>.offset(of: field) else { return nil }
        return (UnsafeMutableRawPointer(self) + offset).assumingMemoryBound(to: Field.self)
    }
}

@discardableResult
fileprivate func errnoQ<Result: SignedInteger>(retryOnInterrupt: Bool, _ body: () -> Result) throws -> Result {
    repeat {
        let result = body()
        if result >= 0 {
            return result
        }
        if retryOnInterrupt && errno == Errno.interrupted.rawValue {
            continue
        }
        throw Errno(rawValue: errno)
    } while true
}