import time
import traceback

from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.common.by import By
from selenium.webdriver.chrome.service import Service
from selenium.common.exceptions import NoSuchElementException, WebDriverException
from webdriver_manager.chrome import ChromeDriverManager
import undetected_chromedriver as uc

print('Sleeping 10s before loading page...')
time.sleep(10)

ser = Service('/usr/lib/chromium-browser/chromedriver')
chrome_options = uc.ChromeOptions()
chrome_options.add_argument('--headless')
chrome_options.add_argument('start-maximized')
chrome_options.add_argument('--no-sandbox')


try:
    driver = uc.Chrome(service=ser, options=chrome_options)
except WebDriverException as e:
    print('Wrong chrome driver version, extracting correct version...')
    tb = traceback.format_exc()
    version_string = tb.split('Current browser version is ')[1]
    major_version = version_string.split('.')[0]
    print('Trying version:', major_version)

    chrome_options = uc.ChromeOptions()
    chrome_options.add_argument('--headless')
    chrome_options.add_argument('start-maximized')
    chrome_options.add_argument('--no-sandbox')
    driver = uc.Chrome(service=ser, options=chrome_options, version_main=int(major_version))
    

try:
    driver.get('https://sensors.dns.t0.vc')
    time.sleep(3)

    driver.execute_script("return document.getElementsByClassName('menu')[0].remove();")

    graphs = [
        'Solar_Power',
        'Living_Room_Air',
        'Outside_Temperature',
        'Bedroom_Temperature',
        'Nook_Temperature',
        #'Misc_Temperature',
        'Basement_Temperature',
        'Nook_Thermostat',
        'Gas_Usage',
        'Water_Usage',
        'Living_Room_Lux',
    ]

    for graph in graphs:
        print('Capturing', graph, 'graph...')

        try:
            element = driver.find_element(By.ID, graph)
        except NoSuchElementException:
            print('Graph not found, skipping.')
            continue

        driver.execute_script('arguments[0].scrollIntoView({block: "center"});', element)

        with open('/home/tanner/sensors/export/data/{}.png'.format(graph), 'wb') as f:
            f.write(element.screenshot_as_png)

finally:
    driver.close()
    driver.quit()

print('done.')