//
//  ZYJQNURLProtocol.m
//  AppTest
//
//  Created by yangsen on 2020/4/7.
//  Copyright © 2020 com.qiniu. All rights reserved.
//

#import "ZYJQNURLProtocol.h"
#import "ZYJQNCFHttpClient.h"
#import "NSURLRequest+ZYJQNRequest.h"
#import "ZYJQNDnsPrefetcher.h"
#import "NSObject+ZYJQNSwizzle.h"
#import <objc/runtime.h>

@interface ZYJQNRequestInfo : NSObject
@property(nonatomic, weak)NSURLSession *session;
@property(nonatomic, weak)NSURLSessionDataTask *task;
@end
@implementation ZYJQNRequestInfo
@end

@interface ZYJQNRequestInfoManager : NSObject
@property(nonatomic, strong)NSMutableDictionary <NSString *, ZYJQNRequestInfo *> *infos;
@end
@implementation ZYJQNRequestInfoManager
+ (instancetype)share{
    static ZYJQNRequestInfoManager *manager = nil;
    static dispatch_once_t onceToken;
    dispatch_once(&onceToken, ^{
        manager = [[ZYJQNRequestInfoManager alloc] init];
        [manager setupData];
    });
    return manager;
}
- (void)setupData{
    _infos = [NSMutableDictionary dictionary];
}
- (void)setRequestInfo:(ZYJQNRequestInfo *)info forRequest:(NSURLRequest *)request{
    NSString *requestIdentifier = [request ZYJQN_identifier];
    if (!requestIdentifier || !info) {
        return;
    }
    @synchronized (self) {
        [self.infos setObject:info forKey:requestIdentifier];
    }
}
- (void)removeRequestInfoForRequest:(NSURLRequest *)request{
    NSString *requestIdentifier = [request ZYJQN_identifier];
    if (!requestIdentifier) {
        return;
    }
    @synchronized (self) {
        [self.infos removeObjectForKey:requestIdentifier];
    }
}
- (ZYJQNRequestInfo *)getRequestInfoForRequest:(NSURLRequest *)request{
    NSString *requestIdentifier = [request ZYJQN_identifier];
    if (!requestIdentifier) {
        return nil;
    }
    
    ZYJQNRequestInfo *info = nil;
    @synchronized (self) {
        info = self.infos[requestIdentifier];
    }
    return info;
}

@end



@interface NSURLRequest(ZYJQNHttps)
@property(nonatomic, readonly)NSURLSession *ZYJQN_session;
@property(nonatomic, readonly)NSURLSessionDataTask *ZYJQN_task;
@end
@implementation NSURLRequest(ZYJQNHttps)
- (NSURLSession *)ZYJQN_session{
    return [[ZYJQNRequestInfoManager share] getRequestInfoForRequest:self].session;
}
- (NSURLSessionDataTask *)ZYJQN_task{
    return [[ZYJQNRequestInfoManager share] getRequestInfoForRequest:self].task;
}
- (void)ZYJQN_setSession:(NSURLSession *)session task:(NSURLSessionDataTask *)task{
    ZYJQNRequestInfo *info = [[ZYJQNRequestInfo alloc] init];
    info.session = session;
    info.task = task;
    [[ZYJQNRequestInfoManager share] setRequestInfo:info forRequest:self];
}
- (void)ZYJQN_requestRemoveTask{
    [[ZYJQNRequestInfoManager share] removeRequestInfoForRequest:self];
}
- (BOOL)ZYJQNHttps_shouldInit{
    if ([self ZYJQN_isQiNiuRequest]
        && ([self.URL.absoluteString hasPrefix:@"http://"] || [self.URL.absoluteString hasPrefix:@"https://"])) {
        return YES;
    } else {
        return NO;
    }
}
@end


@interface NSURLSession(ZYJQNURLProtocol)
@end
@implementation NSURLSession(ZYJQNURLProtocol)

+ (void)load{
    static dispatch_once_t onceToken;
    dispatch_once(&onceToken, ^{
        [self ZYJQN_swizzleInstanceMethodsOfSelectorA:@selector(dataTaskWithRequest:)
                                         selectorB:@selector(ZYJQN_dataTaskWithRequest:)];
    });
}
- (NSURLSessionDataTask *)ZYJQN_dataTaskWithRequest:(NSURLRequest *)request{
    NSURLSessionDataTask *task = [self ZYJQN_dataTaskWithRequest:request];
    if ([request ZYJQN_isQiNiuRequest]) {
        [request ZYJQN_setSession:self task:task];
    }
    return task;
}
@end



@interface ZYJQNURLProtocol()<ZYJQNCFHttpClientDelegate>

@property(nonatomic, strong)ZYJQNCFHttpClient *httpsClient;

@end
@implementation ZYJQNURLProtocol

#define kZYJQNRequestIdentifiers @"ZYJQNRequestIdentifiers"

+ (void)registerProtocol{
    [NSURLProtocol registerClass:[ZYJQNURLProtocol class]];
}

+ (void)unregisterProtocol{
    [NSURLProtocol unregisterClass:[ZYJQNURLProtocol class]];
}

