#!/usr/bin/env python3
"""check-wai-corpus-licences.py: every vector of a WAI corpus that keeps licence records
has one, and every record states a licence that permits redistribution.

A corpus listed in CORPORA keeps `licences.json` beside its README:

  {"format": "wai-corpus-licences/1",
   "vectors": {"<name>": {"licence": "<SPDX id>", "content": "generated" | "third-party",
                          "generator": "<path below wai/>", "input": "<what it reads>",
                          "source": "<where a third-party input comes from>",
                          "training": {"content": ..., "generator": ..., "input": ...,
                                       "source": ..., "licence": ...}}}}

What each corpus is held to:

  1. `licences.json` exists, parses, and has this format.
  2. Every vector has a record: each directory under `vectors/`, and each name in the
     corpus's `expected_hashes.json` and `sha256.json`, is a key of `vectors`. Every key
     names a directory under `vectors/`.
  3. `licence` is one of LICENCES, the licences that permit redistribution.
  4. `content` "generated" names a `generator` that exists below `wai/` and a non-empty
     `input`, and no `source`. `content` "third-party" names a non-empty `source`; its
     `licence` is the source's, and any `generator` it names exists.
  5. A vector whose `expected_hashes.json` entry has `"trained": true` has a `training`
     record, held to rule 4, whose `licence`, required when its content is third-party,
     is one of LICENCES; a vector not marked trained has none.
  6. A record has no member other than these.

The check holds what the records state, not their truth: that is a reviewer's, and for
`neural-conformance`'s trained vector also `tools/wai_int_synth_real.py --check`'s.

    check-wai-corpus-licences.py [--root DIR]

Exit 0 when every corpus holds, printing each corpus's vectors by licence; 1 with one
line a finding otherwise. Standard library only.
"""
import json
import os
import sys

# The corpora that keep a licence record for each vector, relative to the root.
CORPORA = ["wai/neural-conformance", "wai/tensor-conformance"]

FORMAT = "wai-corpus-licences/1"

# SPDX identifiers of the licences a committed vector, or the data a committed model was
# trained on, may carry: each permits redistribution.
LICENCES = {"Apache-2.0", "BSD-2-Clause", "BSD-3-Clause", "CC-BY-4.0", "CC0-1.0", "MIT"}

VECTOR_KEYS = {"licence", "content", "generator", "input", "source", "training"}
TRAINING_KEYS = {"content", "generator", "input", "source", "licence"}


def names_in(path):
    """The top-level keys of a JSON object file, or None when the file does not exist."""
    if not os.path.exists(path):
        return None
    with open(path, encoding="utf-8") as fh:
        doc = json.load(fh)
    return doc if isinstance(doc, dict) else {}


def nonempty(v):
    return isinstance(v, str) and v.strip() != ""


def check_content(root, where, rec, findings):
    """Rule 4, for a vector record or a training record."""
    content = rec.get("content")
    gen = rec.get("generator")
    if gen is not None and not (nonempty(gen) and os.path.isfile(os.path.join(root, "wai", gen))):
        findings.append(f"{where}: generator {gen!r} is not a file below wai/")
    if content == "generated":
        if gen is None:
            findings.append(f"{where}: generated content names no generator")
        if not nonempty(rec.get("input")):
            findings.append(f"{where}: generated content names no input")
        if "source" in rec:
            findings.append(f"{where}: generated content has a source")
    elif content == "third-party":
        if not nonempty(rec.get("source")):
            findings.append(f"{where}: third-party content names no source")
    else:
        findings.append(f"{where}: content {content!r} is neither \"generated\" nor \"third-party\"")


def check_corpus(root, corpus, findings, split):
    base = os.path.join(root, corpus)
    path = os.path.join(base, "licences.json")
    if not os.path.isfile(path):
        findings.append(f"{corpus}: no licences.json")
        return
    try:
        with open(path, encoding="utf-8") as fh:
            doc = json.load(fh)
    except (OSError, ValueError) as e:
        findings.append(f"{corpus}/licences.json: does not parse: {e}")
        return
    if not isinstance(doc, dict) or doc.get("format") != FORMAT or not isinstance(doc.get("vectors"), dict):
        findings.append(f"{corpus}/licences.json: not a {FORMAT} document with a `vectors` object")
        return
    records = doc["vectors"]
    vdir = os.path.join(base, "vectors")
    dirs = sorted(d for d in os.listdir(vdir) if os.path.isdir(os.path.join(vdir, d))) if os.path.isdir(vdir) else []
    try:
        hashes = names_in(os.path.join(base, "expected_hashes.json")) or {}
        sha = names_in(os.path.join(base, "sha256.json")) or {}
    except ValueError as e:
        findings.append(f"{corpus}: a hash manifest does not parse: {e}")
        return
    for name in sorted(set(dirs) | set(hashes) | set(sha)):
        if name not in records:
            findings.append(f"{corpus}: vector {name} has no licence record")
    for name in sorted(records):
        where = f"{corpus}/licences.json: {name}"
        rec = records[name]
        if name not in dirs:
            findings.append(f"{where}: no directory vectors/{name}")
        if not isinstance(rec, dict):
            findings.append(f"{where}: not an object")
            continue
        for k in sorted(set(rec) - VECTOR_KEYS):
            findings.append(f"{where}: unknown member {k!r}")
        lic = rec.get("licence")
        if lic not in LICENCES:
            findings.append(f"{where}: licence {lic!r} is not one of {', '.join(sorted(LICENCES))}")
        check_content(root, where, rec, findings)
        trained = isinstance(hashes.get(name), dict) and hashes[name].get("trained") is True
        tr = rec.get("training")
        if trained and tr is None:
            findings.append(f"{where}: a trained model with no training record")
        elif not trained and tr is not None:
            findings.append(f"{where}: a training record on a vector whose model is not marked trained")
        if isinstance(tr, dict):
            tw = where + " training"
            for k in sorted(set(tr) - TRAINING_KEYS):
                findings.append(f"{tw}: unknown member {k!r}")
            check_content(root, tw, tr, findings)
            if (tr.get("content") == "third-party" or "licence" in tr) and tr.get("licence") not in LICENCES:
                findings.append(f"{tw}: licence {tr.get('licence')!r} is not one of {', '.join(sorted(LICENCES))}")
        elif tr is not None:
            findings.append(f"{where}: training is not an object")
        if lic in LICENCES:
            split.setdefault(corpus, {}).setdefault(lic, []).append(name)


def main(argv):
    root = "."
    args = argv[1:]
    if args[:1] == ["--root"] and len(args) == 2:
        root = args[1]
    elif args:
        print("usage: check-wai-corpus-licences.py [--root DIR]", file=sys.stderr)
        return 2
    findings, split = [], {}
    for corpus in CORPORA:
        check_corpus(root, corpus, findings, split)
    for f in findings:
        print(f)
    if findings:
        print(f"check-wai-corpus-licences: {len(findings)} finding(s)")
        return 1
    for corpus in CORPORA:
        for lic, names in sorted(split.get(corpus, {}).items()):
            print(f"{corpus}: {lic}: {', '.join(names)}")
    print("check-wai-corpus-licences: every vector has a licence record")
    return 0


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