import json
import os
import threading
import time
import traceback
import uuid

from urllib.parse import urljoin

errors = None
marionette = None
pytestrunner = None

here = os.path.dirname(__file__)

from .base import (CallbackHandler,
                   CrashtestExecutor,
                   RefTestExecutor,
                   RefTestImplementation,
                   TestharnessExecutor,
                   TimedRunner,
                   WdspecExecutor,
                   WdspecProtocol,
                   get_pages,
                   strip_server)
from .protocol import (ActionSequenceProtocolPart,
                       AssertsProtocolPart,
                       BaseProtocolPart,
                       TestharnessProtocolPart,
                       PrefsProtocolPart,
                       Protocol,
                       StorageProtocolPart,
                       SelectorProtocolPart,
                       ClickProtocolPart,
                       CookiesProtocolPart,
                       SendKeysProtocolPart,
                       TestDriverProtocolPart,
                       CoverageProtocolPart,
                       GenerateTestReportProtocolPart,
                       VirtualAuthenticatorProtocolPart,
                       SetPermissionProtocolPart,
                       PrintProtocolPart,
                       DebugProtocolPart)


def do_delayed_imports():
    global errors, marionette, Addons

    from marionette_driver import marionette, errors
    from marionette_driver.addons import Addons


def _switch_to_window(marionette, handle):
    """Switch to the specified window; subsequent commands will be
    directed at the new window.

    This is a workaround for issue 24924[0]; marionettedriver 3.1.0 dropped the
    'name' parameter from its switch_to_window command, but it is still needed
    for at least Firefox 79.

    [0]: https://github.com/web-platform-tests/wpt/issues/24924

    :param marionette: The Marionette instance
    :param handle: The id of the window to switch to.
    """
    marionette._send_message("WebDriver:SwitchToWindow",
                             {"handle": handle, "name": handle, "focus": True})
    marionette.window = handle


class MarionetteBaseProtocolPart(BaseProtocolPart):
    def __init__(self, parent):
        super(MarionetteBaseProtocolPart, self).__init__(parent)
        self.timeout = None

    def setup(self):
        self.marionette = self.parent.marionette

    def execute_script(self, script, asynchronous=False):
        method = self.marionette.execute_async_script if asynchronous else self.marionette.execute_script
        return method(script, new_sandbox=False, sandbox=None)

    def set_timeout(self, timeout):
        """Set the Marionette script timeout.

        :param timeout: Script timeout in seconds

        """
        if timeout != self.timeout:
            self.marionette.timeout.script = timeout
            self.timeout = timeout

    @property
    def current_window(self):
        return self.marionette.current_window_handle

    def set_window(self, handle):
        _switch_to_window(self.marionette, handle)

    def window_handles(self):
        return self.marionette.window_handles

    def load(self, url):
        self.marionette.navigate(url)

    def wait(self):
        try:
            socket_timeout = self.marionette.client.socket_timeout
        except AttributeError:
            # This can happen if there was a crash
            return
        if socket_timeout:
            try:
                self.marionette.timeout.script = socket_timeout / 2
            except IOError:
                self.logger.debug("Socket closed")
                return

        while True:
            try:
                return self.marionette.execute_async_script("""let callback = arguments[arguments.length - 1];
addEventListener("__test_restart", e => {e.preventDefault(); callback(true)})""")
            except errors.NoSuchWindowException:
                # The window closed
                break
            except errors.ScriptTimeoutException:
                self.logger.debug("Script timed out")
                pass
            except errors.JavascriptException as e:
                # This can happen if we navigate, but just keep going
                self.logger.debug(e)
                pass
            except IOError:
                self.logger.debug("Socket closed")
                break
            except Exception:
                self.logger.warning(traceback.format_exc())
                break
        return False


