#! /usr/bin/env python

import os
import argparse
import rivet
from glob import glob
from datetime import datetime
import json
import logging

## Parse arguments
parser = argparse.ArgumentParser(description="Compile JSON file giving analyses implemented in Rivet for e.g. exposing to HEPData")
parser.add_argument("OUTFILE", nargs="?", type=str, default="analyses.json", help="the file to write into (default: %(default)s)")
parser.add_argument("--schema-version", dest="SCHEMA_VERSION", type=str, default="1.0.0", help="the analyses JSON schema version to use (default: %(default)s)")
parser.add_argument("--check", dest="CHECK", action="store_true", help="whether to check the produced JSON file against the schema (default: %(default)s)")
verbgroup = parser.add_argument_group("verbosity control")
verbgroup.add_argument("-l", dest="LOGLEVEL", default=logging.INFO, help="set a log level")
verbgroup.add_argument("-v", "--verbose", action="store_const", const=logging.DEBUG, dest="LOGLEVEL",
                       default=logging.INFO, help="print debug (very verbose) messages")
verbgroup.add_argument("-q", "--quiet", action="store_const", const=logging.WARNING, dest="LOGLEVEL",
                       default=logging.INFO, help="be very quiet")
verbgroup.add_argument("--silent", action="store_const", const=logging.CRITICAL, dest="LOGLEVEL",
                       default=logging.INFO, help="minimal verbosity (critical messages only)")
args = parser.parse_args()

SCHEMA_URL = f"https://raw.githubusercontent.com/HEPData/hepdata/refs/heads/main/hepdata/templates/analyses_schema/{args.SCHEMA_VERSION}/analyses_schema.json"

## Configure logging
try: # try to convert to int, use str if failed
    args.LOGLEVEL = int(args.LOGLEVEL)
except:
    pass
logging.basicConfig(level=args.LOGLEVEL, format="%(levelname)s: %(message)s")
if not args.OUTFILE.endswith(".json") and args.OUTFILE != "-":
    args.OUTFILE += ".json"
logging.info(f"Using output file name '{args.OUTFILE}' and schema version '{args.SCHEMA_VERSION}'")

## Add info file locations
dirpatt = os.path.join(os.getcwd(), "..", "..", "analyses", "plugin*")
for d in glob(dirpatt):
    logging.debug(f"Found directory '{d}'")
    rivet.addAnalysisDataPath(os.path.abspath(d))

## Retrieve analyses
ins_anas = rivet.AnalysisLoader.analysisInspireIDMap()

## Create dictionary
if args.SCHEMA_VERSION == "0.1.0":
    anas_dic = {iid: [ana] for iid, ana in ins_anas.items()}
    indent = None

else: # SCHEMA_VERSION >= 1.0.0
    anas_dic = {
        "schema_version" : args.SCHEMA_VERSION,
        "tool" : "Rivet",
        "version" : rivet.__version__,
        "date_created" : datetime.now().astimezone().isoformat(),
        "implementations_description" : "Rivet analysis",
        "url_templates" : {
            "main_url" : "https://rivet.hepforge.org/analyses/{name}"
        },
        "analyses" : [],
        "implementations_license" : {
            "name" : "GPLv3",
            "url" : "https://www.gnu.org/licenses/gpl-3.0.en.html"
        }
    }
    indent = 2 # indent file for better readability

    ## Fill analyses into dictionary
    for iid, ana in sorted(ins_anas.items()):
        item = {
            "inspire_id" : int(iid),
            "implementations" : [{
                "name" : ana
            }]
        }
        anas_dic["analyses"].append(item)

## Write out
if args.OUTFILE == "-":
    print(str(anas_dic))
    logging.debug(f"Wrote to stdout.")
else:
    with open(args.OUTFILE, "w") as f:
        json.dump(anas_dic, f, indent=indent)
    logging.debug(f"Wrote to {args.OUTFILE}.")

## Check against analyses JSON schema
## Only sensible if not using stdout as output
if args.CHECK and not args.OUTFILE == "-":
    import jsonschema
    schema_file_name = "analyses_schema.json"
    os.system(f"wget -q {SCHEMA_URL} -O {schema_file_name}")

    with open(schema_file_name) as f:
        schema = json.load(f)
    with open(args.OUTFILE) as f:
        output_json = json.load(f)

    jsonschema.validate(instance=output_json, schema=schema)
    logging.info(f"JSON file {args.OUTFILE} is valid according to the schema '{schema_file_name}' from '{SCHEMA_URL}'.")
