import json
import time
import argparse
import os

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

# To run the script locally ,you need:
# 1. download: chrome
# 2. download: chromedriver, match the version of chrome.
#     https://sites.google.com/chromium.org/driver/downloads?authuser=0
# 3. install python3 and selenium
# 4. set chromdriver to PATH
# 5. run chromdriver to make sure you have permission to run it.
# 6. set bear_token_credentials ENV, e.g. export bear_token_credentials=<user>::<pwd>
# 7. run script: python get_auth_token.py --chromedriver /xxx/xxx --url https://staging.tidbcloud.com/ --tokendir /tmp


def get_driver():
    # run it in background
    chrome_options = Options()
    chrome_options.add_argument('--headless')
    chrome_options.add_argument('--disable-gpu')
    chrome_options.add_argument('--no-sandbox')
    chrome_options.add_argument('--disable-dev-shm-usage')
    chrome_options.set_capability('goog:loggingPrefs', {'performance': 'ALL'})

    return webdriver.Chrome(options=chrome_options)


# convert message to a json object
def process_browser_log_entry(entry):
    return json.loads(entry['message'])['message']


def get_token(events):
    token_prefix = 'Bearer '
    print('Whether it can retrieve security token:', token_prefix in str(events))
    for event in events:
        try:
            headers = event.get("params").get('headers')
            url = headers.get(":path")
            method = headers.get(":method")
            if url.endswith("/me") and method == 'GET':
                auth = headers.get('authorization')
                return auth
        except Exception as err:
            pass


def write_token(user, tokendir, token):
    if not os.path.exists(tokendir):
        os.makedirs(tokendir)

    token_file = os.path.join(tokendir, user)
    if not token:
        return
    print("token is: \n%s" % token)
    with open(token_file, 'w') as f:
        f.write(token)


def get_events(driver):
    time.sleep(10)
    browser_log = driver.get_log('performance')
    return [process_browser_log_entry(entry) for entry in browser_log]


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--url', type=str)
    parser.add_argument('--credential', type=str)
    parser.add_argument('--tokendir', type=str, default='.')
    return parser.parse_args()


def login_tidb_cloud(driver, url, username, password):
    print('opening url >>>> ', url)
    driver.get(url)

    wait = WebDriverWait(driver, 10, 0.5)
    wait.until(EC.presence_of_element_located((By.CLASS_NAME, "email")), message="failed to get elements.")

    driver.find_element(By.CSS_SELECTOR, ".form-input.email").send_keys(username)
    driver.find_element(By.CSS_SELECTOR, ".form-input.password").send_keys(password)
    driver.find_element(By.CLASS_NAME, "btn-submit").click()

    # wait for button 'No' , if is existed, the password is expired.
    try:
        wait = WebDriverWait(driver, 10, 0.5)
        wait.until(EC.presence_of_element_located((By.XPATH, "//*[text()='No']")),message="Password is within validity period")
    except Exception as err:
        print("Exception is : ", err)
        pass
    else:
        # click it and skip changing password.
        driver.find_element(By.XPATH, "//*[text()='No']").click()


def get_users(bear_token_credentials):
    users = {}
    if not bear_token_credentials or len(bear_token_credentials) == 0:
        return users

    sep = '::'
    for line in bear_token_credentials.split():
        if len(line.split(sep)) != 2:
            print('error line:' + line)
            continue
        users[line.split(sep)[0]] = line.split(sep)[1]

    return users


def get_user_token(user, pwd):
    retry_times = 3
    while retry_times > 0:
        try:
            driver = get_driver()
            login_tidb_cloud(driver, get_args().url, user, pwd)
            token = get_token(get_events(driver))
            driver.quit()
            if token and len(token) > 0:
                return token
            raise Exception("token is empty.")
        except Exception as err:
            print(err)
            retry_times = retry_times - 1
            print("Can't get user token. retrying...")
            time.sleep(10)
    else:
        print("failed to get user token. user: " + user)


def main():
    credentials = get_args().credential

    users = get_users(credentials)
    if not users and len(users) == 0:
        print('not credentials found.')
        return

    for user in users:
        print("current user:", user)
        token = get_user_token(user, users[user])
        write_token(user, get_args().tokendir, token)


if __name__ == '__main__':
    main()