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

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

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

PARTS = [
    *[f"TSW-1{i:02d}-07-L-S" for i in range(1, 50)],
    *[f"TSW-1{i:02d}-07-L-D" for i in range(1, 50)],
    *[f"TSW-1{i:02d}-07-L-T" for i in range(1, 50)],
    *[f"TSW-1{i:02d}-07-L-Q" for i in range(1, 50)],
    *[f"TSW-2{i:02d}-07-L-S" for i in range(2, 25)],
    *[f"TSW-2{i:02d}-07-L-D" for i in range(2, 25)],
    *[f"TSW-2{i:02d}-07-L-T" for i in range(2, 25)],
    *[f"TSW-2{i:02d}-07-L-Q" for i in range(2, 25)],
]
URL_STEP = "https://www.snapeda.com/parts/{0}/Samtec/embed/?ref=samtec"


def download_step_samtec(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, 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(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" / "3dshapes" / filename)


now = datetime.now()
INCH = 25.4
SILK_WIDTH = 0.007 * INCH


def footprint_samtec(part: str, lib: Optional[Union[Path, str]] = None) -> str:
    if part[:4].upper() == "TSW-":
        # pinheader
        match = re.match(
            r"(H?)TSW-([12])(\d\d)-(\d\d)-([FLGT])-([SDTQ])(-RA)?(-RE)?(-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 µ" (0.25 µm) Gold on post, Matte Tin on tail',
            "G": '10 µ" (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]

        if lib is None:
            lib = dir_ / "samtec.pretty"
        lib = Path(lib)
        with open(lib / f"{fp_name}.kicad_mod", "w") as f:
            f.write(
                "\n".join(
                    (
                        f'(footprint "{fp_name}" (version {now.strftime("%Y%m%d")}) (generator pcbnew) (layer "F.Cu")',
                        r"  (attr through_hole)",
                        f'  (descr "{description}")',
                        f'  (tags "{" ".join(keywords)}")',
                        f'  (fp_text reference "REF**" (at {center[0]} {-2.032}) (layer "F.SilkS")',
                        r"    (effects (font (size 0.762 0.762) (thickness 0.127)))",
                        r"  )",
                        f'  (fp_text value "VAL**" (at {center[0]} {center[1]} 90) (layer "F.Fab") hide',
                        r"    (effects (font (size 0.762 0.762) (thickness 0.127)))",
                        r"  )",
                        r'  (fp_text user "${REFERENCE}" ' f'(at {center[0]} {center[1]} 90) (layer "F.Fab")',
                        r"    (effects (font (size 1 1) (thickness 0.15)))",
                        r"  )",
                        f"  (fp_rect (start {-0.05 * INCH} {-0.05 * INCH}) "
                        f"(end {0.05 * INCH + strobe[0] * (cols - 1)} {0.05 * INCH + strobe[1] * (rows - 1)})",
                        f'    (stroke (width {SILK_WIDTH}) (type default)) (fill none) (layer "F.SilkS")' r"    )",
                        f"  (fp_line (start {-0.05 * INCH} {0.05 * INCH}) (end {0.05 * INCH} {0.05 * INCH}) "
                        f'(layer "F.SilkS") (width {SILK_WIDTH}))',
                        f"  (fp_line (start {0.05 * INCH} {-0.05 * INCH}) (end {0.05 * INCH} {0.05 * INCH}) "
                        f'(layer "F.SilkS") (width {SILK_WIDTH}))',
                        *[
                            f'  (pad "{pad + 1}" thru_hole {"rect" if pad == 0 else "circle"} '
                            f"(at {strobe[0] * (pad % cols)} {strobe[1] * int(pad / cols)}) "
                            f"(size {pad_size} {pad_size}) (drill {hole_diam}) (layers *.Cu *.Mask))"
                            for pad in range(rows * cols)
                        ],
                        f'  (model "/home/brendan/Documents/projects/kicad/samtec.pretty/3dshapes/{part}.stp"',
                        f"    (offset (xyz {center[0]} {-center[1]} 0))",
                        r"    (scale (xyz 1 1 1))",
                        r"    (rotate (xyz -90 0 90))",
                        r"  )",
                        r")",
                    )
                )
                + "\n"
            )
    else:
        raise ValueError(f"Unknown part family for {part}")


if __name__ == "__main__":
    for part in PARTS:
        time.sleep(rng.random() * 3)
        download_step_samtec(part)
        footprint_samtec(part)

    # from multiprocessing import Pool
    # pool = Pool(10)
    # pool.map(download_step_samtec, PARTS)
    # pool.map(footprint_samtec, PARTS)
    # pool.close()