#!/usr/bin/env python3
"""Emissar Mandate v0.1 example checks.

    python validate.py [--a2a-schema PATH_OR_URL ...]

1. Validates the example files written by `node mandate-demo.mjs --write .` against
   ../../schemas/mandate-v0.1.schema.json (Draft 2020-12, with format checks): the processed mandate
   payload (root), JOSE headers, KB-JWT payload, the A2A extension metadata entries, the Agent Card params
   and the Status List Token payload.
2. Negative tests: objects the schema MUST reject.
3. Validates the A2A shapes (SendMessageRequest params, SendMessageResponse results, AgentCard) against the
   A2A JSON Schema bundle generated from a2a.proto. Default source: the bundle published at
   https://a2a-protocol.org/latest/spec/a2a.json (uses $defs). The older bundle at /v1.0.0/spec/a2a.json uses a
   different layout (definitions keyed by title, unresolved external $refs) and is not supported here.

Needs: pip install --user jsonschema rfc3339-validator rfc3986-validator
"""
import copy
import json
import pathlib
import sys
import urllib.request

from jsonschema import Draft202012Validator, FormatChecker
from referencing import Registry, Resource

HERE = pathlib.Path(__file__).resolve().parent
SCHEMA = json.loads((HERE.parent.parent / "schemas" / "mandate-v0.1.schema.json").read_text(encoding="utf-8"))
EXT = "https://emissar.ai/ext/mandate/v1"
FORMATS = FormatChecker()
failures = []


def load(name):
    return json.loads((HERE / name).read_text(encoding="utf-8"))


def validator(schema_doc, pointer):
    sid = schema_doc["$id"]
    registry = Registry().with_resource(sid, Resource.from_contents(schema_doc))
    return Draft202012Validator({"$ref": f"{sid}#{pointer}"}, registry=registry, format_checker=FORMATS)


def check(label, schema_doc, pointer, instance, expect_valid=True):
    errors = sorted(validator(schema_doc, pointer).iter_errors(instance), key=lambda e: list(e.path))
    ok = (not errors) if expect_valid else bool(errors)
    print(f"  {'ok  ' if ok else 'FAIL'} {label}" + ("" if expect_valid else " (rejected as expected)" if ok else " (was accepted)"))
    if expect_valid:
        for e in errors[:5]:
            print(f"         {list(e.path)}: {e.message[:160]}")
    if not ok:
        failures.append(label)


def decode_jwt(compact):
    import base64
    h, p, _ = compact.split(".")
    pad = lambda s: s + "=" * (-len(s) % 4)
    return json.loads(base64.urlsafe_b64decode(pad(h))), json.loads(base64.urlsafe_b64decode(pad(p)))


print("Mandate schema: examples")
mandate = load("mandate.json")
presentation = load("presentation.json")
request = load("send-message-request.json")
completed = load("send-message-response-completed.json")
auth_required = load("auth-required-response.json")
card = load("agent-card.json")
status_token = load("status-list-token.json")

check("mandate.json processedPayloadAsPresentedInB (root)", SCHEMA, "", mandate["processedPayloadAsPresentedInB"])
full = copy.deepcopy(mandate["processedPayloadAsPresentedInB"])
full["authorization_details"] = [d["decoded"][1] for d in mandate["disclosures"] if d["label"].startswith("authorization_details")]
full["principal"].update({d["decoded"][1]: d["decoded"][2] for d in mandate["disclosures"] if d["label"].startswith("principal.")})
check("mandate.json with every disclosure applied (root)", SCHEMA, "", full)
check("mandate.json header ($defs/issuerJwtHeader)", SCHEMA, "/$defs/issuerJwtHeader", mandate["header"])
h, p = decode_jwt(mandate["compact"].split("~")[0])
print(f"  {'ok  ' if (h, p) == (mandate['header'], mandate['payloadAsIssued']) else 'FAIL'} compact issuer JWT decodes to the listed header and payload")
if (h, p) != (mandate["header"], mandate["payloadAsIssued"]):
    failures.append("compact decode")
