//
//  SBRProcessBundleBridge.m
//  SBrowser
//
//  Created by James Magahern on 7/22/20.
//

#import "SBRProcessBundleBridge.h"

#import "SBRScriptPolicy.h"
#import "Hacks.h"

#import <OSLog/OSLog.h>
#import <stdatomic.h>

#import <WebKit/_WKRemoteObjectInterface.h>
#import <WebKit/_WKRemoteObjectRegistry.h>
#import <WebKit/_WKProcessPoolConfiguration.h>

#import <WebKit/WKProcessPoolPrivate.h>
#import <WebKit/WKWebViewPrivate.h>
#import <WebKit/WKWebViewConfigurationPrivate.h>
#import <WebKit/WKContentRuleListStore.h>

#define WKUserStyleSheet id
#define WKUserStyleSheetEncodedClassName "X1dLVXNlclN0eWxlU2hlZXQ="

NSArray<NSString *> *CommonCDNList(void) {
    static dispatch_once_t onceToken;
    static NSArray<NSString *> *commonCDNList = nil;
    dispatch_once(&onceToken, ^{
        commonCDNList = @[
            @"cdn.jsdelivr.net",
             @"fsdn.net",
             @"cdnjs.com",
             @"osscdn.com",
             @"code.jquery.com",
             @"bootstrapcdn.com"
        ];
    });
    
    return commonCDNList;
}

@interface StyleSheet : NSObject <NSCopying>
@property (nonatomic, readonly, copy) NSString *source;
@property (nonatomic, readonly, copy) NSURL *baseURL;
@property (nonatomic, readonly, getter=isForMainFrameOnly) BOOL forMainFrameOnly;

- (instancetype)initWithSource:(NSString *)source forMainFrameOnly:(BOOL)forMainFrameOnly;
@end

// Note: We no longer proxy http/https via WKURLSchemeHandler; we use
// WebKit content blocking rules instead.

@interface WKUserContentController (Private)
- (void)__addUserStyleSheet:(WKUserStyleSheet)userStyleSheet;
- (void)__removeUserStyleSheet:(WKUserStyleSheet)userStyleSheet;
- (void)__addUserScriptImmediately:(WKUserScript *)userScript;
@end

@implementation WKUserContentController (Private)

- (void)__addUserStyleSheet:(id)userStyleSheet
{
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Warc-performSelector-leaks"
    [self performSelector:DecodedSelector("X2FkZFVzZXJTdHlsZVNoZWV0Og==") withObject:userStyleSheet];
#pragma clang diagnostic pop
}

- (void)__removeUserStyleSheet:(id)userStyleSheet
{
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Warc-performSelector-leaks"
    [self performSelector:DecodedSelector("X3JlbW92ZVVzZXJTdHlsZVNoZWV0Og==") withObject:userStyleSheet];
#pragma clang diagnostic pop
}

- (void)__addUserScriptImmediately:(WKUserScript *)userScript
{
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Warc-performSelector-leaks"
    [self performSelector:DecodedSelector("X2FkZFVzZXJTY3JpcHRJbW1lZGlhdGVseTo=") withObject:userScript];
#pragma clang diagnostic pop
}

@end

#define LOG_DEBUG(format, ...) os_log_debug(_log, format, ##__VA_ARGS__)
#define LOG_ERROR(format, ...) os_log_error(_log, format, ##__VA_ARGS__)

static os_log_t _log;

__attribute__((constructor))
static void initialize_log(void) {
    _log = os_log_create("net.buzzert.attractor.webview", "bridge");
}

@interface NSURLResponse (BridgeAdditions)
@property (nonatomic, readonly) BOOL isJavascriptResponse;
@end

@implementation NSURLResponse (BridgeAdditions)

- (BOOL)isJavascriptResponse
{
    NSString *extension = [[self URL] pathExtension];
    if ([[extension lowercaseString] isEqualToString:@"js"]) {
        return YES;
    }
    
    NSString *MIMEType = [self MIMEType];
    if ([[MIMEType lowercaseString] containsString:@"javascript"]) {
        return YES;
    }
    
    return NO;
}

