 
 import UIKit
 
-class SuperMockURLProtocol: NSURLProtocol {
+class SuperMockURLProtocol: URLProtocol {
     
-    override class func canInitWithRequest(request: NSURLRequest) -> Bool {
+    override class func canInit(with request: URLRequest) -> Bool {
         
         if request.hasMock() {
-            print("Requesting MOCK for : \(request.URL)")
+            print("Requesting MOCK for : \(String(describing: request.url))")
             return true
         }
-        print("Passing Through WITHOUT MOCK : \(request.URL)")
+        print("Passing Through WITHOUT MOCK : \(String(describing: request.url))")
         return false
     }
     
     
-    override class func canonicalRequestForRequest(request: NSURLRequest) -> NSURLRequest {
+    override class func canonicalRequest(for request: URLRequest) -> URLRequest {
         return request
     }
 

    }
}

class SuperMockRecordingURLProtocol: NSURLProtocol {
    
    var connection : NSURLConnection?
    var mutableData : NSMutableData?
    
    override class func canInitWithRequest(request: NSURLRequest) -> Bool {
        
        if let _ = NSURLProtocol.propertyForKey("SuperMockRecordingURLProtocol", inRequest: request) {
            return false
        }
        if SuperMockResponseHelper.sharedHelper.recording  {

        return false
    }
    
    override class func canonicalRequest(for request: URLRequest) -> URLRequest {
        return request
    }
    
    override class func requestIsCacheEquivalent(_ a: URLRequest, to b: URLRequest) -> Bool {
        return super.requestIsCacheEquivalent(a, to:b)
    }
    
    override func startLoading() {
        
        if let copyRequest = (request as NSURLRequest).mutableCopy() as? NSMutableURLRequest {
            
            URLProtocol.setProperty(request.url!, forKey: "SuperMockRecordingURLProtocol", in: copyRequest)
            let configuration = URLSessionConfiguration.background(withIdentifier: UUID().uuidString)
            let session = URLSession(configuration: configuration, delegate: self, delegateQueue: OperationQueue())
            
            dataTask = session.dataTask(with: copyRequest as URLRequest)
            dataTask?.resume()
            self.dataTask!.resume()
        }
    }
    
    override func stopLoading() {
        dataTask?.cancel()
        response = nil
        mutableData = nil
        
        connection?.cancel()
    }
}
extension SuperMockRecordingURLProtocol: URLSessionDataDelegate {
    
    func urlSession(_ session: URLSession, dataTask: URLSessionDataTask, didReceive data: Data) {
        client?.urlProtocol(self, didLoad: data)
        mutableData?.append(data)
    }
}
extension SuperMockRecordingURLProtocol: URLSessionTaskDelegate {
    
    func urlSession(_ session: URLSession, dataTask: URLSessionDataTask, didReceive response: URLResponse, completionHandler: @escaping (URLSession.ResponseDisposition) -> Void) {
        mutableData = NSMutableData()
        client?.urlProtocol(self, didReceive: response, cacheStoragePolicy: URLCache.StoragePolicy.notAllowed)
        completionHandler(.allow)
    }
    
    func urlSession(_ session: URLSession, task: URLSessionTask, didCompleteWithError error: Error?) {
        if let error = error {
            client?.urlProtocol(self, didFailWithError: error)
            return
        }
        SuperMockResponseHelper.sharedHelper.recordDataForRequest(mutableData! as Data, request: request)
        client?.urlProtocolDidFinishLoading(self)
        
    }
}

extension SuperMockRecordingURLProtocol: NSURLConnectionDataDelegate {
    
    func connection(_ connection: NSURLConnection, didReceive response: URLResponse) {
        if let httpResponse = response as? HTTPURLResponse {
            let headers = httpResponse.allHeaderFields
            SuperMockResponseHelper.sharedHelper.recordResponseHeadersForRequest(headers, request: request, response: httpResponse)
        }
        client?.urlProtocol(self, didReceive: response, cacheStoragePolicy: URLCache.StoragePolicy.notAllowed)
    }
    
    func connection(_ connection: NSURLConnection, didReceive data: Data) {
        client?.urlProtocol(self, didLoad: data)
        mutableData?.append(data)
    }
    
    func connectionDidFinishLoading(_ connection: NSURLConnection) {
        client?.urlProtocolDidFinishLoading(self)
        // TODO: fix the !
        SuperMockResponseHelper.sharedHelper.recordDataForRequest(mutableData! as Data, request: request)
    }
    
    func connection(_ connection: NSURLConnection, didFailWithError error: Error) {
        client?.urlProtocol(self, didFailWithError: error)
    }
}