class MarionetteTestharnessProtocolPart(TestharnessProtocolPart):
    def __init__(self, parent):
        super(MarionetteTestharnessProtocolPart, self).__init__(parent)
        self.runner_handle = None
        with open(os.path.join(here, "runner.js")) as f:
            self.runner_script = f.read()
        with open(os.path.join(here, "window-loaded.js")) as f:
            self.window_loaded_script = f.read()

    def setup(self):
        self.marionette = self.parent.marionette

    def load_runner(self, url_protocol):
        # Check if we previously had a test window open, and if we did make sure it's closed
        if self.runner_handle:
            self._close_windows()
        url = urljoin(self.parent.executor.server_url(url_protocol), "/testharness_runner.html")
        self.logger.debug("Loading %s" % url)
        try:
            self.dismiss_alert(lambda: self.marionette.navigate(url))
        except Exception:
            self.logger.critical(
                "Loading initial page %s failed. Ensure that the "
                "there are no other programs bound to this port and "
                "that your firewall rules or network setup does not "
                "prevent access.\n%s" % (url, traceback.format_exc()))
            raise
        self.runner_handle = self.marionette.current_window_handle
        format_map = {"title": threading.current_thread().name.replace("'", '"')}
        self.parent.base.execute_script(self.runner_script % format_map)

    def _close_windows(self):
        handles = self.marionette.window_handles
        runner_handle = None
        try:
            handles.remove(self.runner_handle)
            runner_handle = self.runner_handle
        except ValueError:
            # The runner window probably changed id but we can restore it
            # This isn't supposed to happen, but marionette ids are not yet stable
            # We assume that the first handle returned corresponds to the runner,
            # but it hopefully doesn't matter too much if that assumption is
            # wrong since we reload the runner in that tab anyway.
            runner_handle = handles.pop(0)
            self.logger.info("Changing harness_window to %s" % runner_handle)

        for handle in handles:
            try:
                self.logger.info("Closing window %s" % handle)
                _switch_to_window(self.marionette, handle)
                self.dismiss_alert(lambda: self.marionette.close())
            except errors.NoSuchWindowException:
                # We might have raced with the previous test to close this
                # window, skip it.
                pass
        _switch_to_window(self.marionette, runner_handle)
        return runner_handle

    def close_old_windows(self, url_protocol):
        runner_handle = self._close_windows()
        if runner_handle != self.runner_handle:
            self.load_runner(url_protocol)
        return self.runner_handle

    def dismiss_alert(self, f):
        while True:
            try:
                f()
            except errors.UnexpectedAlertOpen:
                alert = self.marionette.switch_to_alert()
                try:
                    alert.dismiss()
                except errors.NoAlertPresentException:
                    pass
            else:
                break

    def get_test_window(self, window_id, parent, timeout=5):
        """Find the test window amongst all the open windows.
        This is assumed to be either the named window or the one after the parent in the list of
        window handles

        :param window_id: The DOM name of the Window
        :param parent: The handle of the runner window
        :param timeout: The time in seconds to wait for the window to appear. This is because in
                        some implementations there's a race between calling window.open and the
                        window being added to the list of WebDriver accessible windows."""
        test_window = None
        end_time = time.time() + timeout
        while time.time() < end_time:
            if window_id:
                try:
                    # Try this, it's in Level 1 but nothing supports it yet
                    win_s = self.parent.base.execute_script("return window['%s'];" % self.window_id)
                    win_obj = json.loads(win_s)
                    test_window = win_obj["window-fcc6-11e5-b4f8-330a88ab9d7f"]
                except Exception:
                    pass

            if test_window is None:
                handles = self.marionette.window_handles
                if len(handles) == 2:
                    test_window = next(iter(set(handles) - {parent}))
                elif len(handles) > 2 and handles[0] == parent:
                    # Hope the first one here is the test window
                    test_window = handles[1]

            if test_window is not None:
                assert test_window != parent
                return test_window

            time.sleep(0.1)

        raise Exception("unable to find test window")

    def test_window_loaded(self):
        """Wait until the page in the new window has been loaded.

        Hereby ignore Javascript execptions that are thrown when
        the document has been unloaded due to a process change.
        """
        while True:
            try:
                self.parent.base.execute_script(self.window_loaded_script, asynchronous=True)
                break
            except errors.JavascriptException:
                pass