@end

@interface SBRProcessBundleBridge ()

@end

@implementation SBRProcessBundleBridge {
    WKWebView              *_webView;
    WKWebViewConfiguration *_webViewConfiguration;
    WKProcessPool          *_processPool;
    
    WKUserStyleSheet       _darkModeStyleSheet;
    WKUserScript          *_readabilityScript;
    
    NSArray<WKUserScript *> *_userScripts;
    
    // These come from settings. 
    WKUserStyleSheet   _customizedUserStylesheet;
    WKUserScript      *_customizedUserScript;

    // Content blocking
    WKContentRuleList *_activeScriptRuleList;
    void *_urlKVOContext;
}

- (void)tearDown
{
    if (_webView) {
        @try { [_webView removeObserver:self forKeyPath:@"URL" context:_urlKVOContext]; }
        @catch (__unused NSException *ex) {}
    }
}

- (instancetype)initWithWebViewConfiguration:(WKWebViewConfiguration *)webViewConfiguration
{
    self = [super init];
    if (self) {
        if (!webViewConfiguration) {
            webViewConfiguration = [[WKWebViewConfiguration alloc] init];

            // Set up process pool
            _processPool = [[WKProcessPool alloc] init];
            webViewConfiguration.processPool = _processPool;
            
            webViewConfiguration._waitsForPaintAfterViewDidMoveToWindow = NO;
            webViewConfiguration._applePayEnabled = YES;
            // No http/https interception — rely on content blocking rules instead.
        }
        
        _webViewConfiguration = webViewConfiguration;
        
        // User scripts
        WKUserContentController *userContentController = [_webViewConfiguration userContentController];
        for (WKUserScript *script in [self _userScripts]) {
            [userContentController addUserScript:script];
        }
        
        // Reload customized user scripts/stylesheets from settings
        [self reloadCustomizedUserScriptsAndStylesheets];
        
        // Instantiate web view
        WKWebView *webView = [[WKWebView alloc] initWithFrame:CGRectZero configuration:webViewConfiguration];
        if (@available(iOS 16.0, *)) {
            webView.findInteractionEnabled = YES;
        }
        
        _webView = webView;
        _urlKVOContext = &_urlKVOContext; // unique context pointer
        [_webView addObserver:self forKeyPath:@"URL" options:(NSKeyValueObservingOptionNew) context:_urlKVOContext];

        // Initialize content blocking rules for current host
        [self rebuildContentBlockingRulesForCurrentHost];
    }
    
    return self;
}

- (WKUserScript *)_loadScriptForResource:(NSString *)resourceName withExtension:(NSString *)extension
{
    NSURL *url = [[NSBundle mainBundle] URLForResource:resourceName withExtension:extension];
    NSString *source = [NSString stringWithContentsOfURL:url encoding:NSUTF8StringEncoding error:nil];
    
    return [[WKUserScript alloc] initWithSource:source injectionTime:WKUserScriptInjectionTimeAtDocumentEnd forMainFrameOnly:YES];
}

- (NSArray<WKUserScript *> *)_userScripts
{
    if (!_userScripts) {
        _userScripts = @[
            [self _loadScriptForResource:@"Tagger" withExtension:@"js"],
        ];
    }
    
    return _userScripts;
}

- (void)reloadCustomizedUserScriptsAndStylesheets
{
    WKUserContentController *userContentController = [_webViewConfiguration userContentController];
    [userContentController removeAllUserScripts];
        
    NSString *scriptSource = [[NSUserDefaults standardUserDefaults] stringForKey:@"userScript"];
    if ([scriptSource length]) {
        _customizedUserScript = [[WKUserScript alloc] initWithSource:scriptSource injectionTime:WKUserScriptInjectionTimeAtDocumentEnd forMainFrameOnly:YES];
        [userContentController addUserScript:_customizedUserScript];
    }
    
    if (_customizedUserStylesheet) {
        [userContentController __removeUserStyleSheet:_customizedUserStylesheet];
    }
    
    NSString *stylesheetSource = [[NSUserDefaults standardUserDefaults] stringForKey:@"userStylesheet"];
    if ([stylesheetSource length]) {
        _customizedUserStylesheet = [[DecodedClass(WKUserStyleSheetEncodedClassName) alloc] initWithSource:stylesheetSource forMainFrameOnly:YES];
        [userContentController __addUserStyleSheet:_customizedUserStylesheet];
    }
}

