import datetime
import os
from selenium import webdriver
from selenium.common.exceptions import WebDriverException
import config


class Browser(object):
    def __init__(self):
        pass

    def capture(self):
        pass


class Chrome(Browser):
    def __init__(self, html):
        self.html = html
        self.browser = self._runChrome()

    def _runChrome(self):
        options = webdriver.ChromeOptions()
        options.add_argument('headless')
        options.add_argument('no-sandbox')
        driver = webdriver.Chrome(options=options)
        driver.set_window_size(config.WIDTH, config.HEIGHT)
        return driver

    def capture(self):
        """
        현재 상태를 크롬으로 캡쳐한 파일 반환
        :return: image_path
        """
        date_time = datetime.datetime.now()
        date_time = date_time.strftime('%Y-%m-%d-%H-%M-%S')
        file_path = 'image/' + str(date_time) + '.png'
        now_path = 'file://' + os.path.join(os.getcwd(), self.html)
        try:
            self.browser.get(now_path)
        except WebDriverException:
            self.browser = self._runChrome()
            self.browser.get(now_path)
        self.browser.save_screenshot(file_path)
        return file_path