Gemini/GeminiProtocol/GeminiDataTask.swift

//
//  GeminiDataTask.swift
//  GeminiProtocol
//
//  Created by Shadowfacts on 7/14/20.
//

import Foundation
import Network

public class GeminiDataTask {
    public typealias Completion = (Result<GeminiResponse, Error>) -> Void
    
    private static let queue = DispatchQueue(label: "GeminiDataTask", qos: .default)
    
    public let request: GeminiRequest
    private let completion: Completion
    private var state = State.unstarted
    private let connection: NWConnection
    // todo: remove stupid hack when deployment target is >= iOS 15/macOS 12
    private var _attribution: Any? = nil
    @available(iOS 15.0, *)
    public var attribution: NWParameters.Attribution {
        get {
            _attribution as? NWParameters.Attribution ?? .developer
        }
        set {
            _attribution = newValue
        }
    }
    
    public init(request: GeminiRequest, completion: @escaping Completion) {
        self.request = request
        self.completion = completion
        let port = request.url.port != nil ? UInt16(request.url.port!) : 1965
        let endpoint: NWEndpoint = .hostPort(host: NWEndpoint.Host(request.url.host!), port: NWEndpoint.Port(rawValue: port)!)
        self.connection = NWConnection(to: endpoint, using: .gemini)
        
        if #available(iOS 15.0, *) {
            connection.parameters.attribution = attribution
        }
        
        self.connection.stateUpdateHandler = { (newState) in
            switch newState {
            case .ready:
                self.sendRequest()
            case let .failed(error):
                self.state = .completed
                self.connection.cancel()
                self.completion(.failure(.connectionError(error)))
            default:
                break
            }
        }
        
    }
    
    public convenience init(url: URL, completion: @escaping Completion) throws {
        self.init(request: try GeminiRequest(url: url), completion: completion)
    }
    
    deinit {
        self.cancel()
    }
    
    public func resume() {
        guard state == .unstarted else { return }
        state = .started
        connection.start(queue: GeminiDataTask.queue)
    }
    
    public func cancel() {
        guard state != .completed else { return }
        connection.cancel()
        state = .completed
    }
    
    private func sendRequest() {
        let message = NWProtocolFramer.Message(geminiRequest: request)
        let context = NWConnection.ContentContext(identifier: "GeminiRequest", metadata: [message])
        connection.send(content: nil, contentContext: context, isComplete: true, completion: .contentProcessed({ (error) in
            if let error = error {
                self.state = .completed
                self.connection.cancel()
                self.completion(.failure(.connectionError(error)))
            }
        }))
        receive()
    }
    
    private func receive() {
        connection.receiveMessage { (data, context, isComplete, error) in
            if let error = error {
                self.completion(.failure(.connectionError(error)))
            } else if let message = context?.protocolMetadata(definition: GeminiProtocol.definition) as? NWProtocolFramer.Message,
                      let header = message.geminiResponseHeader {
                guard isComplete else { fatalError() }
                let response = GeminiResponse(header: header, body: data)
                self.completion(.success(response))
            }
            
            self.connection.cancel()
            self.state = .completed
        }
    }
}

extension GeminiDataTask {
    enum State {
        case unstarted, started, completed
    }
}

public extension GeminiDataTask {
    enum Error: Swift.Error {
        case noData
        case connectionError(NWError)
    }
}