# coding:utf-8
from selenium import webdriver
from selenium.webdriver.chrome.options import Options
import time
from tools_hjh.ThreadPool import ThreadPool
import os
from tools_hjh import Tools
from tools_hjh.Tools import mkdir, rm


def main():
    pass


class Chrome():
    """ 使用浏览器解析url，返回源码
        __init__.param：
            chrome_path: chrome.exe路径
            chromedriver_path: chromedriver.exe路径
    """

    def __init__(self, chrome_path, chromedriver_path, is_hidden=False, is_display_picture=True, headers=None):
        chrome_options = Options()
        # chrome_options.add_argument("--single-process")
        # chrome_options.add_argument("--disable-site-isolation-trials")

        if is_hidden:
            chrome_options.add_argument("--headless")
        chrome_options.binary_location = chrome_path
        if not is_display_picture:
            chrome_options.add_experimental_option('prefs', {'profile.managed_default_content_settings.images': 2})
        if headers is not None:
            for key, value in headers.items():
                header = key + "=" + value
                chrome_options.add_argument(header)
        self.chrome = webdriver.Chrome(chromedriver_path, options=chrome_options)
        self.status = 0  # 未被占用
    
    def set_cookies(self, cookies, cookie_domian):
        self.get(cookie_domian)
        for cookie in cookies:
            self.chrome.add_cookie(cookie)
        
    def close(self):
        self.chrome.quit()
        
    def get(self, url):
        self.status = 1  # 被占用
        self.chrome.get(url)
        text = self.chrome.page_source
        self.status = 0
        return text
    
    def download(self, url, filepath, replace=False):
        """ 下载请求的文件, 返回文件大小 """
        filepath = filepath.replace('\\', '/')
        
        if '/' in filepath:
            path = filepath.rsplit('/', 1)[0] + '/'
            mkdir(path)
            
        if os.path.exists(filepath) and replace:
            rm(filepath)
        elif os.path.exists(filepath) and not replace:
            return 0

        try:
            self.get(url)
            screenshot_data = self.chrome.get_screenshot_as_png()  # 获取二进制数据
            with open(filepath, "wb") as f:
                f.write(screenshot_data)
        except Exception as _:
            pass
        finally:
            try:
                f.close()
            except:
                pass
        
        download_size = os.path.getsize(filepath)

        return download_size
    
    def get_status(self):
        return self.status


class ChromePool():

    def __init__(self, maxSize, chrome_path, chromedriver_path, is_hidden=False, is_display_picture=True, cache=False, headers=None):
        self.cache = cache
        self.maxSize = maxSize
        self.pool = []
        for _ in range(0, maxSize):
            self.pool.append(Chrome(chrome_path, chromedriver_path, is_hidden, is_display_picture, headers))
        self.openSize = 0
        
    def set_cookies(self, cookies, cookie_domian):

        def set_one(chrome, cookies, cookie_domian):
            chrome.set_cookies(cookies, cookie_domian)

        tp = ThreadPool(self.maxSize)
        for chrome in self.pool:
            tp.run(set_one, (chrome, cookies, cookie_domian))
        tp.wait()
        
    def get(self, url):
        while self.openSize == self.maxSize:
            time.sleep(0.2)
        for chrome in self.pool:
            if chrome.get_status() == 0:
                self.openSize = self.openSize + 1
                if self.cache:
                    url_split = url.split('/')
                    host = url_split[0] + '//' + url_split[2] + '/'
                    path_ = 'tmp/' + url.replace(host, '').rstrip('/')
                    if os.path.exists(path_):
                        text = Tools.cat(path_)
                    else:
                        text = chrome.get(url)
                        Tools.echo(text, path_)
                else:
                    text = chrome.get(url)
                self.openSize = self.openSize - 1
                return text
        time.sleep(0.2)
        return self.get(url)

    def download(self, url, filepath, replace=False):
        while self.openSize == self.maxSize:
            time.sleep(0.2)
        for chrome in self.pool:
            if chrome.get_status() == 0:
                self.openSize = self.openSize + 1
                size = chrome.download(url, filepath, replace)
                self.openSize = self.openSize - 1
                return size
        time.sleep(0.2)
        return self.download(url, filepath, replace)
    
    def close(self):
        while self.openSize > 0:
            time.sleep(0.2)
        for c in self.pool:
            c.close()
            

if __name__ == '__main__':
    main()