//MARK: -- overload
+ (BOOL)canInitWithRequest:(NSURLRequest *)request {

    if ([NSURLProtocol propertyForKey:kZYJQNRequestIdentifiers inRequest:request]) {
        return NO;
    }

    id <ZYJQNInetAddressDelegate> address = [kZYJQNDnsPrefetcher getInetAddressByHost:request.URL.host].firstObject;
    
    if ([request ZYJQNHttps_shouldInit]
        && address.ipValue && address.ipValue.length > 0) {
        return YES;
    } else {
        [request ZYJQN_requestRemoveTask];
        return NO;
    }
}


+ (NSURLRequest *)canonicalRequestForRequest:(NSURLRequest *)request {
    
    NSString *host = request.URL.host;
    NSString *ip = [kZYJQNDnsPrefetcher getInetAddressByHost:host].firstObject.ipValue;
    
    if (!ip || ip.length == 0) {
        return request;
    }
    
    NSString *urlString = request.URL.absoluteString;
    urlString = [urlString stringByReplacingOccurrencesOfString:host withString:ip];
    
    NSMutableURLRequest *requestNew = [request mutableCopy];
    requestNew.URL = [NSURL URLWithString:urlString];
    
    return [requestNew copy];
}


- (void)startLoading {
    
    [self loadingRequest:self.request];
}

- (void)stopLoading {

    [self.httpsClient stopLoading];
    self.httpsClient = nil;
}


- (void)loadingRequest:(NSURLRequest *)request{
    
    self.httpsClient = [ZYJQNCFHttpClient client:request];
    self.httpsClient.delegate = self;
    
    [NSURLProtocol setProperty:@(YES)
                        forKey:kZYJQNRequestIdentifiers
                     inRequest:self.httpsClient.request];
    
    [self.httpsClient startLoading];
}

//MARK: -- delegate
- (void)didSendBodyData:(int64_t)bytesSent
         totalBytesSent:(int64_t)totalBytesSent
totalBytesExpectedToSend:(int64_t)totalBytesExpectedToSend{
    
    id <NSURLSessionTaskDelegate> sessionDelegate = (id <NSURLSessionTaskDelegate>)self.request.ZYJQN_session.delegate;
    if ([sessionDelegate respondsToSelector:@selector(URLSession:
                                                      task:
                                                      didSendBodyData:
                                                      totalBytesSent:
                                                      totalBytesExpectedToSend:)]) {
        
        [sessionDelegate URLSession:self.request.ZYJQN_session
                               task:self.request.ZYJQN_task
                    didSendBodyData:bytesSent
                     totalBytesSent:totalBytesSent
           totalBytesExpectedToSend:totalBytesExpectedToSend];
    }
}

- (void)didFinish {
    
    [self.client URLProtocolDidFinishLoading:self];
    [self.request ZYJQN_requestRemoveTask];
}

- (void)didLoadData:(nonnull NSData *)data {
    
    [self.client URLProtocol:self didLoadData:data];
}

- (BOOL)evaluateServerTrust:(SecTrustRef)serverTrust forDomain:(NSString *)domain {

    
    NSMutableArray *policies = [NSMutableArray array];
    if (domain) {
        [policies addObject:(__bridge_transfer id)SecPolicyCreateSSL(true, (__bridge CFStringRef)domain)];
    } else {
        [policies addObject:(__bridge_transfer id)SecPolicyCreateBasicX509()];
    }

    SecTrustSetPolicies(serverTrust, (__bridge CFArrayRef)policies);

    SecTrustResultType result = kSecTrustResultInvalid;
    
    OSStatus status = SecTrustEvaluate(serverTrust, &result);
    if (status != errSecSuccess) {
        return NO;
    }
    
    if (result == kSecTrustResultUnspecified || result == kSecTrustResultProceed) {
        return YES;
    } else {
        return NO;
    }
}

- (void)onError:(nonnull NSError *)error {
    
    [self.client URLProtocol:self didFailWithError:error];
}

- (void)onReceiveResponse:(nonnull NSURLResponse *)response {
    
    [self.client URLProtocol:self
          didReceiveResponse:response
          cacheStoragePolicy:NSURLCacheStorageNotAllowed];
}

- (void)redirectedToRequest:(nonnull NSURLRequest *)request redirectResponse:(nonnull NSURLResponse *)redirectResponse {
    
    if ([self.client respondsToSelector:@selector(URLProtocol:wasRedirectedToRequest:redirectResponse:)]) {
        [self.client URLProtocol:self wasRedirectedToRequest:request redirectResponse:redirectResponse];
        [self.client URLProtocolDidFinishLoading:self];
    } else {
        [self.httpsClient stopLoading];
        [self loadingRequest:request];
    }
}

@end


@implementation NSURLSessionConfiguration(ZYJQNURLProtocol)
+ (NSURLSessionConfiguration *)ZYJQN_sessionConfiguration{

    NSURLSessionConfiguration *config = [NSURLSessionConfiguration defaultSessionConfiguration];
    config.protocolClasses = @[[ZYJQNURLProtocol class]];
    return config;
}
@end


