import time
import zipfile
from pathlib import Path
from tempfile import TemporaryDirectory
from typing import Optional, Union

import numpy as np
import regex as re
from selenium import webdriver
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.support.select import By
from selenium.webdriver.support.wait import WebDriverWait

from fp_generator import INCH, Footprint

dir_ = Path(__file__).parent
rng = np.random.default_rng()


PARTS = [
    # Pin Headers
    *[f"TSW-1{i:02d}-07-L-S" for i in range(1, 5)],  # quick for testing
    # *[f"TSW-1{i:02d}-07-L-S" for i in range(1, 51)],
    # *[f"TSW-1{i:02d}-07-L-D" for i in range(1, 51)],
    # # *[f"TSW-1{i:02d}-07-L-T" for i in range(1, 51)],
    # # *[f"TSW-1{i:02d}-07-L-Q" for i in range(1, 51)],
    # *[f"TSW-2{i:02d}-07-L-S" for i in range(2, 26)],
    # *[f"TSW-2{i:02d}-07-L-D" for i in range(2, 26)],
    # # *[f"TSW-2{i:02d}-07-L-T" for i in range(2, 26)],
    # # *[f"TSW-2{i:02d}-07-L-Q" for i in range(2, 26)],
    # Sockets
    # SSW
    # SMH
    # Discrete Wire
    # TFM
    # SFM
    # High Speed Board to Board
    # *np.flatten(
    #     [[f"LSHM-1{i:02d}-{h:02.1f}-L-DV-A-S-TR" for i in [5, 10, 20, 30, 40, 50]] for h in [2.5, 3.0, 4.0, 6.0]]
    # ),
]


def download_step(part: str, headless: bool = True) -> None:
    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)},
        )
        if headless:
            chrome_options.add_argument("--headless")

        with webdriver.Chrome(
            options=chrome_options,
        ) as browser:
            browser.get(f"https://www.snapeda.com/parts/{part.upper()}/Samtec/embed/?ref=samtec")
            wait = WebDriverWait(browser, 30)
            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)
            # print(list((Path(tmp).glob("*.zip"))))

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

        # move
        filename = f"{part.upper()}.stp"
        # if filename not in list(Path(tmp).glob("*.stp")):
        #     raise ValueError(
        #         f"Name mismatch: {filename} does not exist. Found {[f.name for f in list(Path(tmp).glob('*.stp'))]}"
        #     )
        (Path(tmp) / filename).rename(dir_ / "samtec.pretty" / "step" / filename)


def footprint(part: str, lib: Optional[Union[Path, str]] = None) -> str:
    if lib is None:
        lib = dir_ / "samtec.pretty"
    lib = Path(lib)

    if any([part.startswith(s) for s in ["TSW-", "HTSW-"]]):
        # pinheader
        match = re.match(
            r"(H?)TSW-([12])(\d\d)-(\d\d)-([FLGT])-([SDTQ])(-R[AE])?(-NA)?(-LL)?(-LC)?(-LA)?(-(\d\d\d))?",
            part.upper(),
        )
        high_temp = match[1] != ""
        spacing = int(match[2])
        rows = int(match[3])
        # lead_style = {
        #     5:
        # }[int(match[4])]
        lead_style = "vertical"
        plating = {
            "F": "Gold flash on post, Matte Tin on tail",
            "L": "10 µin (0.25 µm) Gold on post, Matte Tin on tail",
            "G": "10 µin (0.25 µm) Gold on post, Gold flash on balance",
            "T": "Matte Tin",
        }[match[5]]
        cols = {"S": 1, "D": 2, "T": 3, "Q": 4}[match[6]]

        strobe = [(2 if cols == 4 else 1) * 0.1 * INCH, spacing * 0.1 * INCH]
        if cols == 4:
            # only outer columns filled
            cols = 2
        fp_name = f"PinHeader_{cols}x{rows:02d}_0.1in_{part}"
        description = (
            f"Header, Male, {cols}x{rows:02d}, {lead_style}, {plating}{', high temperature' if high_temp else ''}"
        )
        keywords = ["header", "pin"]
        hole_diam = 0.040 * INCH
        pad_size = 0.075 * INCH

        center = [strobe[0] * (cols - 1) / 2, strobe[1] * (rows - 1) / 2]

        fp = Footprint(fp_name, lib=lib, description=description, keywords=keywords)
        fp.add_text("reference", "REF**", center[0], -2.032, layer="F.SilkS")
        fp.add_text("value", "VAL**", center[0] - 0.025 * INCH, center[1], rotation=90, layer="F.Fab")
        fp.add_text("user", r"${REFERENCE}**", center[0] + 0.025 * INCH, center[1], rotation=90, layer="F.Fab")
        fp.add_rect(
            [-0.05 * INCH] * 2,
            (0.05 * INCH + strobe[0] * (cols - 1), 0.05 * INCH + strobe[1] * (rows - 1)),
            layer="F.SilkS",
        )
        fp.add_line((-0.05 * INCH, 0.05 * INCH), (0.05 * INCH, 0.05 * INCH), layer="F.SilkS")
        fp.add_line((0.05 * INCH, -0.05 * INCH), (0.05 * INCH, 0.05 * INCH), layer="F.SilkS")
        for pad in range(rows * cols):
            fp.add_pad(
                pad + 1,
                strobe[0] * (pad % cols),
                strobe[1] * int(pad / cols),
                size=pad_size,
                hole=hole_diam,
                shape="rect" if pad == 0 else "circle",
            )
            fp.add_model(
                lib / "step" / f"{part}.stp",
                (center[0], center[1], 0),
                rotate=(-90, 0, 90),
            )
        fp.write()

    elif any([part.startswith(s) for s in ["SSW-"]]):
        # socket
        raise NotImplementedError()
    else:
        raise ValueError(f"Unknown part family for {part}")


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

    pool = Pool(10)

    # for part in PARTS:
    #     # time.sleep(rng.random() * 3)
    #     # download_step(part)
    #     footprint(part)

    # pool.map(download_step, PARTS)
    pool.map(footprint, PARTS)
    pool.close()