import Foundation
import Network

enum UDPState: String {
    case initial
    case connectionCreated
    case preparingConnection
    case sending
    case waitingForResponse
}

final class UDPClient {
    
    let host: NWEndpoint.Host
    
    let port: NWEndpoint.Port
    
    private var connection: NWConnection?
    
    private var completion: ((Data?) -> Void)?
    
    private var state: UDPState = .initial
    
    init(host: String, port: UInt16) {
        self.host = .init("192.168.188.118")
        self.port = .init(rawValue: port)!
    }
    
    deinit {
        print("Destroying UDP Client")
        finish()
    }
    
    func begin() {
        guard state == .initial else {
            print("Invalid state for begin(): \(state)")
            return
        }
        connection = NWConnection(host: host, port: port, using: .udp)
        state = .connectionCreated
        print("Created connection: \(connection != nil)")
    }
    
    func send(message: Data, completion: @escaping (Data?) -> Void) {
        print("Sending message to \(host) at port \(port)")
        guard state == .connectionCreated else {
            print("Invalid state preparing for send: \(state)")
            return
        }
        guard let connection else {
            print("Failed to send, no connection")
            completion(nil)
            return
        }
        
        self.completion = completion
        connection.stateUpdateHandler = { [weak self] (newState) in
            switch (newState) {
            case .ready:
                print("State: Ready\n")
                self?.send(message, over: connection)
            case .setup:
                print("State: Setup\n")
            case .cancelled:
                print("Cancelled UDP connection")
                self?.finish()
            case .preparing:
                print("Preparing UDP connection")
            case .failed(let error):
                print("Failed to start UDP connection: \(error)")
                self?.finish()
//            default:
//                print("ERROR! State not defined!\n")
//                self?.finish()
            case .waiting(_):
                print("Waiting for UDP connection path change")
            @unknown default:
                print("Unknown UDP connection state: \(newState)")
                self?.finish()
            }
        }
        print("Preparing connection")
        state = .preparingConnection
        connection.start(queue: .global())
        
        DispatchQueue.main.asyncAfter(deadline: .now() + .seconds(5)) {
            guard self.state == .preparingConnection else {
                return
            }
            print("Timed out preparing connection")
            self.finish()
        }
    }
    
    private func finish(_ data: Data? = nil) {
        completion?(data)
        completion = nil
        connection?.stateUpdateHandler = nil
        connection?.cancel()
        connection = nil
        state = .initial
    }
    
    private func send(_ data: Data, over connection: NWConnection) {
        guard state == .preparingConnection else {
            print("Invalid state for send: \(state)")
            return
        }
        
        connection.stateUpdateHandler = nil
        
        let completion = NWConnection.SendCompletion.contentProcessed { [weak self] error in
            if let error {
                print("Failed to send UDP packet: \(error)")
                self?.finish()
            } else {
                print("Finished sending message")
                self?.waitForResponse(over: connection)
            }
        }
        state = .sending
        connection.send(content: data, completion: completion)
        print("Started to send message")
        
        DispatchQueue.main.asyncAfter(deadline: .now() + .seconds(5)) {
            guard self.state == .sending else {
                return
            }
            print("Timed out waiting for for send to complete")
            self.finish()
        }
    }
    
    private func waitForResponse(over connection: NWConnection) {
        guard state == .sending else {
            print("Invalid state for send: \(state)")
            return
        }
        state = .waitingForResponse
        connection.receiveMessage { [weak self] (data, context, isComplete, error) in
            guard self?.state == .waitingForResponse else {
                return
            }
            guard isComplete else {
                if let error {
                    print("Failed to receive UDP message: \(error)")
                } else {
                    print("Failed to receive complete UDP message without error")
                }
                self?.finish()
                return
            }
            guard let data else {
                print("Received UDP message without data")
                self?.finish()
                return
            }
            print("Received \(data.count) bytes")
            self?.finish(data)
        }
        
        DispatchQueue.main.asyncAfter(deadline: .now() + .seconds(5)) {
            guard self.state == .waitingForResponse else {
                return
            }
            print("Timed out waiting for response")
            self.finish()
        }
    }
}