#!/usr/bin/env python3

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
"""Regenerate `docker_requirements.txt` from `docker_requirements.in` using `uv`, targeting the
Docker image runtime, and filter out `torch` + CUDA wheels so Docker doesn't try to reinstall
PyTorch.

Usage:
  python3 kimodo/scripts/lock_requirements.py

Optional args:
  --python-version 3.10
  --python-platform x86_64-manylinux2014
  --in docker_requirements.in
  --out docker_requirements.txt
"""

import argparse
import shutil
import subprocess
from pathlib import Path
from typing import Iterable

DEFAULT_PYTHON_VERSION = "3.10"
DEFAULT_PYTHON_PLATFORM = "x86_64-manylinux2014"

# Packages to omit from the lockfile because the Docker base image already provides torch+CUDA.
OMIT_NAMES = {"torch", "triton", "networkx", "sympy", "mpmath"}
OMIT_PREFIXES = ("nvidia-",)


def _run(cmd: list[str]) -> None:
    print("+", " ".join(cmd))
    subprocess.run(cmd, check=True)


def _ensure_uv() -> None:
    if shutil.which("uv") is None:
        raise SystemExit(
            "ERROR: `uv` is not installed or not on PATH.\n"
            "Install it (one of):\n"
            "  - pipx install uv\n"
            "  - python -m pip install --user uv\n"
            "Then rerun this script."
        )


def _parse_req_name(line: str) -> str:
    # uv emits `name==version` lines.
    s = line.strip()
    if "==" in s:
        return s.split("==", 1)[0].strip()
    # Fallback: treat the whole token before space as name.
    return s.split()[0].strip()


def _iter_blocks(lines: list[str]) -> Iterable[list[str]]:
    """Split a docker_requirements.txt into blocks: [top-level req line + indented comments]."""
    i = 0
    n = len(lines)
    while i < n:
        line = lines[i]
        # Header/comments/blank
        if line.startswith("#") or line.strip() == "":
            yield [line]
            i += 1
            continue

        # Top-level requirement line
        if not line.startswith(" "):
            block = [line]
            i += 1
            while i < n and (lines[i].startswith(" ") or lines[i].strip() == "" or lines[i].startswith("#")):
                # Stop if we hit another top-level requirement line
                if not lines[i].startswith(" ") and not lines[i].startswith("#") and lines[i].strip() != "":
                    break
                block.append(lines[i])
                i += 1
            yield block
            continue

        # Indented line without a requirement header (shouldn't happen, but keep)
        yield [line]
        i += 1


def _should_omit(req_line: str) -> bool:
    name = _parse_req_name(req_line)
    if name in OMIT_NAMES:
        return True
    for pfx in OMIT_PREFIXES:
        if name.startswith(pfx):
            return True
    return False


def filter_lockfile(path: Path) -> None:
    lines = path.read_text(encoding="utf-8").splitlines(True)
    out: list[str] = []

    inserted_note = False
    for block in _iter_blocks(lines):
        first = block[0]

        # After the uv header lines, insert a short note once.
        if (not inserted_note) and first.startswith("# This file was autogenerated by uv"):
            out.extend(block)
            out.append(
                "# NOTE: `torch` (and its CUDA wheels) are intentionally omitted from this lockfile.\n"
                "# The Docker base image (nvcr.io/nvidia/pytorch) already provides a tested PyTorch build.\n"
                "#\n"
            )
            inserted_note = True
            continue

        if first.startswith("#") or first.strip() == "":
            out.extend(block)
            continue

        if _should_omit(first):
            continue

        out.extend(block)

    path.write_text("".join(out), encoding="utf-8")


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--in", dest="in_file", default="docker_requirements.in")
    ap.add_argument("--out", dest="out_file", default="docker_requirements.txt")
    ap.add_argument("--python-version", default=DEFAULT_PYTHON_VERSION)
    ap.add_argument("--python-platform", default=DEFAULT_PYTHON_PLATFORM)
    args = ap.parse_args()

    _ensure_uv()

    in_path = Path(args.in_file)
    out_path = Path(args.out_file)
    if not in_path.exists():
        raise SystemExit(f"ERROR: missing {in_path}")

    _run(
        [
            "uv",
            "pip",
            "compile",
            "-U",
            str(in_path),
            "-o",
            str(out_path),
            "--python-version",
            args.python_version,
            "--python-platform",
            args.python_platform,
        ]
    )
    filter_lockfile(out_path)
    print(f"OK: wrote {out_path} (filtered torch/CUDA wheels)")


if __name__ == "__main__":
    main()