#!/usr/bin/env python3
"""Validate sample WAI manifests against the generated v2 vocabulary.

Each manifest is expanded through vocab/v2.jsonld exactly as written on the
wire, with no @type added, and validated against vocab/v2.shapes.ttl with
v2.ttl as the ontology graph. Every case states whether it must conform. Exit
status 0 when every case behaves as stated, 1 otherwise.

Needs rdflib, pyld and pyshacl (CI pins their versions). Run from anywhere:

    python3 wai/vocab/check_shapes.py
"""

import copy
import json
import pathlib
import sys

import rdflib
from pyld import jsonld
from pyshacl import validate

VOCAB = pathlib.Path(__file__).resolve().parent


def wai1(capability, fallback=None, media="image", intent="replicate", **extra):
    m = {
        "wai": "1.0",
        "media": media,
        "intent": intent,
        "model_requirement": {"capability": capability},
        "conditioning": {"kind": "png"},
        "target": {"note": "opaque JSON, carried as a literal"},
    }
    if fallback is not None:
        m["model_requirement"]["fallback"] = fallback
    m.update(extra)
    return m


def wai2(*components, intent="improve"):
    renditions = []
    for c in components:
        if isinstance(c, tuple):
            cap, base = c
            renditions.append({"capability": cap, "kind": "bytes", "role": "companion", "companion_of": base})
        else:
            renditions.append({"capability": c, "kind": "bytes"})
    return {"wai": "1.1", "media": "video", "intent": intent, "renditions": renditions, "target": None}


SHA_A = "ab" * 32
SHA_B = "cd" * 32


def pin(**over):
    """A parameter-set pin (SPEC section 3.1); a member set to None is dropped."""
    p = {"id": "m", "sha256": SHA_A, "determinism": "int", "form": "wai.prior.set",
         "tier": "decode-equivalence", "supersedes": {"sha256": SHA_B}}
    p.update(over)
    return {k: v for k, v in p.items() if v is not None}


def pinned(prior):
    m = wai1("wai.neural.int_mlicpp")
    m["conditioning"] = {"kind": "int_mlicpp", "prior": prior}
    return m


def pinned2(prior):
    m = wai2("wai.neural.int_mlicpp", intent="replicate")
    m["renditions"][0]["prior"] = prior
    return m


# (name, manifest, must conform)
CASES = [
    ("WAI1 png, untyped", wai1("wai.image.png"), True),
    ("WAI1 png with a png fallback", wai1("wai.image.jxl", "wai.image.png"), True),
    ("WAI1 media value no revision lists (a reader accepts it)", wai1("wai.mesh.glb", media="model"), True),
    ("WAI1 candidate capability", wai1("wai.splat.sog", media="volumetric"), False),
    ("WAI1 candidate capability, the other candidate", wai1("wai.tensor.nnc", media="binary"), False),
    ("WAI1 candidate fallback", wai1("wai.image.png", "wai.tensor.nnc"), False),
    ("WAI1 companion-only fallback", wai1("wai.video.av1", "wai.meta.hdr10plus", media="video"), False),
    ("WAI1 unregistered capability", wai1("wai.image.nope"), False),
    ("WAI1 record as capability", wai1("wai.video.manifest", media="video"), False),
    ("WAI1 with no intent", {k: v for k, v in wai1("wai.image.png").items() if k != "intent"}, False),
    ("WAI2 rendition + companion, untyped", wai2("wai.video.av1", ("wai.meta.hdr10plus", "wai.video.av1")), True),
    ("WAI2 candidate rendition", wai2("wai.splat.sog", "wai.image.png"), False),
    ("WAI2 candidate companion", wai2("wai.video.av1", ("wai.splat.sog", "wai.video.av1")), False),
    ("WAI2 companion bound to a candidate base", wai2("wai.video.av1", ("wai.meta.hdr10plus", "wai.tensor.nnc")), False),
    ("WAI1 and WAI2 at once", dict(wai1("wai.image.png"), renditions=[{"capability": "wai.image.png", "kind": "png"}]), False),
    # The parameter-set pin (SPEC section 3.1), under conditioning and on a WAI2 component.
    ("WAI1 pin with every member", pinned(pin()), True),
    ("WAI1 pin with only the required members", pinned(pin(form=None, tier=None, supersedes=None)), True),
    ("WAI1 pin with a tier token no revision registers", pinned(pin(tier="decode_equivalence")), False),
    ("WAI1 pin whose digest is not lowercase hex", pinned(pin(sha256=SHA_A.upper())), False),
    ("WAI1 pin with no determinism", pinned(pin(determinism=None)), False),
    ("WAI1 pin whose supersedes names no digest", pinned(pin(supersedes={"form": "wai.prior.set"})), False),
    ("WAI2 component pin", pinned2(pin()), True),
    ("WAI2 component pin with an unregistered tier token", pinned2(pin(tier="exact")), False),
]


def main():
    context = json.loads((VOCAB / "v2.jsonld").read_text())["@context"]
    shapes = rdflib.Graph().parse(VOCAB / "v2.shapes.ttl", format="turtle")
    ontology = rdflib.Graph().parse(VOCAB / "v2.ttl", format="turtle")
    failures = 0
    for name, manifest, want in CASES:
        doc = copy.deepcopy(manifest)
        doc["@context"] = context
        data = rdflib.Graph().parse(data=json.dumps(jsonld.expand(doc)), format="json-ld")
        conforms, _, text = validate(data, shacl_graph=shapes, ont_graph=ontology, inference="none")
        ok = conforms == want
        failures += not ok
        print(f"{'ok  ' if ok else 'FAIL'}  {name}: conforms={conforms}, expected {want}")
        if not ok:
            print(text)
    # The typed form validates the same way as the untyped one.
    typed = wai1("wai.splat.sog")
    typed["@type"] = "Manifest"
    typed["@context"] = context
    data = rdflib.Graph().parse(data=json.dumps(jsonld.expand(typed)), format="json-ld")
    conforms, _, _ = validate(data, shacl_graph=shapes, ont_graph=ontology, inference="none")
    ok = conforms is False
    failures += not ok
    print(f"{'ok  ' if ok else 'FAIL'}  WAI1 candidate capability, typed: conforms={conforms}, expected False")
    if failures:
        print(f"{failures} case(s) did not behave as stated")
        return 1
    print("v2 shapes: every case behaves as stated")
    return 0


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