import time
from multiprocessing import Pool
from pathlib import Path
from tempfile import TemporaryDirectory
from zipfile import ZipFile

from retry import retry
from selenium import webdriver
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.support.select import By, Select
from selenium.webdriver.support.wait import WebDriverWait

dir_ = Path(__file__).parent

PARTS = [
    *[f"TSW-1{i:02d}-05-G-S" for i in range(1, 10)],
]
URL_STEP = "https://www.snapeda.com/parts/{0}/Samtec/embed/?ref=samtec"


@retry(tries=2, delay=1)
def download_step(part: str):
    print(part)
    with TemporaryDirectory(prefix=str(dir_ / "tmp")) as tmp:
        tmp = Path(tmp)

        chrome_options = webdriver.ChromeOptions()
        chrome_options.add_experimental_option(
            "prefs",
            {"download.default_directory": str(tmp)},
        )
        chrome_options.add_argument("--headless")

        with webdriver.Chrome(
            options=chrome_options,
        ) as browser:
            browser.get(URL_STEP.format(part.upper()))
            wait = WebDriverWait(browser, 10)
            wait.until(
                EC.element_to_be_clickable(
                    browser.find_element(by=By.ID, value="download_traceparts_3d_model")
                )
            ).click()
            wait.until(
                EC.element_to_be_clickable(
                    browser.find_element(
                        by=By.ID,
                        value="samtec-checkbox-3d-modal-download-individual-btn",
                    )
                )
            ).click()

            # wait for download to finish
            while len(list((Path(tmp).glob("*.zip")))) == 0:
                time.sleep(0.1)

        # unzip
        with ZipFile(list((Path(tmp).glob("*.zip")))[0]) as zip:
            zip.extractall(path=str(tmp))

        # move
        filename = f"{part.upper()}.stp"
        (Path(tmp) / filename).rename(dir_ / "3dshapes" / filename)


for part in PARTS:
    download_step(part)