import pickle
import asyncio
import aiofiles
from pathlib import Path


pkg_path = Path("__file__").resolve().parent / "src"
print('pkg_path = {}'.format(pkg_path))
ckip_path = pkg_path / "models" / "ckip"
print('ckip_path = {}'.format(ckip_path))
pkg_path.parent.mkdir(parents=True, exist_ok=True)
ckip_path.parent.mkdir(parents=True, exist_ok=True)
ckip_path.mkdir(parents=True, exist_ok=True)
permission_tester = ckip_path / f"drivers_permission_tester.txt"
with open(permission_tester, 'w') as f_test: pass  # check permission


async def write_drivers(nlp_model: str) -> None:
    """The write drivers function writes the ckip drivers to pickle files asynchronously.
    Args:
        nlp_model (str): the nlp model name
    """
    driver_path = ckip_path / f"{nlp_model}_drivers.pickle"
    with open(driver_path, 'w') as f_test: pass  # check permission
    from ckip_transformers.nlp import (
        CkipWordSegmenter,
        CkipPosTagger,
        CkipNerChunker,
    )

    drivers = (
        CkipWordSegmenter(model=nlp_model),
        CkipPosTagger(model=nlp_model),
        CkipNerChunker(model=nlp_model),
    )

    driver_path = ckip_path / f"{nlp_model}_drivers.pickle"
    async with aiofiles.open(driver_path, mode="wb") as file:
        result = pickle.dumps(drivers)
        await file.write(result)
        print(f"{nlp_model}_drivers.pickle done!")


async def download_ckip_drivers(ckip_nlp_models):
    await asyncio.gather(*list(map(write_drivers, ckip_nlp_models)))