from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC

from bs4 import BeautifulSoup
from chromedriver_py import binary_path

import tarfile 
import requests
import time
import os.path
import os
import shutil
from timeout import timeout


def download_data(mode, version, n_max_tries, max_time_per_dl, start, stop, dl_root):

    # paths and variables
    variables = ["co", "ch4"]
    var = ["CO", "CH4"] # name of directories

    login_url = 'https://gaw.kishou.go.jp/login/user'
    url = 'https://gaw.kishou.go.jp/search/gas_species/{}/latest'
    
    if mode == "all":
        os.makedirs('{}/WMO_WDCGG/original_files/{}/'.format(dl_root, version), exist_ok=True)
        download_location = "{}/WMO_WDCGG/original_files/{}/".format(dl_root, version)

    elif mode == "nrt":
        os.makedirs('{}/WMO_WDCGG/original_files/nrt'.format(dl_root), exist_ok=True)
        download_location = "{}/WMO_WDCGG/original_files/nrt/".format(dl_root)

    # set up driver
    options = Options()
    prefs = {'download.default_directory' : download_location}
    options.add_experimental_option('prefs', prefs)
    options.add_argument("--no-sandbox")
    options.add_argument("--headless")

    i=0
    for variable in variables:

        svc = webdriver.ChromeService(executable_path=binary_path)
        driver = webdriver.Chrome(service=svc, options=options)

        # open url
        driver.get(login_url)
        time.sleep(2)

        # login
        email = driver.find_element(By.ID, "uid") 
        email.send_keys("raphael.grodofzig@bsc.es")
        passwd = driver.find_element(By.ID, "pwd") 
        passwd.send_keys("IbRGaH66")
        time.sleep(1)
        driver.find_element(By.ID, "loginbt").click()
        print("Successfully logged in")
        time.sleep(2)

        driver.get(url.format(variable))

        WebDriverWait(driver, max_time_per_dl).until(EC.presence_of_all_elements_located((By.CSS_SELECTOR, "button[class='btn'][type='submit']"))) # wait till loaded

        # start download
        driver.find_element(By.CSS_SELECTOR, "button[class='btn'][type='submit']").click()
        print(variable)

        # wait until download finished
        while not os.path.exists("{}{}_latest.txt.tar.gz".format(download_location, variable)):
            time.sleep(1)

            if os.path.isfile("{}{}_latest.txt.tar.gz".format(download_location, variable)):
                print('{} download successful'.format(variable))
                time.sleep(2)

                # untar
                with tarfile.open("{}{}_latest.txt.tar.gz".format(download_location, variable)) as file:
                    file.extractall(download_location)
                time.sleep(2)

                # move files
                shutil.rmtree(download_location+var[i]) # remove old data
                shutil.move(download_location+"txt/"+variable, download_location+var[i])
                shutil.rmtree(download_location+"txt/")
                
        
        os.remove("{}{}_latest.txt.tar.gz".format(download_location, variable))
        time.sleep(2)
        i=i+1      

        driver.close()
        time.sleep(4)