import asyncio
import hashlib
from concurrent.futures import ProcessPoolExecutor
from functools import partial
from pathlib import Path
from typing import Dict
from urllib.parse import urlparse

import aiofiles
import aiohttp
import yaml
from conda_package_handling import api

cheap = yaml.load(Path("./conda-lock.yml").open(), Loader=yaml.SafeLoader)
root = Path("~/miniconda3/pkgs").expanduser()
root.mkdir(exist_ok=True)


def extract(sha256: str, file: str, target: str):
    hasher = hashlib.new("sha256")
    with open(file, "rb") as fh:
        for chunk in iter(partial(fh.read, 8192), b""):
            hasher.update(chunk)
    assert hasher.hexdigest() == sha256
    print("Extracting", file, target)
    api.extract(file, target)


async def main(packages: Dict):
    async with aiohttp.ClientSession() as session:
        tasks = []
        urls = []
        for package in packages:
            if package["platform"] == "linux-aarch64" and package["manager"] == "conda":
                url = package["url"]
                parsed = urlparse(url)
                pth = Path(parsed.path)
                pth.name
                tasks.append(
                    download_url(
                        sha256=package["hash"]["sha256"],
                        session=session,
                        fp=root / pth.name,
                        url=package["url"],
                    )
                )
                urls.append(url)
        await asyncio.gather(*tasks)
        async with aiofiles.open(root / Path("urls.txt").expanduser(), mode="w") as f:
            for url in urls:
                await f.write(url + "\n")
        async with aiofiles.open(root / Path("urls").expanduser(), mode="w") as f:
            await f.write("\n")


async def download_url(session: aiohttp.ClientSession, sha256: str, fp: Path, url: str):
    print("Downloading", url)
    async with session.get(url) as resp:
        async with aiofiles.open(fp, mode="wb") as f:
            async for chunk in resp.content.iter_chunked(10000):
                await f.write(chunk)

    loop = asyncio.get_running_loop()
    target_dir = (
        str(fp).removesuffix(".tar.bz2").removesuffix(".tar").removesuffix(".conda")
    )
    await loop.run_in_executor(p, extract, sha256, str(fp), target_dir)


if __name__ == "__main__":
    with ProcessPoolExecutor(max_workers=10) as p:
        asyncio.run(main(cheap["package"]))