check("presentation.json kbJwtHeader ($defs/kbJwtHeader)", SCHEMA, "/$defs/kbJwtHeader", presentation["kbJwtHeader"])
check("presentation.json kbJwtPayload ($defs/kbJwtPayload)", SCHEMA, "/$defs/kbJwtPayload", presentation["kbJwtPayload"])
kb_h, kb_p = decode_jwt(presentation["presentation"].split("~")[-1])
print(f"  {'ok  ' if (kb_h, kb_p) == (presentation['kbJwtHeader'], presentation['kbJwtPayload']) else 'FAIL'} KB-JWT in the presentation decodes to the listed header and payload")
envelope = request["params"]["message"]["metadata"][EXT]
check("send-message-request.json metadata entry ($defs/presentationEnvelope)", SCHEMA, "/$defs/presentationEnvelope", envelope)
print(f"  {'ok  ' if envelope['presentation'] == presentation['presentation'] else 'FAIL'} request carries the presentation from presentation.json")
check("auth-required-response.json status metadata ($defs/challengeEnvelope)", SCHEMA, "/$defs/challengeEnvelope",
      auth_required["result"]["task"]["status"]["message"]["metadata"][EXT])
check("send-message-response-completed.json artifact metadata ($defs/resultEnvelope)", SCHEMA, "/$defs/resultEnvelope",
      completed["result"]["task"]["artifacts"][0]["metadata"][EXT])
ext = next(e for e in card["capabilities"]["extensions"] if e["uri"] == EXT)
check("agent-card.json extension params ($defs/cardParams)", SCHEMA, "/$defs/cardParams", ext["params"])
check("status-list-token.json payload ($defs/statusListTokenPayload)", SCHEMA, "/$defs/statusListTokenPayload", status_token["payload"])

print("\nMandate schema: negative tests")
bad = copy.deepcopy(full); bad["authorization_details"][0]["limits"]["max_amount_usd"] = 1
check("detail limits with an unknown member", SCHEMA, "", bad, expect_valid=False)
bad = copy.deepcopy(full); bad["authorization_details"][0]["locations"] = ["https://x.example"]
check("detail with a member v0.1 does not define (locations)", SCHEMA, "", bad, expect_valid=False)
bad = copy.deepcopy(full); bad["_sd"] = ["x"]
check("processed payload that still contains _sd", SCHEMA, "", bad, expect_valid=False)
bad = copy.deepcopy(full); del bad["status"]
check("mandate without a status claim", SCHEMA, "", bad, expect_valid=False)
bad = copy.deepcopy(full); bad["cnf"] = {"jwk": presentation["kbJwtHeader"]["jwk"]}
check("mandate bound with cnf.jwk instead of cnf.jkt", SCHEMA, "", bad, expect_valid=False)
bad = copy.deepcopy(full); bad["authorization_details"][0]["limits"]["max_amount"]["amount"] = 200.5
check("amount that is not an integer in minor units", SCHEMA, "", bad, expect_valid=False)
bad = copy.deepcopy(presentation["kbJwtHeader"]); bad["jwk"]["d"] = "AAAA"
check("KB-JWT header carrying a private key", SCHEMA, "/$defs/kbJwtHeader", bad, expect_valid=False)
bad = copy.deepcopy(mandate["header"]); bad["jku"] = "https://attacker.example/jwks.json"
check("issuer JWT header with jku", SCHEMA, "/$defs/issuerJwtHeader", bad, expect_valid=False)
check("presentation without a KB-JWT", SCHEMA, "/$defs/presentationEnvelope",
      {"presentation": "~".join(presentation["presentation"].split("~")[:-1]) + "~"}, expect_valid=False)

print("\nA2A JSON Schema (generated from a2a.proto)")
sources = sys.argv[sys.argv.index("--a2a-schema") + 1:] if "--a2a-schema" in sys.argv else [
    "https://a2a-protocol.org/latest/spec/a2a.json"]
for n, src in enumerate(sources):
    try:
        if src.startswith("https://"):
            with urllib.request.urlopen(src, timeout=30) as r:
                a2a = json.load(r)
        else:
            a2a = json.loads(pathlib.Path(src).read_text(encoding="utf-8"))
    except Exception as e:  # network or file problem: report, do not hide
        print(f"  FAIL could not load {src}: {e}")
        failures.append(src)
        continue
    a2a["$id"] = f"https://a2a-bundle.invalid/{n}/a2a.json"  # local key only; the bundle's own refs are relative
    print(f"  source {src}")
    check("send-message-request.json params ($defs/SendMessageRequest)", a2a, "/$defs/SendMessageRequest", request["params"])
    check("send-message-response-completed.json result ($defs/SendMessageResponse)", a2a, "/$defs/SendMessageResponse", completed["result"])
    check("auth-required-response.json result ($defs/SendMessageResponse)", a2a, "/$defs/SendMessageResponse", auth_required["result"])
    check("agent-card.json ($defs/AgentCard)", a2a, "/$defs/AgentCard", card)

print(f"\n{'All checks passed.' if not failures else f'{len(failures)} check(s) failed: {failures}'}")
sys.exit(1 if failures else 0)