import json
import logging as logger
import os
import random

from PIL import Image
from selenium.webdriver import DesiredCapabilities
from selenium.webdriver.firefox.options import Options
from selenium.webdriver.remote.webdriver import WebDriver

from WebsiteSpider.utils.http_utils import http_post
from WebsiteSpider.utils.string_utils import get_legal_file_name, get_str_md5


def check_status(drive_path):
    api = drive_path + '/graphql'
    post_body = '{"query": "{ nodesInfo { nodes { status } } }"}'
    response = http_post(api, post_body)
    data_body = json.loads(response.content.decode())
    status = data_body['data']['nodesInfo']['nodes'][0]['status']
    return status


def check_session(drive_path):
    api = drive_path + '/graphql'
    post_body = '{"query": "{ grid { maxSession, sessionCount } }"}'
    try:
        response = http_post(api, post_body)
        data_body = json.loads(response.content.decode())
        session_info = data_body['data']['grid']
        return session_info
    except Exception as e:
        logger.warning("获取地址为 {} 的 Selenium 信息失败，错误信息是：".format(drive_path))
        logger.warning(repr(e))
        return None


def get_best_driver(path_list):
    max_available_session_num = 0
    best_driver = random.choice(path_list)
    try:
        for path in path_list:
            session_info = check_session(path)
            if session_info:
                available_session_num = session_info['maxSession'] - session_info['sessionCount']
                if available_session_num > max_available_session_num:
                    max_available_session_num = available_session_num
                    best_driver = path
    except Exception as e:
        logger.warning("获取最佳 Selenium 服务信息失败，将随机选取一个服务，错误信息是：")
        logger.warning(repr(e))
    logger.info("选择的 Selenium 服务地址为 {}".format(best_driver))
    return best_driver


def get_screenshot(driver, url, save_path, save_page_title=True):
    driver.maximize_window()
    driver.get(url)
    width = driver.execute_script("return document.documentElement.scrollWidth")
    height = driver.execute_script("return document.documentElement.scrollHeight")
    driver.set_window_size(width, height)
    if save_page_title:
        page_title = driver.title
        save_file_name = save_path + '/' + get_legal_file_name(page_title) + '.png'
    else:
        save_file_name = save_path + '/' + get_str_md5(url) + '.png'
    driver.save_screenshot(save_file_name)
    return save_file_name


def init_remote_firefox_driver(remote_path):
    driver_options = Options()
    driver_options.headless = True
    driver_options.add_argument('--no-sandbox')
    driver_options.add_argument('--disable-dev-shm-usage')
    driver_options.add_argument('--disable-gpu')
    driver_type = DesiredCapabilities.FIREFOX
    driver = WebDriver(command_executor=remote_path, desired_capabilities=driver_type,
                       options=driver_options)
    return driver


def pic_to_pdf(img_path, pdf_path):
    if "jpg" in img_path or 'png' in img_path or 'jpeg' in img_path:
        pdf_name = img_path.split('/')[-1].split('.')[0]
        image_rgba = Image.open(img_path)
        image_rgb = Image.new('RGB', image_rgba.size, (255, 255, 255))  # white background
        image_rgb.paste(image_rgba, mask=image_rgba.split()[3])  # paste using alpha channel as mask
        image_rgb.save(pdf_path + '/' + pdf_name + '.pdf', "PDF", resolution=100.0, quality=95)


if __name__ == "__main__":
    _driver = init_remote_firefox_driver('http://144.34.185.108:28095')
    _url = "https://www.mcc.gov/"
    _save_path = 'E:/yuxin'
    try:
        screenshot_file_path = get_screenshot(_driver, _url, _save_path, save_page_title=False)
        pic_to_pdf(screenshot_file_path, _save_path)
    except Exception as e:
        print(repr(e))
    _driver.quit()
