# SPDX-License-Identifier: CC0-1.0

"""
Haketilo unit tests - building pattern tree and putting it in a content script
"""

# This file is part of Haketilo
#
# Copyright (C) 2021, Wojtek Kosior <koszko@koszko.org>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the CC0 1.0 Universal License as published by
# the Creative Commons Corporation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# CC0 1.0 Universal License for more details.

import pytest
import re
import json
from selenium.webdriver.support.ui import WebDriverWait

from ..script_loader import load_script

def simple_sample_mapping(patterns, fruit):
    if type(patterns) is not list:
        patterns = [patterns]
    payloads = dict([(p, {'identifier': f'{fruit}-{p}'}) for p in patterns])
    return {
        'source_copyright': [],
        'type': 'mapping',
        'identifier': f'inject-{fruit}',
        'payloads': payloads
    }

content_script_tree_re = re.compile(r'this.haketilo_pattern_tree = (.*);')
def extract_tree_data(content_script_text):
    return json.loads(content_script_tree_re.search(content_script_text)[1])

content_script_mapping_re = re.compile(r'this.haketilo_mappings = (.*);')
def extract_mappings_data(content_script_text):
    return json.loads(content_script_mapping_re.search(content_script_text)[1])

# Fields that are not relevant for testing are omitted from these mapping
# definitions.
sample_mappings = [simple_sample_mapping(pats, fruit) for pats, fruit in [
    (['https://gotmyowndoma.in/index.html',
      'http://gotmyowndoma.in/index.html'], 'banana'),
    (['https://***.gotmyowndoma.in/index.html',
      'https://**.gotmyowndoma.in/index.html',
      'https://*.gotmyowndoma.in/index.html',
      'https://gotmyowndoma.in/index.html'], 'orange'),
    ('https://gotmyowndoma.in/index.html/***', 'grape'),
    ('http://gotmyowndoma.in/index.html/***', 'melon'),
    ('https://gotmyowndoma.in/index.html', 'peach'),
    ('https://gotmyowndoma.in/*', 'pear'),
    ('https://gotmyowndoma.in/**', 'raspberry'),
    ('https://gotmyowndoma.in/***', 'strawberry'),
    ('https://***.gotmyowndoma.in/index.html', 'apple'),
    ('https://***.gotmyowndoma.in/*', 'avocado'),
    ('https://***.gotmyowndoma.in/**', 'papaya'),
    ('https://***.gotmyowndoma.in/***', 'kiwi')
]]

# Even though patterns_query_manager.js is normally meant to run from background
# page, tests can be as well performed running it from extension's bundled page.
@pytest.mark.get_page('https://gotmyowndoma.in')
def test_pqm_tree_building(driver, execute_in_page):
    """
    patterns_query_manager.js tracks Haketilo's internal database and builds a
    constantly-updated pattern tree based on its contents. Mock the database and
    verify tree building works properly.
    """
    execute_in_page(load_script('background/patterns_query_manager.js'))
    # Mock IndexedDB and build patterns tree.
    execute_in_page(
        '''
        const initial_mappings = arguments[0]
        let mappingchange;
        function track_mock(cb)
        {
            mappingchange = cb;

            return [{}, initial_mappings];
        }
        haketilodb.track.mappings = track_mock;

        let last_script;
        let unregister_called = 0;
        async function register_mock(injection)
        {
            await new Promise(resolve => setTimeout(resolve, 1));
            last_script = injection.js[0].code;
            return {unregister: () => unregister_called++};
        }
        browser = {contentScripts: {register: register_mock}};

        returnval(start());
        ''',
        sample_mappings[0:2])

    found, tree, content_script, deregistrations = execute_in_page(
        '''
        returnval([pqt.search(tree, arguments[0]).next().value,
                   tree, last_script, unregister_called]);
        ''',
        'https://gotmyowndoma.in/index.html')
    best_pattern = 'https://gotmyowndoma.in/index.html'
    assert found == \
        dict([(f'inject-{fruit}', {'identifier': f'{fruit}-{best_pattern}'})
              for fruit in ('banana', 'orange')])
    assert tree == extract_tree_data(content_script)
    assert deregistrations == 0

    def condition_mappings_added(driver):
        last_script = execute_in_page('returnval(last_script);')
        return all([m['identifier'] in last_script for m in sample_mappings])

    execute_in_page(
        '''
        for (const mapping of arguments[0])
            mappingchange({key: mapping.identifier, new_val: mapping});
        ''',
        sample_mappings[2:])
    WebDriverWait(driver, 10).until(condition_mappings_added)

    odd = [m['identifier'] for i, m in enumerate(sample_mappings) if i % 2]
    even = [m['identifier'] for i, m in enumerate(sample_mappings) if 1 - i % 2]

    def condition_odd_removed(driver):
        last_script = execute_in_page('returnval(last_script);')
        return (all([id not in last_script for id in odd]) and
                all([id in last_script for id in even]))

    def condition_all_removed(driver):
        content_script = execute_in_page('returnval(last_script);')
        return extract_tree_data(content_script) == {}

    execute_in_page(
        '''
        arguments[0].forEach(identifier => mappingchange({key: identifier}));
        ''',
        odd)

    WebDriverWait(driver, 10).until(condition_odd_removed)

    execute_in_page(
        '''
        arguments[0].forEach(identifier => mappingchange({key: identifier}));
        ''',
        even)

    WebDriverWait(driver, 10).until(condition_all_removed)

