#!/usr/bin/env python3
"""Inspect, export, and safely import RIFT JAX compilation caches."""

import argparse
import json
import os
import sys

from RIFT.jax_cache import (
    compatibility_key,
    configure_persistent_cache,
    export_bundle,
    import_bundle,
    runtime_compatibility,
)


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--cache-root", help="cache root (default: RIFT/XDG cache root)")
    commands = parser.add_subparsers(dest="command", required=True)
    commands.add_parser("fingerprint", help="print this runtime/device cache identity")
    export = commands.add_parser("export", help="export the active warmed cache")
    export.add_argument("output")
    export.add_argument("--profile")
    export.add_argument("--shape", action="append", default=[], metavar="NAME=VALUE")
    ingest = commands.add_parser("import", help="validate and import a cache bundle")
    ingest.add_argument("bundle")
    ingest.add_argument("--expect-profile")
    args = parser.parse_args(argv)

    # Keep ``--help`` usable in RIFT's base installation, where JAX is an
    # optional dependency.  Operational subcommands require JAX, but argparse
    # exits for help before this import is reached.
    import jax

    compatibility = runtime_compatibility(jax)
    if args.command == "fingerprint":
        print(json.dumps({"compatibility_key": compatibility_key(compatibility),
                          "compatibility": compatibility}, indent=2, sort_keys=True))
        return 0

    configure_args = (["--jax-cache-dir", args.cache_root] if args.cache_root else [])
    active = configure_persistent_cache(jax, configure_args)
    if active is None:
        parser.error("the JAX cache directory is unavailable; choose a writable --cache-root")
    if args.command == "export":
        shapes = {}
        for item in args.shape:
            if "=" not in item:
                parser.error("--shape must be NAME=VALUE")
            key, value = item.split("=", 1)
            shapes[key] = value
        manifest = export_bundle(active, args.output, compatibility, args.profile, shapes)
        print(json.dumps(manifest, indent=2, sort_keys=True))
    else:
        root = args.cache_root or str(active.parent)
        # ``active`` is authoritative even when the standard JAX environment
        # variable names an exact (non-namespaced) directory. Importing into a
        # derived sibling would succeed but the next ILE would never read it.
        destination = import_bundle(args.bundle, root, compatibility,
                                    args.expect_profile, destination=active)
        print(destination)
    return 0


if __name__ == "__main__":
    sys.exit(main())
