//
//  SBRProcessPlugin.m
//  SBrowserProcessBundle
//
//  Created by James Magahern on 7/22/20.
//

#import "SBRProcessPlugin.h"

#import "SBRWebProcessDelegate.h"
#import "SBRWebProcessProxy.h"

#import <WebKit/_WKRemoteObjectInterface.h>
#import <WebKit/_WKRemoteObjectRegistry.h>
#import <WebKit/WKWebProcessPlugInBrowserContextController.h>
#import <WebKit/WKWebProcessPlugInBrowserContextControllerPrivate.h>
#import <WebKit/WKWebProcessPlugInLoadDelegate.h>

@interface SBRProcessPlugin () <WKWebProcessPlugInLoadDelegate, SBRWebProcessProxy>
@property (nonatomic, strong) id<SBRWebProcessDelegate> processDelegate;
@property (nonatomic, strong) NSMutableSet<NSString *> *allowedResourceOrigins;
@end

@implementation SBRProcessPlugin

+ (void)initialize
{
    [super initialize];
    NSLog(@"SBRProcessPlugin Initialized");
}

#pragma mark <SBRWebProcessProxy>

- (void)hello
{
    NSLog(@"SBRProcessPlugin: Helloooooo");
}

- (void)syncAllowedResourceOrigins:(NSSet<NSString *> *)allowedOrigins
{
    if (!_allowedResourceOrigins) {
        _allowedResourceOrigins = [allowedOrigins mutableCopy];
    } else {
        [_allowedResourceOrigins unionSet:allowedOrigins];
    }
}

#pragma mark <WKWebProcessPlugIn>

- (void)webProcessPlugIn:(WKWebProcessPlugInController *)plugInController didCreateBrowserContextController:(WKWebProcessPlugInBrowserContextController *)browserContextController
{
    _WKRemoteObjectInterface *proxyInterface = [_WKRemoteObjectInterface remoteObjectInterfaceWithProtocol:@protocol(SBRWebProcessProxy)];
    [browserContextController._remoteObjectRegistry registerExportedObject:self interface:proxyInterface];

    _WKRemoteObjectInterface *eventListenerInterface = [_WKRemoteObjectInterface remoteObjectInterfaceWithProtocol:@protocol(SBRWebProcessDelegate)];
    self.processDelegate = [browserContextController._remoteObjectRegistry remoteObjectProxyWithInterface:eventListenerInterface];

    browserContextController.loadDelegate = self;
    
    [[self processDelegate] webProcessDidConnect];
}

#pragma mark <WKWebProcessPlugInLoadDelegate>

- (NSURLRequest *)webProcessPlugInBrowserContextController:(WKWebProcessPlugInBrowserContextController *)controller frame:(WKWebProcessPlugInFrame *)frame willSendRequest:(NSURLRequest *)request redirectResponse:(NSURLResponse *)redirectResponse
{
    NSLog(@"SBRProcessPlugin: Sending request: %@", request);
    
    NSURL *requestURL = [request URL];
    NSString *originString = [requestURL host];
    NSString *requestExtension = [requestURL pathExtension];
    if (requestExtension.length > 0 && [requestExtension isEqualToString:@"js"]) {
        [[self processDelegate] webProcessDidLoadScriptWithOrigin:originString];
        
        if ([_allowedResourceOrigins containsObject:originString]) {
            NSLog(@"SBRProcessPlugin: Allowing whitelisted requestURL: %@", requestURL);
        } else {
            NSLog(@"SBRProcessPlugin: Blocking requestURL: %@", requestURL);
            request = nil;
        }
    }
    
    return request;
}

@end