content_js = '''
let already_run = false;
this.haketilo_content_script_main = function() {
    if (already_run)
        return;
    already_run = true;
    document.documentElement.innerHTML = "<body><div id='tree-json'>";
    document.getElementById("tree-json").innerText =
        JSON.stringify(this.haketilo_pattern_tree);
}
if (this.haketilo_pattern_tree !== undefined)
    this.haketilo_content_script_main();
'''

def background_js():
    pqm_js = load_script('background/patterns_query_manager.js',
                         "#IMPORT background/broadcast_broker.js")
    return pqm_js + '; broadcast_broker.start(); start();'

@pytest.mark.ext_data({
    'content_script':    content_js,
    'background_script': background_js
})
@pytest.mark.usefixtures('webextension')
def test_pqm_script_injection(driver, execute_in_page):
    # Let's open a normal page in a second window. Window 0 will be used to make
    # changed to IndexedDB and window 1 to test the working of content scripts.
    driver.execute_script('window.open("about:blank", "_blank");')
    windows = [*driver.window_handles]
    assert len(windows) == 2

    def run_content_script():
        driver.switch_to.window(windows[1])
        driver.get('https://gotmyowndoma.in/index.html')
        windows[1] = driver.window_handles[1]
        return driver.execute_script(
            '''
            return (document.getElementById("tree-json") || {}).innerText;
            ''')

    for attempt in range(10):
        json_txt = run_content_script()
        if json.loads(json_txt) == {}:
            break;
        assert attempt != 9

    driver.switch_to.window(windows[0])
    execute_in_page(load_script('common/indexeddb.js'))

    sample_data = {
        'mappings': dict([(sm['identifier'], {'1.0': sm})
                          for sm in sample_mappings]),
        'resources': {},
        'files': {}
    }
    execute_in_page('returnval(save_items(arguments[0]));', sample_data)

    for attempt in range(10):
        tree_json = run_content_script()
        json.loads(tree_json)
        if all([m['identifier'] in tree_json for m in sample_mappings]):
            break
        assert attempt != 9

    driver.switch_to.window(windows[0])
    execute_in_page(
        '''{
        const identifiers = arguments[0];
        async function remove_items()
        {
            const ctx = await start_items_transaction(["mappings"], {});
            for (const id of identifiers)
                await remove_mapping(id, ctx);
            await finalize_transaction(ctx);
        }
        returnval(remove_items());
        }''',
        [sm['identifier'] for sm in sample_mappings])

    for attempt in range(10):
        if json.loads(run_content_script()) == {}:
            break
        assert attempt != 9