class MarionettePrefsProtocolPart(PrefsProtocolPart):
    def setup(self):
        self.marionette = self.parent.marionette

    def set(self, name, value):
        if value.lower() not in ("true", "false"):
            try:
                int(value)
            except ValueError:
                value = "'%s'" % value
        else:
            value = value.lower()

        self.logger.info("Setting pref %s (%s)" % (name, value))

        script = """
            let prefInterface = Components.classes["@mozilla.org/preferences-service;1"]
                                          .getService(Components.interfaces.nsIPrefBranch);
            let pref = '%s';
            let type = prefInterface.getPrefType(pref);
            let value = %s;
            switch(type) {
                case prefInterface.PREF_STRING:
                    prefInterface.setCharPref(pref, value);
                    break;
                case prefInterface.PREF_BOOL:
                    prefInterface.setBoolPref(pref, value);
                    break;
                case prefInterface.PREF_INT:
                    prefInterface.setIntPref(pref, value);
                    break;
                case prefInterface.PREF_INVALID:
                    // Pref doesn't seem to be defined already; guess at the
                    // right way to set it based on the type of value we have.
                    switch (typeof value) {
                        case "boolean":
                            prefInterface.setBoolPref(pref, value);
                            break;
                        case "string":
                            prefInterface.setCharPref(pref, value);
                            break;
                        case "number":
                            prefInterface.setIntPref(pref, value);
                            break;
                        default:
                            throw new Error("Unknown pref value type: " + (typeof value));
                    }
                    break;
                default:
                    throw new Error("Unknown pref type " + type);
            }
            """ % (name, value)
        with self.marionette.using_context(self.marionette.CONTEXT_CHROME):
            self.marionette.execute_script(script)

    def clear(self, name):
        self.logger.info("Clearing pref %s" % (name))
        script = """
            let prefInterface = Components.classes["@mozilla.org/preferences-service;1"]
                                          .getService(Components.interfaces.nsIPrefBranch);
            let pref = '%s';
            prefInterface.clearUserPref(pref);
            """ % name
        with self.marionette.using_context(self.marionette.CONTEXT_CHROME):
            self.marionette.execute_script(script)

    def get(self, name):
        script = """
            let prefInterface = Components.classes["@mozilla.org/preferences-service;1"]
                                          .getService(Components.interfaces.nsIPrefBranch);
            let pref = '%s';
            let type = prefInterface.getPrefType(pref);
            switch(type) {
                case prefInterface.PREF_STRING:
                    return prefInterface.getCharPref(pref);
                case prefInterface.PREF_BOOL:
                    return prefInterface.getBoolPref(pref);
                case prefInterface.PREF_INT:
                    return prefInterface.getIntPref(pref);
                case prefInterface.PREF_INVALID:
                    return null;
            }
            """ % name
        with self.marionette.using_context(self.marionette.CONTEXT_CHROME):
            self.marionette.execute_script(script)


class MarionetteStorageProtocolPart(StorageProtocolPart):
    def setup(self):
        self.marionette = self.parent.marionette

    def clear_origin(self, url):
        self.logger.info("Clearing origin %s" % (url))
        script = """
            let url = '%s';
            let uri = Components.classes["@mozilla.org/network/io-service;1"]
                                .getService(Ci.nsIIOService)
                                .newURI(url);
            let ssm = Components.classes["@mozilla.org/scriptsecuritymanager;1"]
                                .getService(Ci.nsIScriptSecurityManager);
            let principal = ssm.createContentPrincipal(uri, {});
            let qms = Components.classes["@mozilla.org/dom/quota-manager-service;1"]
                                .getService(Components.interfaces.nsIQuotaManagerService);
            qms.clearStoragesForPrincipal(principal, "default", null, true);
            """ % url
        with self.marionette.using_context(self.marionette.CONTEXT_CHROME):
            self.marionette.execute_script(script)


