#!/usr/bin/env python3
"""Generate the GRIP/1.0-draft.2 canonical test vectors.

Builds the Appendix C refund chain end to end with real Ed25519 keys and
real digests, then derives the negative vectors that must FAIL specific
rules. Deterministic: keys are seeded, so digests are stable across runs.
"""
from __future__ import annotations

import hashlib
import json
import os
import shutil
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from gripcore import (  # noqa: E402
    canonicalize, digest_bytes, seal, sign_detached, b64u,
)
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey  # noqa: E402
import base64  # noqa: E402

ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "vectors"))
GRIP = "1.0-draft.2"

# Deterministic keys — seeded so the published vectors are reproducible.
IDENTITIES = {
    "dana":        "https://id.acme.com/people/dana#key1",
    "coordinator": "https://runtime.acme.com/issuers/coordinator#key3",
    "gate":        "https://runtime.acme.com/issuers/gate#key1",
    "bot":         "https://agents.acme.com/support-bot#key2",
    "settle":      "https://runtime.acme.com/issuers/settle#key4",
    "lee":         "https://id.acme.com/people/lee#key1",
    "mallory":     "https://evil.example.com/issuers/mallory#key1",
}
KEYS = {
    n: Ed25519PrivateKey.from_private_bytes(hashlib.sha256(f"grip-vector-seed:{n}".encode()).digest())
    for n in IDENTITIES
}

T0 = 1787000000


def bundle(extra_att=None) -> dict:
    return {
        "grip": GRIP,
        "keys": [
            {"kid": IDENTITIES[n],
             "public_key": base64.b64encode(
                 KEYS[n].public_key().public_bytes_raw()).decode(),
             "valid_from": T0 - 86400, "valid_until": T0 + 86400 * 365}
            for n in IDENTITIES
        ],
        "head_attestations": extra_att or [],
    }


SETTLEMENT_IAT = T0 + 3000  # inside the Grant window, so a post-settlement
                            # receipt can also be inside it (surgical vector)


