from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from PIL import Image
import io
import time
import os
import sys
import numpy as np
from fpdf import FPDF

WEBSITES = """https://google.com
https://yahoo.com"""
WEBSITES = WEBSITES.splitlines()
# screenshot size
WIDTH = 800
HEIGHT = 600
# screenshot margin
MARGIN = 10
# screenshot padding
PADDING = 10
options = Options()
options.add_argument('--headless')
options.add_argument('--disable-gpu')
options.add_argument('--no-sandbox')
driver = webdriver.Chrome(options=options, executable_path=os.getcwd() + '/chromedriver')
driver.set_window_size(WIDTH, HEIGHT)
for website in WEBSITES:
    try:
        driver.get(website)
        time.sleep(1)
        # zoom out the image before taking screenshot
        driver.execute_script("document.body.style.zoom='50%'")
        # get the screenshot
        time.sleep(0.2)
        img = Image.open(io.BytesIO(driver.get_screenshot_as_png()))
        img = img.resize((WIDTH, HEIGHT))
    except:
        print('error', website)
    img.save('/tmp/screenshots/' + website.replace('https://', '').replace('http://', '').replace('/', '_') + '.png')
driver.quit()

pdf = FPDF()
images = os.listdir('/tmp/screenshots')
images.sort()
for image in images:
    if not image.endswith('.png'):
        continue
    pdf.add_page()
    pdf.set_font('Arial', 'B', 16)
    pdf.cell(40, 10, image.replace('.png', ''))
    pdf.image('/tmp/screenshots/' + image, x=0, y=20, w=210, h=0, type='png')
pdf.output('/tmp/screenshots.pdf', 'F')