"""Check a dataset directory against the croissant-rai profile.

Usage: python check_croissant.py DATA_DIR > predicate.json
"""
import hashlib
import json
import sys
from pathlib import Path

import mlcroissant as mlc

RAI_FIELDS = [
    "dataCollection", "dataCollectionType", "dataAnnotationProtocol",
    "dataAnnotationPlatform", "dataLimitations", "dataBiases",
    "personalSensitiveInformation", "dataUseCases",
]


def sha256(path):
    return hashlib.sha256(path.read_bytes()).hexdigest()


def check(data):
    meta_path = data / "croissant.json"
    meta = json.loads(meta_path.read_text())
    results = []

    def record(requirement, ok, evidence):
        results.append({"requirement": requirement,
                        "status": "pass" if ok else "fail",
                        "evidence": evidence})

    # 1. The official Croissant library loads the document without errors.
    try:
        mlc.Dataset(jsonld=str(meta_path))
        record("cr-valid", True, {"conformsTo": meta.get("conformsTo")})
    except Exception as e:
        record("cr-valid", False, {"error": str(e)[:200]})

    # 2. Every packed data file has a matching sha256 in the metadata.
    declared = {f.get("name"): f.get("sha256") for f in meta.get("distribution", [])}
    files = {p.name: sha256(p) for p in sorted(data.glob("*.parquet"))}
    matches = {n: "match" if declared.get(n) == h else "mismatch" for n, h in files.items()}
    record("cr-files-match", bool(files) and set(matches.values()) == {"match"}, matches)

    # 3. The chosen Responsible AI fields are present and non-empty.
    for field in RAI_FIELDS:
        value = meta.get(f"rai:{field}")
        record(f"rai-{field}", bool(value), {f"rai:{field}": str(value)[:200] if value else None})

    # 4. The metadata says where the data came from (PROV-O).
    prov = meta.get("prov:wasGeneratedBy") or meta.get("prov:wasDerivedFrom")
    record("prov-provenance", bool(prov), {"prov": "present" if prov else None})
    return results


if __name__ == "__main__":
    results = check(Path(sys.argv[1]))
    me = Path(__file__)
    json.dump({
        "profile": {
            "id": "croissant-rai",
            "version": "1.1+rai-1.0",
            "uri": "https://docs.mlcommons.org/croissant/docs/croissant-spec-1.1.html",
        },
        "level": "pass" if all(r["status"] == "pass" for r in results) else "fail",
        "checker": {"name": me.name, "digest": "sha256:" + sha256(me)},
        "results": results,
    }, sys.stdout, indent=2)
    print()