class MarionetteAssertsProtocolPart(AssertsProtocolPart):
    def setup(self):
        self.assert_count = {"chrome": 0, "content": 0}
        self.chrome_assert_count = 0
        self.marionette = self.parent.marionette

    def get(self):
        script = """
        debug = Cc["@mozilla.org/xpcom/debug;1"].getService(Ci.nsIDebug2);
        if (debug.isDebugBuild) {
          return debug.assertionCount;
        }
        return 0;
        """

        def get_count(context, **kwargs):
            try:
                context_count = self.marionette.execute_script(script, **kwargs)
                if context_count:
                    self.parent.logger.info("Got %s assert count %s" % (context, context_count))
                    test_count = context_count - self.assert_count[context]
                    self.assert_count[context] = context_count
                    return test_count
            except errors.NoSuchWindowException:
                # If the window was already closed
                self.parent.logger.warning("Failed to get assertion count; window was closed")
            except (errors.MarionetteException, IOError):
                # This usually happens if the process crashed
                pass

        counts = []
        with self.marionette.using_context(self.marionette.CONTEXT_CHROME):
            counts.append(get_count("chrome"))
        if self.parent.e10s:
            counts.append(get_count("content", sandbox="system"))

        counts = [item for item in counts if item is not None]

        if not counts:
            return None

        return sum(counts)


class MarionetteSelectorProtocolPart(SelectorProtocolPart):
    def setup(self):
        self.marionette = self.parent.marionette

    def elements_by_selector(self, selector):
        return self.marionette.find_elements("css selector", selector)


class MarionetteClickProtocolPart(ClickProtocolPart):
    def setup(self):
        self.marionette = self.parent.marionette

    def element(self, element):
        return element.click()


class MarionetteCookiesProtocolPart(CookiesProtocolPart):
    def setup(self):
        self.marionette = self.parent.marionette

    def delete_all_cookies(self):
        self.logger.info("Deleting all cookies")
        return self.marionette.delete_all_cookies()


class MarionetteSendKeysProtocolPart(SendKeysProtocolPart):
    def setup(self):
        self.marionette = self.parent.marionette

    def send_keys(self, element, keys):
        return element.send_keys(keys)


class MarionetteActionSequenceProtocolPart(ActionSequenceProtocolPart):
    def setup(self):
        self.marionette = self.parent.marionette

    def send_actions(self, actions):
        actions = self.marionette._to_json(actions)
        self.logger.info(actions)
        self.marionette._send_message("WebDriver:PerformActions", actions)


class MarionetteTestDriverProtocolPart(TestDriverProtocolPart):
    def setup(self):
        self.marionette = self.parent.marionette

    def send_message(self, cmd_id, message_type, status, message=None):
        obj = {
            "cmd_id": cmd_id,
            "type": "testdriver-%s" % str(message_type),
            "status": str(status)
        }
        if message:
            obj["message"] = str(message)
        self.parent.base.execute_script("window.postMessage(%s, '*')" % json.dumps(obj))

    def _switch_to_frame(self, index_or_elem):
        try:
            self.marionette.switch_to_frame(index_or_elem)
        except (errors.NoSuchFrameException,
                errors.StaleElementException) as e:
            raise ValueError from e

    def _switch_to_parent_frame(self):
        self.marionette.switch_to_parent_frame()


