from selenium.webdriver.remote.webdriver import WebDriver
from selenium.webdriver.common.by import By
from PIL import Image
import time
import os
_scroll_to_bottom = "window.scroll({top:document.body.clientHeight,left:0,behavior:'auto'});"
_scroll_to_y = "window.scroll({top:%d,left:0,behavior:'auto'});"
_tmp_dir = "./tmp"
_tmp_file = "%s/%%s_%%s.png" % _tmp_dir
_body = "//body"
def screenshot(driver:WebDriver, save_path):
    file_name = os.path.basename(save_path)
    if not os.path.exists(_tmp_dir):
        os.mkdir(_tmp_dir)

    tmp_file_list = []
    try:
        tmp_file = _tmp_file % (file_name, 0)
        driver.save_screenshot(tmp_file)
        tmp_file_list.append(tmp_file)
        body_h = driver.execute_script("return document.documentElement.scrollHeight")
        current_h = driver.execute_script("return document.documentElement.clientHeight")
        for i in range(1, int(body_h / current_h)):
            driver.execute_script(_scroll_to_y % (current_h * i))
            time.sleep(0.5)
            driver.save_screenshot(_tmp_file % (file_name, i))
            tmp_file_list.append(_tmp_file % (file_name, i))
            merge_images(tmp_file, _tmp_file % (file_name, i), tmp_file)

        driver.execute_script(_scroll_to_bottom)
        end_file = _tmp_file % (file_name, "bottom")
        driver.save_screenshot(end_file)
        tmp_file_list.append(end_file)
        merge_images(tmp_file, end_file,save_path,int(current_h - int(body_h % current_h)))
    finally:
        # 删除临时文件
        for tmp in tmp_file_list:
            os.remove(tmp)


def merge_images(image1: str, image2: str,output: str,overlap_size=0):
    size = overlap_size
    img1,img2 = Image.open(image1), Image.open(image2)
    size1, size2 = img1.size, img2.size
    merge = Image.new("RGB", (size1[0], size1[1] + size2[1] - size))
    loc1, loc2 = (0, 0), (0, size1[1] - size)
    merge.paste(img1, loc1)
    merge.paste(img2, loc2)
    merge.save(output)