#!/usr/bin/env python3

import os
from pathlib import Path
from typing import List

BASE_URL = "https://huggingface.co/csukuangfj/k2/resolve/main/"


def generate_url(files: List[str]) -> List[str]:
    ans = []
    base = BASE_URL
    for f in files:
        ans.append(base + str(f))
    return ans


def get_all_files(d: str, suffix: str) -> List[str]:
    if "*" in suffix:
        suffix = suffix.split(".")[-1]

    d_list = d

    if isinstance(d_list, str):
        d_list = [d_list]

    ss = []
    for d in d_list:
        for root, _, files in os.walk(d):
            for f in files:
                if f.endswith(suffix):
                    ss.append(os.path.join(root, f))

    return sorted(ss, reverse=True)


def to_markdown(urls: List[str]) -> List[str]:
    ans = []
    for url in urls:
        name = url.rsplit("/", maxsplit=1)[1]
        ans.append(f"[{name}]({url})")
    return ans


def to_file(filename: str, md: List[str]):
    with open(filename, "w") as f:
        f.write("<br/>\n".join(md))


def main():
    cpu = get_all_files("cpu", suffix="*.whl")
    cpu_wheels = generate_url(cpu)
    cpu_wheels = to_markdown(cpu_wheels)
    to_file("cpu.md", cpu_wheels)

    cuda = get_all_files("cuda", suffix="*.whl")
    cuda_wheels = generate_url(cuda)
    cuda_wheels = to_markdown(cuda_wheels)
    to_file("cuda.md", cuda_wheels)


if __name__ == "__main__":
    main()