#!/usr/bin/env python3

import base64
import hashlib
import os
from pathlib import Path
import re
import subprocess
import sys


URL_RE = re.compile(r'^(?P<indent>\s*)url = "(?P<url>https://.*/packages\.adb)";$')
HASH_RE = re.compile(r'^(?P<indent>\s*)hash = ".*";$')
NAME_RE = re.compile(r'^\s*name = "(?P<name>[^"]+)";$')


def vendor_block(lines, start, end, nix_file, indexes_dir):
    url_match = None
    hash_index = None
    name = None

    for index in range(start, end):
        stripped = lines[index].rstrip("\n")
        if match := URL_RE.match(stripped):
            url_match = (index, match)
        if HASH_RE.match(stripped):
            hash_index = index
        if match := NAME_RE.match(stripped):
            name = match.group("name")

    if url_match is None:
        return False
    if hash_index is None or name is None:
        raise RuntimeError(f"incomplete sourceInfo in {nix_file}")

    url_index, match = url_match
    destination = indexes_dir / name
    temporary = destination.with_suffix(destination.suffix + ".tmp")
    subprocess.run(
        ["curl", "--fail", "--location", "--silent", "--show-error",
         "--output", temporary, match.group("url")],
        check=True,
    )
    temporary.replace(destination)

    digest = base64.b64encode(hashlib.sha256(destination.read_bytes()).digest()).decode()
    hash_indent = HASH_RE.match(lines[hash_index].rstrip("\n")).group("indent")
    relative = os.path.relpath(destination, nix_file.parent)
    if not relative.startswith("."):
        relative = f"./{relative}"

    lines[hash_index] = f'{hash_indent}hash = "sha256-{digest}";\n'
    lines[url_index] = f'{match.group("indent")}url = "file://${{{relative}}}";\n'
    return True


def vendor_file(nix_file, indexes_dir):
    lines = nix_file.read_text().splitlines(keepends=True)
    block_start = None
    changed = 0

    for index, line in enumerate(lines):
        if line.strip() == "sourceInfo = {":
            block_start = index + 1
        elif block_start is not None and line.strip() == "};":
            changed += vendor_block(lines, block_start, index, nix_file, indexes_dir)
            block_start = None

    if changed:
        nix_file.write_text("".join(lines))

    return changed


def main():
    if len(sys.argv) != 2:
        raise SystemExit(f"usage: {sys.argv[0]} CACHE-DIRECTORY")

    cache_dir = Path(sys.argv[1]).resolve()
    if not (cache_dir / "default.nix").is_file():
        raise SystemExit(f"not an OpenWrt cache directory: {cache_dir}")

    indexes_dir = cache_dir / "indexes"
    indexes_dir.mkdir(exist_ok=True)

    changed = 0
    for nix_file in sorted(cache_dir.rglob("*.nix")):
        changed += vendor_file(nix_file, indexes_dir)

    print(f"Vendored {changed} repository indexes in {cache_dir}", file=sys.stderr)


if __name__ == "__main__":
    main()