class MarionetteCoverageProtocolPart(CoverageProtocolPart):
    def setup(self):
        self.marionette = self.parent.marionette

        if not self.parent.ccov:
            self.is_enabled = False
            return

        script = """
            const {PerTestCoverageUtils} = ChromeUtils.import("chrome://remote/content/marionette/PerTestCoverageUtils.jsm");
            return PerTestCoverageUtils.enabled;
            """
        with self.marionette.using_context(self.marionette.CONTEXT_CHROME):
            self.is_enabled = self.marionette.execute_script(script)

    def reset(self):
        script = """
            var callback = arguments[arguments.length - 1];

            const {PerTestCoverageUtils} = ChromeUtils.import("chrome://remote/content/marionette/PerTestCoverageUtils.jsm");
            PerTestCoverageUtils.beforeTest().then(callback, callback);
            """
        with self.marionette.using_context(self.marionette.CONTEXT_CHROME):
            try:
                error = self.marionette.execute_async_script(script)
                if error is not None:
                    raise Exception('Failure while resetting counters: %s' % json.dumps(error))
            except (errors.MarionetteException, IOError):
                # This usually happens if the process crashed
                pass

    def dump(self):
        if len(self.marionette.window_handles):
            handle = self.marionette.window_handles[0]
            _switch_to_window(self.marionette, handle)

        script = """
            var callback = arguments[arguments.length - 1];

            const {PerTestCoverageUtils} = ChromeUtils.import("chrome://remote/content/marionette/PerTestCoverageUtils.jsm");
            PerTestCoverageUtils.afterTest().then(callback, callback);
            """
        with self.marionette.using_context(self.marionette.CONTEXT_CHROME):
            try:
                error = self.marionette.execute_async_script(script)
                if error is not None:
                    raise Exception('Failure while dumping counters: %s' % json.dumps(error))
            except (errors.MarionetteException, IOError):
                # This usually happens if the process crashed
                pass

class MarionetteGenerateTestReportProtocolPart(GenerateTestReportProtocolPart):
    def setup(self):
        self.marionette = self.parent.marionette

    def generate_test_report(self, config):
        raise NotImplementedError("generate_test_report not yet implemented")

class MarionetteVirtualAuthenticatorProtocolPart(VirtualAuthenticatorProtocolPart):
    def setup(self):
        self.marionette = self.parent.marionette

    def add_virtual_authenticator(self, config):
        raise NotImplementedError("add_virtual_authenticator not yet implemented")

    def remove_virtual_authenticator(self, authenticator_id):
        raise NotImplementedError("remove_virtual_authenticator not yet implemented")

    def add_credential(self, authenticator_id, credential):
        raise NotImplementedError("add_credential not yet implemented")

    def get_credentials(self, authenticator_id):
        raise NotImplementedError("get_credentials not yet implemented")

    def remove_credential(self, authenticator_id, credential_id):
        raise NotImplementedError("remove_credential not yet implemented")

    def remove_all_credentials(self, authenticator_id):
        raise NotImplementedError("remove_all_credentials not yet implemented")

    def set_user_verified(self, authenticator_id, uv):
        raise NotImplementedError("set_user_verified not yet implemented")


class MarionetteSetPermissionProtocolPart(SetPermissionProtocolPart):
    def setup(self):
        self.marionette = self.parent.marionette

    def set_permission(self, descriptor, state, one_realm):
        body = {
            "descriptor": descriptor,
            "state": state,
        }
        if one_realm is not None:
            body["oneRealm"] = one_realm
        try:
            self.marionette._send_message("WebDriver:SetPermission", body)
        except errors.UnsupportedOperationException:
            raise NotImplementedError("set_permission not yet implemented")


class MarionettePrintProtocolPart(PrintProtocolPart):
    def setup(self):
        self.marionette = self.parent.marionette
        self.runner_handle = None

    def load_runner(self):
        url = urljoin(self.parent.executor.server_url("http"), "/print_reftest_runner.html")
        self.logger.debug("Loading %s" % url)
        try:
            self.marionette.navigate(url)
        except Exception as e:
            self.logger.critical(
                "Loading initial page %s failed. Ensure that the "
                "there are no other programs bound to this port and "
                "that your firewall rules or network setup does not "
                "prevent access.\n%s" % (url, traceback.format_exc(e)))
            raise
        self.runner_handle = self.marionette.current_window_handle

    def render_as_pdf(self, width, height):
        margin = 0.5 * 2.54
        body = {
            "page": {
                "width": width,
                "height": height
            },
            "margin": {
                "left": margin,
                "right": margin,
                "top": margin,
                "bottom": margin,
            },
            "shrinkToFit": False,
            "printBackground": True,
        }
        return self.marionette._send_message("WebDriver:Print", body, key="value")

    def pdf_to_png(self, pdf_base64, page_ranges):
        handle = self.marionette.current_window_handle
        _switch_to_window(self.marionette, self.runner_handle)
        try: