import zipfile
from pathlib import Path
from tempfile import TemporaryDirectory
from typing import Union

import requests

from fp_generator import INCH, Footprint

dir_ = Path(__file__).parent

PARTS = [
    # # Micro-Fit 3.0 Single Row Headers
    # *[f"43650-{i:02d}10" for i in range(2, 13)],  # 1 row, RA, press-fit retention clip
    # *[f"43650-{i:02d}13" for i in range(2, 13)],  # 1 row, RA, solder tab
    # *[f"43650-{i:02d}22" for i in range(2, 13)],  # 1 row, Vertical, press-fit retention clip
    # *[f"43650-{i:02d}25" for i in range(2, 13)],  # 1 row, RA, solder tab
    # # Micro-Fit 3.0 Dual Row Headers
    # *[f"43045-{i:02d}07" for i in range(2, 25, 2)],  # 2 row, RA, press-fit retention clip
    # *[f"43045-{i:02d}10" for i in range(2, 25, 2)],  # 2 row, RA, solder tab
    # *[f"43045-{i:02d}16" for i in range(2, 25, 2)],  # 2 row, Vertical, press-fit retention clip
    # *[f"43045-{i:02d}19" for i in range(2, 25, 2)],  # 2 row, RA, solder tab
    # # Eurostyle 5.08mm Headers
    "39531-0002",
    # Eurostyle 5.08mm Plugs
    # "39530-0002",
]


def download_step_molex(part: Union[str, int]) -> None:
    # remove dash
    part = f'{int(str(part).replace("-", "")):09d}'

    print(part)

    with TemporaryDirectory(prefix=str(dir_ / "tmp") + "/") as tmp:
        tmp = Path(tmp)

        for p in [part, part[:5] + "-" + part[5:]]:
            try:
                zipname = f"{part}_stp.zip"

                # download
                with open(tmp / zipname, "wb") as f:
                    f.write(requests.get(f"https://www.molex.com/pdm_docs/stp/{p}_stp.zip").content)

                # unzip
                with zipfile.ZipFile(tmp / zipname) as zip:
                    zip.extractall(path=str(tmp))

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

                return
            except zipfile.BadZipFile:
                print(f"zip does not exist for part {p}")

        raise ValueError(f"Could not download .stp for part {part}")


if __name__ == "__main__":
    from multiprocessing import Pool

    pool = Pool(10)
    pool.map(download_step_molex, PARTS)
    # for part in PARTS:
    #     download_step_molex(part)
    pool.close()