# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.

import contextlib
import os
import re
import textwrap

from marionette_driver.addons import Addons
from marionette_driver.errors import MarionetteException
from marionette_driver.wait import Wait
from marionette_driver import By, keys
from marionette_harness import MarionetteTestCase
from marionette_harness.runner.mixins.window_manager import WindowManagerMixin

from ping_server import PingServer


CANARY_CLIENT_ID = "c0ffeec0-ffee-c0ff-eec0-ffeec0ffeec0"
UUID_PATTERN = re.compile(
    r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$"
)


class TelemetryTestCase(WindowManagerMixin, MarionetteTestCase):
    def __init__(self, *args, **kwargs):
        """Initialize the test case and create a ping server."""
        super(TelemetryTestCase, self).__init__(*args, **kwargs)

        self.ping_server = PingServer(
            self.testvars["server_root"], self.testvars["server_url"]
        )

    def setUp(self, *args, **kwargs):
        """Set up the test case and start the ping server."""
        super(TelemetryTestCase, self).setUp(*args, **kwargs)

        # Store IDs of addons installed via self.install_addon()
        self.addon_ids = []

        with self.marionette.using_context(self.marionette.CONTEXT_CONTENT):
            self.marionette.navigate("about:about")

        self.ping_server.start()

    def disable_telemetry(self):
        """Disable the Firefox Data Collection and Use in the current browser."""
        self.marionette.instance.profile.set_persistent_preferences(
            {"datareporting.healthreport.uploadEnabled": False}
        )
        self.marionette.set_pref("datareporting.healthreport.uploadEnabled", False)

    def enable_telemetry(self):
        """Enable the Firefox Data Collection and Use in the current browser."""
        self.marionette.instance.profile.set_persistent_preferences(
            {"datareporting.healthreport.uploadEnabled": True}
        )
        self.marionette.set_pref("datareporting.healthreport.uploadEnabled", True)

    @contextlib.contextmanager
    def new_tab(self):
        """Perform operations in a new tab and then close the new tab."""

        with self.marionette.using_context(self.marionette.CONTEXT_CHROME):
            start_tab = self.marionette.current_window_handle
            new_tab = self.open_tab(focus=True)
            self.marionette.switch_to_window(new_tab)

            yield

            self.marionette.close()
            self.marionette.switch_to_window(start_tab)

    def search(self, text):
        """Perform a search via the browser's URL bar."""

        # Reload newtab to prevent urlbar from not accepting correct input
        with self.marionette.using_context(self.marionette.CONTEXT_CONTENT):
            self.marionette.navigate("about:newtab")

        with self.marionette.using_context(self.marionette.CONTEXT_CHROME):
            self.marionette.execute_script("gURLBar.select();")
            urlbar = self.marionette.find_element(By.ID, "urlbar-input")
            urlbar.send_keys(keys.Keys.DELETE)
            urlbar.send_keys(text + keys.Keys.ENTER)
        # This script checks that the search terms used for searching
        # appear in the URL when the page loads.
        script = """\
        let location = document.location.toString()
        function validate(term){
            return location.includes(term)
        }
        return arguments[0].every(validate)
        """
        # Wait for search page to load
        with self.marionette.using_context(self.marionette.CONTEXT_CONTENT):
            Wait(self.marionette, 30, 0.5).until(
                lambda driver: driver.execute_script(
                    script, script_args=[text.split()]
                ),
                message="Search terms not found, maybe the page didn't load?",
            )

    def search_in_new_tab(self, text):
        """Open a new tab and perform a search via the browser's URL bar,
        then close the new tab."""

        with self.new_tab():
            self.search(text)

    def assertIsValidUUID(self, value):
        """Check if the given UUID is valid."""

        self.assertIsNotNone(value)
        self.assertNotEqual(value, "")

        # Check for client ID that is used when Telemetry upload is disabled
        self.assertNotEqual(value, CANARY_CLIENT_ID, msg="UUID is CANARY CLIENT ID")

        self.assertIsNotNone(
            re.match(UUID_PATTERN, value),
            msg="UUID does not match regular expression",
        )

    def wait_for_pings(self, action_func, ping_filter, count, ping_server=None):
        """Call the given action and wait for pings to come in and return
        the `count` number of pings, that match the given filter.
        """

        if ping_server is None:
            ping_server = self.ping_server

        # Keep track of the current number of pings
        current_num_pings = len(ping_server.pings)

        # New list to store new pings that satisfy the filter
        filtered_pings = []

        def wait_func(*args, **kwargs):
            # Ignore existing pings in ping_server.pings
            new_pings = ping_server.pings[current_num_pings:]

            # Filter pings to make sure we wait for the correct ping type
            filtered_pings[:] = [p for p in new_pings if ping_filter(p)]

            return len(filtered_pings) >= count

        self.logger.info(
            "wait_for_pings running action '{action}'.".format(
                action=action_func.__name__
            )
        )

        # Call given action and wait for a ping
        action_func()

        try:
            Wait(self.marionette, 60).until(wait_func)
        except Exception as e:
            self.fail("Error waiting for ping: {}".format(e.message))

        return filtered_pings[:count]

    def wait_for_ping(self, action_func, ping_filter, ping_server=None):
        """Call wait_for_pings() with the given action_func and ping_filter and
        return the first result.
        """
        [ping] = self.wait_for_pings(
            action_func, ping_filter, 1, ping_server=ping_server
        )
        return ping

    def restart_browser(self):
        """Restarts browser while maintaining the same profile."""
        return self.marionette.restart(clean=False, in_app=True)

    def start_browser(self):
        """Start the browser."""
        return self.marionette.start_session()

    def quit_browser(self):
        """Quit the browser."""
        return self.marionette.quit(in_app=True)

    def install_addon(self):
        """Install a minimal addon and add its ID to self.addon_ids."""

        resources_dir = os.path.join(os.path.dirname(__file__), "resources")
        addon_path = os.path.abspath(os.path.join(resources_dir, "helloworld"))

        try:
            # Ensure the Environment has init'd so the installed addon
            # triggers an "environment-change" ping.
            script = """\
            let [resolve] = arguments;
            Cu.import("resource://gre/modules/TelemetryEnvironment.jsm");
            TelemetryEnvironment.onInitialized().then(resolve);
            """

            with self.marionette.using_context(self.marionette.CONTEXT_CHROME):
                self.marionette.execute_async_script(textwrap.dedent(script))

            addons = Addons(self.marionette)
            addon_id = addons.install(addon_path, temp=True)
        except MarionetteException as e:
            self.fail("{} - Error installing addon: {} - ".format(e.cause, e.message))
        else:
            self.addon_ids.append(addon_id)

    def set_persistent_profile_preferences(self, preferences):
        """Wrapper for setting persistent preferences on a user profile"""
        return self.marionette.instance.profile.set_persistent_preferences(preferences)

    def set_preferences(self, preferences):
        """Wrapper for setting persistent preferences on a user profile"""
        return self.marionette.set_prefs(preferences)

    @property
    def client_id(self):
        """Return the ID of the current client."""
        with self.marionette.using_context(self.marionette.CONTEXT_CHROME):
            return self.marionette.execute_script(
                'Cu.import("resource://gre/modules/ClientID.jsm");'
                "return ClientID.getCachedClientID();"
            )

    @property
    def subsession_id(self):
        """Return the ID of the current subsession."""
        with self.marionette.using_context(self.marionette.CONTEXT_CHROME):
            ping_data = self.marionette.execute_script(
                'Cu.import("resource://gre/modules/TelemetryController.jsm");'
                "return TelemetryController.getCurrentPingData(true);"
            )
            return ping_data[u"payload"][u"info"][u"subsessionId"]

    def tearDown(self, *args, **kwargs):
        """Stop the ping server and tear down the testcase."""
        super(TelemetryTestCase, self).tearDown()
        self.ping_server.stop()
        self.marionette.quit(clean=True)
