#!/usr/bin/env python3
import os
import re
from pathlib import Path

from util import generate_url, get_all_files


class Wheel:
    def __init__(self, full_name: str, url: str):
        """
        Args:
          full_name:
            Example: k2_sherpa-1.3.dev20230725+cpu.torch2.0.1-cp310-cp310-macosx_10_9_x86_64.whl
        """
        self.full_name = full_name
        pattern = r"k2_sherpa-(\d)\.(\d)+((\.)(\d))?\.dev(\d{8})\+cpu\.torch(\d\.\d+\.\d)-cp(\d+)"
        m = re.search(pattern, full_name)

        self.sherpa_major = int(m.group(1))
        self.sherpa_minor = int(m.group(2))
        self.sherpa_patch = int(m.group(5)) if m.group(5) else 0
        self.sherpa_date = int(m.group(6))
        self.torch_version = m.group(7)
        self.py_version = int(m.group(8))
        self.url = url

    def __str__(self):
        return self.url

    def __repr__(self):
        return self.url


def sort_by_wheel(x: Wheel):
    torch_major, torch_minor, torch_patch = x.torch_version.split(".")
    torch_major = int(torch_major)
    torch_minor = int(torch_minor)
    torch_patch = int(torch_patch)
    return (
        x.sherpa_major,
        x.sherpa_minor,
        x.sherpa_patch,
        x.sherpa_date,
        torch_major,
        torch_minor,
        torch_patch,
        x.py_version,
    )


def main():
    wheels = get_all_files("macos", suffix="*.whl")
    wheels += get_all_files("ubuntu-cpu", suffix="*.whl")
    wheels += get_all_files("cpu", suffix="*.whl")
    urls = generate_url(wheels)

    wheels = []
    for url in urls:
        full_name = url.rsplit("/", maxsplit=1)[1]
        wheels.append(Wheel(full_name, url))

    wheels.sort(reverse=True, key=sort_by_wheel)
    with open("./cpu.html", "w") as f:
        for w in wheels:
            f.write(f'<a href="{w.url}">{w.full_name}</a><br/>\n')

    with open("./cpu-cn.html", "w") as f:
        for w in wheels:
            url = w.url.replace("huggingface.co", "hf-mirror.com")
            f.write(f'<a href="{url}">{w.full_name}</a><br/>\n')


if __name__ == "__main__":
    main()