#!/usr/bin/env python3
# Old Nvidia drivers do not have a JSON with supported models,
# but JSONs of newer drivers specify GPUs supported by older drivers.
# This scripts takes a JSON from a newer driver and converts it
# into a JSON for an older driver which does not ship a native JSON.

from signal import signal, SIGPIPE, SIG_DFL
signal(SIGPIPE,SIG_DFL)

import os
import json
import sys

EXIT_CODE_EINPUT = 10
EXIT_CODE_PARSEERROR = 20

FILE = os.getenv("FILE")
if not FILE:
    print("Env FILE is not set! Set FILE=supported-gpus.json", file=sys.stderr)
    exit(EXIT_CODE_EINPUT)

BRANCH = os.getenv("BRANCH")
if not BRANCH:
    print("Env BRANCH is not set! Set e.g. BRANCH=390 to generate JSON for branch 390.xx", file=sys.stderr)
    exit(EXIT_CODE_EINPUT)

try:
    with open(FILE) as f:
        data = f.read()
        array = json.loads(data)
except:
    print(f"Error decoding JSON from file {FILE}!", file=sys.stderr)
    exit(EXIT_CODE_EINPUT)

chips_arr = []

for chip in array["chips"]:
    if not chip.get("legacybranch"):
        continue
    if not chip["legacybranch"].startswith(f"{BRANCH}."):
        continue
    # devid (Device ID) cannot be empty, other params can be empty
    if not chip.get("devid"):
        print("Device ID cannot be empty!", file=sys.stderr)
        exit(EXIT_CODE_PARSEERROR)
    arr = {
        "devid": chip["devid"],
        "subdevid": chip.get("subdevid"),
        "subvendorid": chip.get("subvendorid"),
        "name": chip["name"],
        "features": chip.get("features"),
    }
    chips_arr.append(arr)

chips_arr = {"chips": chips_arr}
print(json.dumps(chips_arr))

