import argparse
import random
import time
from pathlib import Path

from huggingface_hub import HfApi
from selenium.webdriver.common.by import By
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.support.ui import WebDriverWait

from lib.hf import list_spaces, login
from lib.webdriver import create_driver, load_cookies, save_cookies

parser = argparse.ArgumentParser()
parser.add_argument("--cookies", required=True, help="Path to the file containing your HF cookies")
parser.add_argument("--token", required=True, help="Your HF API token")
parser.add_argument("--username", required=True, help="Your HF username")
parser.add_argument("--password", required=True, help="Your HF password")
parser.add_argument("--headless", action='store_true', help="Run Selinium in headless mode")
args = parser.parse_args()

cookies_file = Path(args.cookies).resolve().expanduser().absolute()

login(args.token)

driver = create_driver(args.headless)

driver.get(f'https://huggingface.co')
if cookies_file.exists():
    load_cookies(cookies_file)

driver.get('https://huggingface.co/settings/profile')
if driver.current_url != 'https://huggingface.co/settings/profile':
    print('Not logged in!')
    username = WebDriverWait(driver, 10).until(EC.presence_of_element_located((By.NAME, "username")))
    username.send_keys(args.username)
    password = driver.find_element(By.NAME, "password")
    password.send_keys(args.password)
    login_button = WebDriverWait(driver, 10).until(EC.presence_of_element_located((By.XPATH, '//button[@type="submit"]')))
    login_button.click()
    save_cookies(driver, cookies_file)
    print('Manually logged in.')
else:
    print('Auth is correct')

while True:
    our_spaces = [x.id for x in list_spaces(args.username) if 'proxy' in x.id]
    random.shuffle(our_spaces)

    driver.get(f'https://huggingface.co')
    time.sleep(5)
    for space in our_spaces:
        print(space)
        driver.get(f'https://huggingface.co/spaces/{space}')
        time.sleep(100)

        # try:
        if len(driver.find_elements(By.ID, 'captcha-container')):
            print('CAPTCHA!!!')
        else:
            for iframe in driver.find_elements(By.TAG_NAME, 'iframe'):
                src = iframe.get_attribute('src')
                if 'hf.space' in src and '?__sign' in src:
                    print('Space is online')
                    break
            else:
                api = HfApi()
                try:
                    api.restart_space(repo_id=space)
                    print('Restarted space')
                except Exception as e:
                    print('Failed to restart space!', e)

        time.sleep(30)

        # Avoid the captcha
        driver.get('https://huggingface.co')
        time.sleep(30)
    sleep_time = 300
    print(f'Sleeping {sleep_time} seconds...')
    time.sleep(sleep_time)