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

import numpy as np
import regex as re
from selenium import webdriver
from selenium.webdriver.support.select import By

from fp_generator import INCH, Footprint

dir_ = Path(__file__).parent

PARTS = [
    # JST-SH, Right angle
    *[f"SM{i:02d}B-SRSS-TB" for i in range(2, 16)],
    # JST-SH, Vertical
    # *[f"BM{i:02d}B-SRSS-TB" for i in range(2, 16)],
]


def round2grid(val: float, direction: Optional[Literal["up", "down"]] = None, grid: float = 0.005 * INCH):
    if direction == "up":
        val = val + grid / 2
    elif direction == "down":
        val = val - grid / 2

    return np.round(val / grid) * grid


def download_step(part: Union[str, int], 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.jst-mfg.com/product/index.php?type=10&series=231&doc=2&filename={part.upper()}.zip"
            )
            browser.implicitly_wait(10)
            browser.find_element(by=By.ID, value="company").send_keys("N/A")
            browser.find_element(by=By.ID, value="section").send_keys("N/A")
            browser.find_element(by=By.ID, value="name").send_keys("Brendan Haines")
            browser.find_element(by=By.ID, value="address").send_keys("7415 Spy Glass Ct")
            browser.find_element(by=By.ID, value="tel").send_keys("(505) 738-6174")
            browser.find_element(by=By.ID, value="email").send_keys("brendan.haines@gmail.com")
            # browser.find_element(by=By.ID, value="url").send_keys("")
            browser.find_element(by=By.XPATH, value="//input[@type='button' and @value='I agree']").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(tmp / f"{part.upper()}.zip") as zip:
            zip.extractall(path=tmp)

        # move
        (Path(tmp) / f"{part.upper()}.STEP").rename(dir_ / "jst.pretty" / "step" / f"{part.upper()}.stp")


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

    if part.startswith("SM"):
        # Right angle
        match = re.match(r"SM(\d\d)B-SRSS-TB", part)
        pins = int(match[1])

        fp = Footprint(part, lib=lib, description=f"connector, JST-SH, 1x{pins:02f}, RA")
        fp.add_text("reference", "REF**", (5.55 + (0.025 + 0.005) * INCH), 0, rotation=90, layer="F.SilkS")
        fp.add_text("value", "VAL**", 4 * 0.025 * INCH, 0, rotation=90, layer="F.Fab")
        fp.add_text("user", r"${REFERENCE}**", 2 * 0.025 * INCH, 0, rotation=90, layer="F.Fab")
        a, b = {
            2: (1, 4),
            3: (2, 5),
            4: (3, 6),
            5: (4, 7),
            6: (5, 8),
            7: (6, 9),
            8: (7, 10),
            9: (8, 11),
            10: (9, 12),
            11: (10, 13),
            12: (11, 14),
            13: (12, 15),
            14: (13, 16),
            15: (14, 17),
            20: (19, 22),
        }[pins]
        # fp.add_line((0.05 * INCH, -0.05 * INCH), (0.05 * INCH, 0.05 * INCH), layer="F.SilkS")
        for pad in range(pins):
            fp.add_pad(
                pad + 1,
                4 + (5.55 - 4) / 2,
                -1.0 * (pins - 1) / 2 + 1.0 * pad,
                size=(5.55 - 4, 0.6),
                shape="rect",
            )
        # mechanical holes
        for s in [1, -1]:
            fp.add_pad(0, 1.8 / 2, s * (1.0 * (pins - 1) / 2 + 0.7 + 1.2 / 2), size=(1.8, 1.2), shape="rect")
        fp.add_model(
            lib / "step" / f"{part}.stp",
            (4, 0, -0.3),
            rotate=(-90, 0, 90),
        )
        fp.add_rect((0.2, -b / 2), (0.2 + 4.25, b / 2), layer="F.Fab")
        for s in [1, -1]:
            fp.add_line(
                (0.190 * INCH, s * round2grid((a + 0.6) / 2 + 0.010 * INCH, "up")),
                (0.190 * INCH, s * round2grid(b / 2 + 0.010 * INCH, "up")),
                layer="F.SilkS",
            )
            fp.add_line(
                (0.190 * INCH, s * round2grid(b / 2 + 0.010 * INCH, "up")),
                (round2grid(1.8 + 0.010 * INCH), s * round2grid(b / 2 + 0.010 * INCH, "up")),
                layer="F.SilkS",
            )
        fp.write()
    else:
        raise NotImplementedError


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

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