#!python
####
from selenium import webdriver
from selenium.webdriver.common.by import By
import requests
from time import sleep
from os import system, path, makedirs, chdir
from sarge import get_stdout
import sys
import signal

allthethings = True
baseurl = "https://www.tiktok.com/"
watermark = False
headless = True
headers = {}

def signal_handler(signum, frame):
  try:
    driver.quit()
  except Exception:
    pass

  sys.exit(1)

signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler)

def show_help(stat):
    print("""
Usage: 
       tiktokget <args> username
       or
       tiktokget <args> url
       some args:
       [-l url] to download a url, one-off.
       [-f] to disable headless, to help with download failures 
       or
       [-w] to download with a watermark

       url should be in quotes.

       examples:
       download all of a usernames videos (the default):
       tiktokget someuser 
       download all of a usernames videos, but open the browser window:
       tiktokget -f someuser
""")
    exit(stat)
i=1
if 1 < len(sys.argv) <= 5:
    while i < len(sys.argv):
        if sys.argv[i] == '-w':
            watermark=True
        if sys.argv[i] == '-h' or sys.argv[i] == '--help':
            show_help(0)
        if sys.argv[i] == '-f':
            headless=False
        if sys.argv[i] == '-l':
            allthethings = False
            try:
                if 'https://www.tiktok.com' in sys.argv[i+1]:
                    url=sys.argv[i+1]
                    i=i+1
                else:
                    show_help(1)
            except IndexError:
                show_help(1)
        else:
            username=sys.argv[i]
        i=i+1
else:
    show_help(1)

if allthethings:
    if '@' in username:
        pass
    else:
        username = '@' + username

opts = webdriver.firefox.options.Options()
if not headless:
    opts.headless = False
else:
    opts.headless = True
binary = webdriver.firefox.firefox_binary.FirefoxBinary(get_stdout('which firefox'))
driver = webdriver.Firefox(options=opts)
driver.set_window_size(1080, 3096)
uagent = driver.execute_script("return navigator.userAgent;")
headers['User-Agent'] = uagent

if allthethings:
    url="{}{}".format(baseurl, username)
    n = 0
    urls=[]
    driver.get(url)
    driver.maximize_window()
    keeplooking = True
    vidurls = driver.find_elements(By.XPATH,"//div[contains(@class, 'DivWrapper')]/a")
    
    while keeplooking == True:
        prev_vidurls = len(vidurls)
        for vidurl in vidurls:
            h = vidurl.get_attribute('href')
            if h in urls:
                pass
            else:
                urls.append(h)
    
        driver.execute_script("window.scrollBy(0,document.body.scrollHeight);")
        sleep(3)
        vidurls = driver.find_elements(By.XPATH,"//div[contains(@class, 'DivWrapper')]/a")
        for vidurl in vidurls:
            h = vidurl.get_attribute('href')
            if h in urls:
                pass
            else:
                urls.append(h)
        print("Found " + str(len(vidurls)) + " videos.", end="\r")
        if prev_vidurls == len(vidurls):
            keeplooking = False
            n = len(vidurls)
    
    print("Downloading...")
    
    dirname = username[1:] #remove the leading '@'
    makedirs(dirname, exist_ok=True)
    chdir(dirname)
    
    for url in urls:
        driver.get(url)
        sleep(5)
        title = driver.find_element(By.XPATH,"//span[contains(@class, 'SpanText')]").text
        title = title.replace('/', '\u2044') #replace forward slash with division symbol
        p = '{0:03d}'.format(n)
        title = p + '_' + title + '.mp4'
        print(title)
        title_len = len(title)
        if path.exists(title):
            pass
        else:
            if watermark:
                vid = driver.find_element(By.XPATH,"//video")
                url = vid.get_attribute("src")
                video = requests.get(url, headers=headers)
                with open(title, 'wb') as v:
                    v.write(video.content)
            else:
                system('yt-dlp -q --no-warnings -o "{}" "{}"'.format(title, url))
        n = n - 1
else:
    driver.get(url)
    sleep(5)
    title = driver.find_element(By.XPATH,"//span[contains(@class, 'SpanText')]").text
    title = title.replace('/', '\u2044') #replace forward slash with division symbol
    title = title + '.mp4'
    print(title)
    title_len = len(title)
    if path.exists(title):
        pass
    else:
        print("Downloading...")
        if watermark:
            vid = driver.find_element(By.XPATH,"//video")
            url = vid.get_attribute("src")
            video = requests.get(url, headers=headers)
            with open(title, 'wb') as v:
                v.write(video.content)
        else:
            system('yt-dlp -q --no-warnings -o "{}" "{}"'.format(title, url))

driver.quit()
