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

def download(beatmap_id, download_directory, session):
    beatmap_id = str(beatmap_id)
    os.makedirs(f"/tmp/osudd/{beatmap_id}/", exist_ok=True)
    options = webdriver.ChromeOptions()
    options.add_argument("--headless")
    options.add_argument("--no-sandbox")
    options.add_argument("--disable-dev-shm-usage")
    options.add_experimental_option('excludeSwitches', ['enable-logging'])
    prefs = {"profile.default_content_settings.popups": 0,    
            "download.default_directory":f"/tmp/osudd/{beatmap_id}/", ### Set the path accordingly
            "download.prompt_for_download": False, ## change the downpath accordingly
            "download.directory_upgrade": True}
    options.add_experimental_option("prefs", prefs)
    driver = webdriver.Chrome(options=options)
    driver.get("https://osu.ppy.sh/")
    new_cookie = {'name': 'osu_session',
                'value': session,
                'domain': 'osu.ppy.sh'
                } 
    driver.add_cookie(new_cookie)
    driver.get(f"https://osu.ppy.sh/beatmapsets/{beatmap_id}")
    time.sleep(2)
    element = driver.find_element(By.XPATH,"//*[contains(text(), 'Download')]")
    parent_a_tag = element.find_element(By.XPATH,'ancestor::a[1]')
    parent_a_tag.click()
    timeout = 60
    elapsed_time = 0
    while not any(file.endswith(".osz") for file in os.listdir(f"/tmp/osudd/{beatmap_id}/")):
        time.sleep(1)
        elapsed_time += 1
        if elapsed_time > timeout:
            break
    shutil.move(f"/tmp/osudd/{beatmap_id}/", f"{download_directory}/{beatmap_id}/")
    driver.quit()