#!/usr/bin/env python3
"""Two-agent ECT workflow demo: Agent A creates root ECT, Agent B verifies and creates child."""

import time

from ect import (
    Payload,
    create,
    generate_key,
    CreateOptions,
    verify,
    VerifyOptions,
    MemoryLedger,
)

def main():
    ledger = MemoryLedger()
    now = int(time.time())

    key_a = generate_key()
    agent_a = "spiffe://example.com/agent/spec-reviewer"
    agent_b = "spiffe://example.com/agent/implementer"
    kid_a = "agent-a-key"

    # 1) Agent A creates root ECT (task id = jti per spec)
    root_jti = "550e8400-e29b-41d4-a716-446655440001"
    payload_a = Payload(
        iss=agent_a,
        aud=[agent_b],
        iat=now,
        exp=now + 600,
        jti=root_jti,
        wid="wf-demo-001",
        exec_act="review_requirements_spec",
        pred=[],
        ext={
            "pol": "spec_review_policy_v2",
            "pol_decision": "approved",
        },
    )
    ect_a = create(payload_a, key_a, CreateOptions(key_id=kid_a))
    print("Agent A created root ECT (jti=550e8400-..., review_requirements_spec)")

    # 2) Agent B verifies
    def resolve_key(kid):
        if kid == kid_a:
            return key_a.public_key()
        return None

    opts = VerifyOptions(
        verifier_id=agent_b,
        resolve_key=resolve_key,
        store=ledger,
        now=now,
    )
    parsed = verify(ect_a, opts)
    ledger.append(ect_a, parsed.payload)
    print("Agent B verified root ECT and appended to ledger")

    # 3) Agent B creates child ECT (pred contains predecessor jti values per spec)
    key_b = generate_key()
    kid_b = "agent-b-key"
    child_jti = "550e8400-e29b-41d4-a716-446655440002"
    payload_b = Payload(
        iss=agent_b,
        aud=["spiffe://example.com/system/ledger"],
        iat=now + 1,
        exp=now + 600,
        jti=child_jti,
        wid="wf-demo-001",
        exec_act="implement_module",
        pred=[root_jti],
        ext={
            "pol": "coding_standards_v3",
            "pol_decision": "approved",
        },
    )
    ect_b = create(payload_b, key_b, CreateOptions(key_id=kid_b))
    print("Agent B created child ECT (jti=550e8400-...002, implement_module, pred=[predecessor jti])")

    # 4) Verify child ECT with DAG
    def resolver_b(kid):
        if kid == kid_b:
            return key_b.public_key()
        if kid == kid_a:
            return key_a.public_key()
        return None

    opts_b = VerifyOptions(
        verifier_id="spiffe://example.com/system/ledger",
        resolve_key=resolver_b,
        store=ledger,
        now=now + 2,
    )
    parsed_b = verify(ect_b, opts_b)
    ledger.append(ect_b, parsed_b.payload)
    print("Verified child ECT with DAG validation and appended to ledger")
    print(f"Ledger entries: {parsed.payload.jti} ({parsed.payload.exec_act}), {parsed_b.payload.jti} ({parsed_b.payload.exec_act})")


if __name__ == "__main__":
    main()