def build_chain(bad_approval=False, over_budget=False, forged_grant=False,
                extra=None):
    """Appendix C, end to end. Options each introduce exactly one defect."""
    o = {}

    o["01-mandate"] = seal({
        "grip": GRIP, "type": "mandate",
        "iss": IDENTITIES["dana"],
        "goal": "Resolve open billing tickets, refunds under $50 each",
        "constraints": ["never contact the customer's bank",
                        "refunds only on tickets tagged verified"],
        "scope": [
            {"verb": "tool.call", "resource": "mcp://crm/*", "class": "observe"},
            {"verb": "tool.call", "resource": "mcp://crm/refund", "class": "settle"},
            {"verb": "model.call", "resource": "class:reasoning", "class": "egress"}],
        "issuers": [IDENTITIES["coordinator"]],
        "acceptance": [
            {"code": "tickets-closed", "check": "all selected tickets in state closed"},
            {"code": "refund-ledger", "check": "every refund posted with ticket reference",
             "independent": True}],
        "budget": {"usd_micro": 200000000, "tokens_in": 5000000,
                   "tokens_out": 1000000, "ms": 86400000},
        "iat": T0, "exp": T0 + 86400,
    }, KEYS["dana"], IDENTITIES["dana"])
    M = o["01-mandate"]["id"]

    o["02-grant"] = seal({
        "grip": GRIP, "type": "grant", "parents": [M],
        "iss": IDENTITIES["coordinator"], "actor": IDENTITIES["bot"],
        "allow": [
            {"verb": "tool.call", "resource": "mcp://crm/*", "class": "observe"},
            {"verb": "tool.call", "resource": "mcp://crm/refund", "class": "settle"},
            {"verb": "model.call", "resource": "class:reasoning", "class": "egress",
             "boundary": ["no-public-cloud"]}],
        "budget": {"usd_micro": 60000000, "tokens_in": 500000,
                   "tokens_out": 100000, "ms": 3600000},
        "iat": T0 + 100, "exp": T0 + 3700,
    }, KEYS["coordinator"], IDENTITIES["coordinator"])
    G = o["02-grant"]["id"]

    crm_history = b'{"ticket":8812,"status":"verified","amount_usd":42.00}'
    runbook = b"# billing runbook\n1. verify tag\n2. refund <= $50\n"
    o["03-context"] = seal({
        "grip": GRIP, "type": "context", "parents": [M],
        "iss": "https://runtime.acme.com/issuers/compiler#key2"
               if False else IDENTITIES["gate"],
        "purpose": "select next act for ticket batch 7",
        "budget_tokens": 4000, "used_tokens": 3112,
        "items": [
            {"ref": "mandate", "digest": M, "trust": "principal"},
            {"ref": "crm-history-8812", "digest": digest_bytes(crm_history), "trust": "external"},
            {"ref": "skill:billing-runbook", "digest": digest_bytes(runbook), "trust": "governed"}],
        "iat": T0 + 200,
    }, KEYS["gate"], IDENTITIES["gate"])
    C = o["03-context"]["id"]

    def actor_sig(act: dict) -> str:
        body = {k: v for k, v in act.items() if k != "actor_sig"}
        return sign_detached(G.encode() + canonicalize(body), KEYS["bot"], IDENTITIES["bot"])

    # -- model.call: reason about the batch (gate + outcome) ---------------
    model_input = b'{"prompt":"select next act for ticket batch 7"}'
    act_model = {"verb": "model.call", "resource": "class:reasoning",
                 "input_digest": digest_bytes(model_input)}
    act_model["actor_sig"] = actor_sig(act_model)
    o["04-receipt-model-gate"] = seal({
        "grip": GRIP, "type": "receipt", "parents": [M, G],
        "iss": IDENTITIES["gate"], "actor": IDENTITIES["bot"],
        "act": act_model, "context": C,
        "decision": {"result": "allowed"},
        "cost": {"usd_micro": 500000, "tokens_in": 413000, "tokens_out": 21100, "ms": 1200},
        "iat": T0 + 210,
    }, KEYS["gate"], IDENTITIES["gate"])
    RM = o["04-receipt-model-gate"]["id"]

    o["05-receipt-model-outcome"] = seal({
        "grip": GRIP, "type": "receipt", "parents": [M, G], "prev": RM,
        "iss": IDENTITIES["gate"], "actor": IDENTITIES["bot"],
        "act": act_model, "context": C, "gate": RM,
        "decision": {"result": "allowed"},
        "outcome": {"status": "ok", "output_digest": digest_bytes(b'{"next":"refund"}'),
                    "resolved": {"provider": "acme-inference", "model": "reasoner-v2",
                                 "boundary": ["no-public-cloud"]}},
        "cost": {"usd_micro": 480000, "tokens_in": 413000, "tokens_out": 21100, "ms": 1180},
        "iat": T0 + 220,
    }, KEYS["gate"], IDENTITIES["gate"])
    RMO = o["05-receipt-model-outcome"]["id"]

    # -- the refund act, denied for want of a Stop ------------------------
    refund_input = b'{"ticket":8812,"amount_usd_micro":42000000,"reason":"verified"}'
    ACTD = digest_bytes(refund_input)
    act_refund = {"verb": "tool.call", "resource": "mcp://crm/refund", "input_digest": ACTD}
    act_refund["actor_sig"] = actor_sig(act_refund)

    o["06-receipt-denied"] = seal({
        "grip": GRIP, "type": "receipt", "parents": [M, G], "prev": RMO,
        "iss": IDENTITIES["gate"], "actor": IDENTITIES["bot"],
        "act": act_refund, "context": C,
        "decision": {"result": "denied", "reason": "GRIP_DENY_STOP_REQUIRED"},
        "cost": {"usd_micro": 0, "tokens_in": 0, "tokens_out": 0, "ms": 4},
        "iat": T0 + 230,
    }, KEYS["gate"], IDENTITIES["gate"])
    RD = o["06-receipt-denied"]["id"]

    o["07-stop"] = seal({
        "grip": GRIP, "type": "stop", "parents": [M, G],
        "iss": IDENTITIES["gate"], "act_digest": ACTD,
        "reason": "settle-class act: refund $42 on ticket 8812",
        "approver": IDENTITIES["dana"], "mode": "signature",
        "iat": T0 + 230, "exp": T0 + 4000,
    }, KEYS["gate"], IDENTITIES["gate"])
    S = o["07-stop"]["id"]

    # bad_approval: signed by mallory but presented under dana's kid.
    _ak = KEYS["mallory"] if bad_approval else KEYS["dana"]
    approval = sign_detached((S + ACTD).encode(), _ak, IDENTITIES["dana"])
    o["08-receipt-approve"] = seal({
        "grip": GRIP, "type": "receipt", "parents": [M, S], "prev": RD,
        "iss": IDENTITIES["gate"], "actor": IDENTITIES["dana"],
        "act": {"verb": "stop.approve", "resource": S},
        "approval": approval,
        "decision": {"result": "allowed"},
        "cost": {"usd_micro": 0, "tokens_in": 0, "tokens_out": 0, "ms": 2},
        "iat": T0 + 900,
    }, KEYS["gate"], IDENTITIES["gate"])
    RA = o["08-receipt-approve"]["id"]

    o["09-receipt-refund-gate"] = seal({
        "grip": GRIP, "type": "receipt", "parents": [M, G], "prev": RA,
        "iss": IDENTITIES["gate"], "actor": IDENTITIES["bot"],
        "act": act_refund, "context": C,
        "decision": {"result": "allowed", "stop": S},
        "cost": {"usd_micro": 42000000, "tokens_in": 0, "tokens_out": 0, "ms": 10},
        "iat": T0 + 910,
    }, KEYS["gate"], IDENTITIES["gate"])
    RG = o["09-receipt-refund-gate"]["id"]

    o["10-receipt-refund-outcome"] = seal({
        "grip": GRIP, "type": "receipt", "parents": [M, G], "prev": RG,
        "iss": IDENTITIES["gate"], "actor": IDENTITIES["bot"],
        "act": act_refund, "context": C, "gate": RG,
        "decision": {"result": "allowed", "stop": S},
        "outcome": {"status": "ok", "effects": ["crm:ticket/8812#refund"],
                    "output_digest": digest_bytes(b'{"refund_id":"rf_8812"}'),
                    "resolved": None},
        "cost": {"usd_micro": 90000000 if over_budget else 42000000,
                 "tokens_in": 0, "tokens_out": 0, "ms": 2398804},
        "iat": T0 + 920,
    }, KEYS["gate"], IDENTITIES["gate"])
    RO = o["10-receipt-refund-outcome"]["id"]

    if forged_grant:
        o["12-grant-forged"] = seal({
            "grip": GRIP, "type": "grant", "parents": [M],
            "iss": IDENTITIES["mallory"], "actor": IDENTITIES["bot"],
            "allow": [{"verb": "tool.call", "resource": "mcp://crm/*", "class": "settle"}],
            "budget": {"usd_micro": 999000000, "tokens_in": 1, "tokens_out": 1, "ms": 1},
            "iat": T0 + 100, "exp": T0 + 3700},
            KEYS["mallory"], IDENTITIES["mallory"])

    prev_id = RO
    for name, spec in (extra or []):
        body = {"grip": GRIP, "type": "receipt", "parents": [M, G], "prev": prev_id,
                "iss": IDENTITIES["gate"], "actor": IDENTITIES["bot"],
                "cost": {"usd_micro": 0, "tokens_in": 0, "tokens_out": 0, "ms": 0}}
        body.update(spec)
        if spec.get("_refund_act"):
            body.pop("_refund_act")
            body["act"] = act_refund
            body["context"] = C
        o[name] = seal(body, KEYS["gate"], IDENTITIES["gate"])
        prev_id = o[name]["id"]

    # -- settlement --------------------------------------------------------
    totals = {"usd_micro": 0, "tokens_in": 0, "tokens_out": 0, "ms": 0}
    superseded = {RM, RG}
    for k, v in o.items():
        if v.get("type") == "receipt" and v["id"] not in superseded:
            for ax, n in v["cost"].items():
                totals[ax] += n

    c1 = {"code": "tickets-closed", "verdict": "accepted", "evidence": [RO],
          "validator": IDENTITIES["settle"]}
    c2body = {"code": "refund-ledger", "verdict": "accepted", "evidence": [RO, RG],
              "validator": IDENTITIES["lee"]}
    c2 = dict(c2body)
    c2["validator_sig"] = sign_detached(canonicalize(c2body), KEYS["lee"], IDENTITIES["lee"])

    o["11-settlement"] = seal({
        "grip": GRIP, "type": "settlement", "parents": [M],
        "iss": IDENTITIES["settle"], "state": "settled",
        "criteria": [c1, c2], "totals": totals,
        "predicate": "identity-inequality",
        "iat": SETTLEMENT_IAT,
    }, KEYS["settle"], IDENTITIES["settle"])

    att = {"mandate": M, "head": RO, "receipt_count": 7, "totals": totals, "iat": SETTLEMENT_IAT + 100}
    return o, att, {"M": M, "G": G, "C": C, "S": S, "ACTD": ACTD, "RD": RD,
                    "RA": RA, "RG": RG, "RO": RO, "SET": o["11-settlement"]["id"],
                    "totals": totals}


def write(dirname: str, objects: dict, att=None, expected=None):
    d = os.path.join(ROOT, dirname)
    if os.path.isdir(d):
        shutil.rmtree(d)
    os.makedirs(os.path.join(d, "objects"))
    for name, obj in objects.items():
        json.dump(obj, open(os.path.join(d, "objects", f"{name}.json"), "w"),
                  indent=2, sort_keys=True)
    json.dump(bundle([att] if att else []),
              open(os.path.join(d, "trust-bundle.json"), "w"), indent=2)
    if expected:
        json.dump(expected, open(os.path.join(d, "expected.json"), "w"), indent=2)
    return d


def main():
    happy, att, ids = build_chain()
    write("refund-happy-path", happy, att, {
        "description": "Appendix C, end to end. All eight rules pass.",
        "expect_clean": True,
        "expect_verdicts": {f"rule{n}": "pass" for n in range(1, 9)},
    })

    # N1 — Stop replay. Zero-cost so budget and totals stay intact and ONLY
    # rule 5 can fail: the replay is the whole defect.
    n1, _, _ = build_chain(extra=[("12-receipt-stop-replay", {
        "_refund_act": True, "decision": {"result": "allowed", "stop": ids["S"]},
        "iat": T0 + 930})])
    write("negative-stop-replay", n1, None, {
        "description": "A second allowed settle-class act references the already-consumed Stop.",
        "expect_clean": False, "must_fail_rules": [5]})

    # N2 — Receipt dated after the terminal Settlement, but still inside the
    # Grant window and zero-cost, so only rule 6 can fail.
    n2, _, _ = build_chain(extra=[("12-receipt-post-settlement", {
        "act": {"verb": "tool.call", "resource": "mcp://crm/list"},
        "decision": {"result": "allowed"}, "iat": SETTLEMENT_IAT + 500})])
    write("negative-post-settlement", n2, None, {
        "description": "A Receipt dated after the Mandate's terminal Settlement.",
        "expect_clean": False, "must_fail_rules": [6]})

    # N3 — root Grant from an issuer the Mandate does not list.
    n3, _, _ = build_chain(forged_grant=True)
    write("negative-forged-grant", n3, None, {
        "description": "A root Grant signed by an issuer the Mandate does not list, "
                       "whose budget also exceeds the Mandate's.",
        "expect_clean": False, "must_fail_rules": [3]})

    # N4 — the approval is not the approver's signature. Rebuilt as a whole
    # chain so the prev-links stay intact and only rule 5 can fail.
    n4, _, _ = build_chain(bad_approval=True)
    write("negative-fabricated-approval", n4, None, {
        "description": "The stop.approve carries a signature that is not the approver's.",
        "expect_clean": False, "must_fail_rules": [5]})

    # N5 — object mutated after signing.
    n5 = {k: dict(v) for k, v in happy.items()}
    n5["01-mandate"]["goal"] = "Resolve open billing tickets, refunds under $5000 each"
    write("negative-tampered-chain", n5, None, {
        "description": "The Mandate's goal was edited after signing; id and sig no longer bind.",
        "expect_clean": False, "must_fail_rules": [1, 2]})

    # N6 — over-budget act, with a SELF-CONSISTENT settlement. The attacker
    # also restated the totals; rule 4 must still catch the breach.
    n6, _, _ = build_chain(over_budget=True)
    write("negative-over-budget", n6, None, {
        "description": "An outcome cost that breaches the Grant's usd_micro ceiling, "
                       "with a settlement whose totals were restated to match.",
        "expect_clean": False, "must_fail_rules": [4]})

    print("VECTORS WRITTEN\n")
    print("Canonical digests (Appendix C):")
    for k in ("M", "G", "C", "S", "RD", "RA", "RG", "RO", "SET"):
        print(f"  {k:4s} {ids[k]}")
    print(f"  ACTD {ids['ACTD']}")
    print(f"\nRecomputed totals: {ids['totals']}")
    json.dump(ids, open(os.path.join(ROOT, "canonical-ids.json"), "w"), indent=2)


if __name__ == "__main__":
    main()