from __future__ import annotations

import argparse
import re
import urllib.request
from pathlib import Path


PACKAGES_URL = "https://raw.githubusercontent.com/mjun0812/flash-attention-prebuild-wheels/main/doc/packages.md"
WHEEL_URL_RE = re.compile(
    r"https://github\.com/mjun0812/flash-attention-prebuild-wheels/releases/download/[^\s)]+\.whl"
)
DEFAULT_OUTPUT = Path(__file__).resolve().parents[1] / "comfy" / "custom_node_facade" / "flash_attention_wheels.py"


def _render_module(urls: list[str]) -> str:
    lines = [
        "from __future__ import annotations",
        "",
        "# Generated by scripts/snapshot_flash_attention_wheels.py.",
        f"# Source: {PACKAGES_URL}",
        "",
        "FLASH_ATTENTION_WHEEL_URLS: tuple[str, ...] = (",
    ]
    lines.extend(f"    {url!r}," for url in urls)
    lines.extend([
        ")",
        "",
    ])
    return "\n".join(lines)


def snapshot(output: Path) -> int:
    with urllib.request.urlopen(PACKAGES_URL, timeout=60) as response:
        markdown = response.read().decode("utf-8")
    urls = sorted(set(WHEEL_URL_RE.findall(markdown)))
    output.write_text(_render_module(urls), encoding="utf-8")
    return len(urls)


def main() -> None:
    parser = argparse.ArgumentParser(description="Snapshot Flash Attention wheel links from upstream packages.md")
    parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
    args = parser.parse_args()
    count = snapshot(args.output)
    print(f"Wrote {count} Flash Attention wheel URLs to {args.output}")


if __name__ == "__main__":
    main()
