from selenium.common.exceptions import TimeoutException
from selenium.webdriver import Firefox, FirefoxProfile
from selenium.webdriver.common.proxy import Proxy, ProxyType
from selenium.webdriver.firefox.options import Options

from setting import *


class Browser(Firefox):

    def __init__(self, user_agent, headless=True, tor=True):
        options = Options()
        options.add_argument('--user-agent="{}"'.format(user_agent))

        if headless:
            options.add_argument('--headless')
            options.add_argument('--window-size=1920,1080')

        profile = None
        if tor:
            proxy_url = PROXY.replace('http://', '', 1)
            proxy = Proxy({
                'proxyType': ProxyType.MANUAL,
                'httpProxy': proxy_url,
                'ftpProxy': proxy_url,
                'sslProxy': proxy_url,
                'noProxy': ''
            })
            profile = FirefoxProfile()
            profile.set_proxy(proxy)

        super().__init__(
            executable_path=GECKODRIVER_PATH, firefox_binary=FIREFOX_PATH, options=options, firefox_profile=profile
        )
        self.implicitly_wait(5)
        self.set_page_load_timeout(20)

    def get(self, url):
        try:
            super().get(url)
        except TimeoutException:
            self.execute_script('window.stop()')