#pragma mark Former <SBRWebProcessDelegate> methods

- (void)webProcessDidAllowScriptWithOrigin:(NSString *)origin
{
    dispatch_async(dispatch_get_main_queue(), ^{
        [[self delegate] webProcess:self didAllowScriptResourceFromOrigin:origin];
    });
}

- (void)webProcessDidBlockScriptWithOrigin:(NSString *)origin
{
    dispatch_async(dispatch_get_main_queue(), ^{
        [[self delegate] webProcess:self didBlockScriptResourceFromOrigin:origin];
    });
}

#pragma mark Actions

- (void)policyDataSourceDidChange
{
    // Rebuild content blocking rules when policy changes.
    [self rebuildContentBlockingRulesForCurrentHost];
}

- (void)setAllowAllScripts:(BOOL)allowAllScripts
{
    _allowAllScripts = allowAllScripts;
    [self rebuildContentBlockingRulesForCurrentHost];
}

- (void)setDarkModeEnabled:(BOOL)darkModeEnabled
{
    _darkModeEnabled = darkModeEnabled;
    
    WKUserContentController *userContentController = [_webViewConfiguration userContentController];
    
    if (darkModeEnabled) {
        if (!_darkModeStyleSheet) {
            NSURL *styleSheetURL = [[NSBundle mainBundle] URLForResource:@"darkmode" withExtension:@"css"];
            NSString *styleSheetSource = [NSString stringWithContentsOfURL:styleSheetURL encoding:NSUTF8StringEncoding error:nil];
            _darkModeStyleSheet = [[DecodedClass(WKUserStyleSheetEncodedClassName) alloc] initWithSource:styleSheetSource forMainFrameOnly:NO];
        }
        
        [userContentController __addUserStyleSheet:_darkModeStyleSheet];
    } else if (_darkModeStyleSheet) {
        [userContentController __removeUserStyleSheet:_darkModeStyleSheet];
    }
}

- (void)parseDocumentForReaderMode:(void (^)(NSString * _Nonnull))completionBlock
{
    WKUserContentController *userContentController = [_webViewConfiguration userContentController];
    
    if (!_readabilityScript) {
        _readabilityScript = [self _loadScriptForResource:@"Readability" withExtension:@"js"];
    }
    
    [userContentController __addUserScriptImmediately:_readabilityScript];
    
    NSString *script = @""
    "var documentClone = document.cloneNode(true);"
    "var article = new Readability(documentClone).parse();"
    "article.content";
    
    os_log_t log = _log;
    [_webView evaluateJavaScript:script completionHandler:^(NSString *result, NSError * _Nullable error) {
        if (error != nil) {
            os_log_error(log, "Bridge: Readability error: %@", error.localizedDescription);
        } else {
            completionBlock(result);
        }
    }];
}

// MARK: - Content Blocking Rules

- (void)observeValueForKeyPath:(NSString *)keyPath ofObject:(id)object change:(NSDictionary<NSKeyValueChangeKey,id> *)change context:(void *)context
{
    if (context == _urlKVOContext && [keyPath isEqualToString:@"URL"]) {
        [self rebuildContentBlockingRulesForCurrentHost];
        return;
    }
    [super observeValueForKeyPath:keyPath ofObject:object change:change context:context];
}

- (void)rebuildContentBlockingRulesForCurrentHost
{
    NSString *hostOrigin = _webView.URL.host;
    if (!hostOrigin) {
        // No page loaded yet; clear any existing rules
        [self applyContentRuleListJSON:nil withName:nil];
        return;
    }

    // Determine policy for this host
    NSNumber *policyType = nil;
    @synchronized (_policyDataSource) {
        NSDictionary<NSString *, NSNumber *> *policyTypes = [_policyDataSource scriptPolicyTypeByOrigin];
        policyType = [policyTypes objectForKey:hostOrigin] ?: @(0);
    }

    // Alpha=0 Bravo=1 Charlie=2 Delta=3 Echo=4
    SBRScriptOriginPolicyType type = [policyType integerValue];
    if (_allowAllScripts || type >= SBRScriptOriginPolicyTypeEcho) {
        // Echo or shields down: no blocking
        [self applyContentRuleListJSON:nil withName:nil];
        return;
    }

    // Base trigger applies only to this page's domain
    NSMutableArray *rules = [NSMutableArray array];
    NSDictionary *baseScriptTrigger = @{
        @"resource-type" : @[ @"script" ],
        @"if-domain"     : @[ hostOrigin ]
    };

    if (type <= SBRScriptOriginPolicyTypeBravo) {
        // Alpha or Bravo: block all external script subresources (inline JS is controlled via preferences elsewhere)
        [rules addObject:@{ @"trigger": baseScriptTrigger, @"action": @{ @"type": @"block" } }];
    } else {
        // Charlie/Delta: block third-party scripts
        NSMutableDictionary *trigger = [baseScriptTrigger mutableCopy];
        trigger[@"load-type"] = @[ @"third-party" ];
        [rules addObject:@{ @"trigger": trigger, @"action": @{ @"type": @"block" } }];
        
        if (type >= SBRScriptOriginPolicyTypeDelta) {
            // Delta: add allowlist for common CDNs
            for (NSString *cdn in CommonCDNList()) {
                NSString *escaped = [NSRegularExpression escapedPatternForString:cdn];
                NSString *pattern = [NSString stringWithFormat:@".*://([^.]*\\.)?%@/", escaped];
                [rules addObject:@{
                    @"trigger": @{
                        @"url-filter": pattern,
                        @"if-domain": @[ hostOrigin ],
                        @"resource-type": @[ @"script" ]
                    },
                    @"action": @{ @"type": @"ignore-previous-rules" }
                }];
            }
            
            // Heuristic: allow cdn.<anything-with-family-name>/* where family name is the second-level label
            NSArray<NSString *> *components = [hostOrigin componentsSeparatedByString:@"."];
            if (components.count > 1) {
                NSString *family = components[components.count - 2];
                NSString *escapedFamily = [NSRegularExpression escapedPatternForString:family];
                NSString *familyPattern = [NSString stringWithFormat:@".*://cdn\\.[^/]*%@[^/]*/", escapedFamily];
                [rules addObject:@{
                    @"trigger": @{
                        @"url-filter": familyPattern,
                        @"if-domain": @[ hostOrigin ],
                        @"resource-type": @[ @"script" ]
                    },
                    @"action": @{ @"type": @"ignore-previous-rules" }
                }];
            }
        }
    }

    NSData *jsonData = [NSJSONSerialization dataWithJSONObject:rules options:0 error:nil];
    NSString *json = [[NSString alloc] initWithData:jsonData encoding:NSUTF8StringEncoding];
    NSString *name = [NSString stringWithFormat:@"net.buzzert.attractor.rules.%@", hostOrigin];
    [self applyContentRuleListJSON:json withName:name];
}

- (void)applyContentRuleListJSON:(NSString *)json withName:(NSString *)name
{
    WKUserContentController *controller = [_webViewConfiguration userContentController];
    if (_activeScriptRuleList) {
        [controller removeContentRuleList:_activeScriptRuleList];
        _activeScriptRuleList = nil;
    }

    if (!json || !name) { return; }

    WKContentRuleListStore *store = [WKContentRuleListStore defaultStore];
    [store compileContentRuleListForIdentifier:name encodedContentRuleList:json completionHandler:^(WKContentRuleList * _Nullable ruleList, NSError * _Nullable error) {
        if (error) {
            LOG_ERROR("Failed to compile content rule list: %@", error.localizedDescription);
            return;
        }
        
        dispatch_async(dispatch_get_main_queue(), ^{
            self->_activeScriptRuleList = ruleList;
            [controller addContentRuleList:ruleList];
            LOG_DEBUG("Applied content rule list: %@", name);
        });
    }];
}

@end