"""ACT Appendix B test vectors.

Generates and validates all 15 test vectors from Appendix B of
draft-nennemann-act-00. Each vector includes description, input
parameters, and expected output or exception.

Reference: ACT Appendix B (Test Vectors).
"""

from __future__ import annotations

import time
import uuid
from dataclasses import dataclass, field
from typing import Any

from .crypto import (
    ACTKeyResolver,
    KeyRegistry,
    PrivateKey,
    PublicKey,
    b64url_sha256,
    compute_sha256,
    generate_ed25519_keypair,
    sign as crypto_sign,
    verify as crypto_verify,
)
from .dag import validate_dag
from .delegation import create_delegated_mandate, verify_capability_subset
from .errors import (
    ACTAudienceMismatchError,
    ACTCapabilityError,
    ACTDAGError,
    ACTDelegationError,
    ACTExpiredError,
    ACTPrivilegeEscalationError,
    ACTSignatureError,
    ACTValidationError,
)
from .ledger import ACTLedger
from .lifecycle import transition_to_record
from .token import (
    ACTMandate,
    ACTRecord,
    Capability,
    Delegation,
    DelegationEntry,
    ErrorClaim,
    Oversight,
    TaskClaim,
    _b64url_encode,
    decode_jws,
    encode_jws,
)
from .verify import ACTVerifier


@dataclass
class TestVector:
    """A single test vector."""

    id: str
    description: str
    valid: bool
    expected_exception: type[Exception] | None = None
    compact: str = ""
    record: ACTMandate | ACTRecord | None = None


def generate_vectors() -> tuple[list[TestVector], dict[str, Any]]:
    """Generate all Appendix B test vectors.

    Returns a list of TestVector objects and a context dict containing
    keys and other state needed for validation.

    Reference: ACT Appendix B.
    """
    # Fixed timestamp for deterministic vectors
    base_time = 1772064000

    # Generate key pairs for test agents
    iss_priv, iss_pub = generate_ed25519_keypair()
    sub_priv, sub_pub = generate_ed25519_keypair()
    agent_c_priv, agent_c_pub = generate_ed25519_keypair()

    # Fixed JTIs for cross-referencing
    jti_b1 = "550e8400-e29b-41d4-a716-446655440001"
    jti_b2 = "550e8400-e29b-41d4-a716-446655440002"
    jti_b3_parent1 = "550e8400-e29b-41d4-a716-446655440003"
    jti_b3_parent2 = "550e8400-e29b-41d4-a716-446655440004"
    jti_b3 = "550e8400-e29b-41d4-a716-446655440005"
    jti_b4 = "550e8400-e29b-41d4-a716-446655440006"
    jti_b5 = "550e8400-e29b-41d4-a716-446655440007"
    wid = "a0b1c2d3-e4f5-6789-abcd-ef0123456789"

    # Key registry
    registry = KeyRegistry()
    registry.register("iss-key", iss_pub)
    registry.register("sub-key", sub_pub)
    registry.register("agent-c-key", agent_c_pub)

    resolver = ACTKeyResolver(registry=registry)

    vectors: list[TestVector] = []
    compacts: dict[str, str] = {}  # jti → compact for delegation refs

    # --- B.1: Phase 1 — Root mandate, Tier 1, Ed25519, no delegation ---
    mandate_b1 = ACTMandate(
        alg="EdDSA",
        kid="iss-key",
        iss="agent-issuer",
        sub="agent-subject",
        aud=["agent-subject", "https://ledger.example.com"],
        iat=base_time,
        exp=base_time + 900,
        jti=jti_b1,
        wid=wid,
        task=TaskClaim(
            purpose="validate_data",
            data_sensitivity="restricted",
        ),
        cap=[
            Capability(action="read.data", constraints={"max_records": 10}),
            Capability(action="write.result"),
        ],
        delegation=Delegation(depth=0, max_depth=2, chain=[]),
    )
    mandate_b1.validate()
    sig_b1 = crypto_sign(iss_priv, mandate_b1.signing_input())
    compact_b1 = encode_jws(mandate_b1, sig_b1)
    compacts[jti_b1] = compact_b1

    vectors.append(TestVector(
        id="B.1",
        description="Phase 1 ACT — root mandate, Tier 1 (Ed25519), no delegation",
        valid=True,
        compact=compact_b1,
        record=mandate_b1,
    ))

    # --- B.2: Phase 2 — Completed execution from B.1 ---
    record_b2, compact_b2 = transition_to_record(
        mandate_b1,
        sub_kid="sub-key",
        sub_private_key=sub_priv,
        exec_act="read.data",
        par=[],
        exec_ts=base_time + 300,
        status="completed",
        inp_hash=b64url_sha256(b"test input data"),
        out_hash=b64url_sha256(b"test output data"),
    )
    compacts[jti_b2] = compact_b2

    vectors.append(TestVector(
        id="B.2",
        description="Phase 2 ACT — completed execution, transition from B.1 mandate",
        valid=True,
        compact=compact_b2,
        record=record_b2,
    ))

    # --- B.3: Phase 2 — Fan-in, two parent jti values ---
    # Create two parent records first
    parent1_mandate = ACTMandate(
        alg="EdDSA", kid="iss-key",
        iss="agent-issuer", sub="agent-subject",
        aud="agent-subject",
        iat=base_time, exp=base_time + 900,
        jti=jti_b3_parent1, wid=wid,
        task=TaskClaim(purpose="branch_a"),
        cap=[Capability(action="compute.result")],
        delegation=Delegation(depth=0, max_depth=1, chain=[]),
    )
    sig_p1 = crypto_sign(iss_priv, parent1_mandate.signing_input())
    compact_p1 = encode_jws(parent1_mandate, sig_p1)

    parent1_record, parent1_compact = transition_to_record(
        parent1_mandate, sub_kid="sub-key", sub_private_key=sub_priv,
        exec_act="compute.result", par=[], exec_ts=base_time + 100,
        status="completed",
    )

    parent2_mandate = ACTMandate(
        alg="EdDSA", kid="iss-key",
        iss="agent-issuer", sub="agent-subject",
        aud="agent-subject",
        iat=base_time, exp=base_time + 900,
        jti=jti_b3_parent2, wid=wid,
        task=TaskClaim(purpose="branch_b"),
        cap=[Capability(action="compute.result")],
        delegation=Delegation(depth=0, max_depth=1, chain=[]),
    )
    sig_p2 = crypto_sign(iss_priv, parent2_mandate.signing_input())
    compact_p2 = encode_jws(parent2_mandate, sig_p2)

    parent2_record, parent2_compact = transition_to_record(
        parent2_mandate, sub_kid="sub-key", sub_private_key=sub_priv,
        exec_act="compute.result", par=[], exec_ts=base_time + 150,
        status="completed",
    )

    # Fan-in record depends on both parents
    fanin_mandate = ACTMandate(
        alg="EdDSA", kid="iss-key",
        iss="agent-issuer", sub="agent-subject",
        aud="agent-subject",
        iat=base_time, exp=base_time + 900,
        jti=jti_b3, wid=wid,
        task=TaskClaim(purpose="merge_results"),
        cap=[Capability(action="compute.result")],
        delegation=Delegation(depth=0, max_depth=1, chain=[]),
    )
    sig_fi = crypto_sign(iss_priv, fanin_mandate.signing_input())

    fanin_record, fanin_compact = transition_to_record(
        fanin_mandate, sub_kid="sub-key", sub_private_key=sub_priv,
        exec_act="compute.result",
        par=[jti_b3_parent1, jti_b3_parent2],
        exec_ts=base_time + 200,
        status="completed",
    )

    vectors.append(TestVector(
        id="B.3",
        description="Phase 2 ACT — fan-in, two parent jti values from parallel branches",
        valid=True,
        compact=fanin_compact,
        record=fanin_record,
    ))

    # --- B.4: Phase 1 — Delegated mandate (depth=1) ---
    delegated_b4, _ = create_delegated_mandate(
        parent_mandate=mandate_b1,
        parent_compact=compact_b1,
        delegator_private_key=iss_priv,
        sub="agent-c",
        kid="iss-key",
        iss="agent-issuer",
        aud="agent-c",
        iat=base_time + 10,
        exp=base_time + 600,
        jti=jti_b4,
        cap=[Capability(action="read.data", constraints={"max_records": 5})],
        task=TaskClaim(purpose="delegated_read"),
    )
    sig_b4 = crypto_sign(iss_priv, delegated_b4.signing_input())
    compact_b4 = encode_jws(delegated_b4, sig_b4)
    compacts[jti_b4] = compact_b4

    vectors.append(TestVector(
        id="B.4",
        description="Phase 1 ACT — delegated mandate (depth=1), chain entry with sig",
        valid=True,
        compact=compact_b4,
        record=delegated_b4,
    ))

    # --- B.5: Phase 2 — Delegated execution record ---
    record_b5, compact_b5 = transition_to_record(
        delegated_b4,
        sub_kid="agent-c-key",
        sub_private_key=agent_c_priv,
        exec_act="read.data",
        par=[],
        exec_ts=base_time + 350,
        status="completed",
    )

    vectors.append(TestVector(
        id="B.5",
        description="Phase 2 ACT — delegated execution record",
        valid=True,
        compact=compact_b5,
        record=record_b5,
    ))

    # --- B.6: del.depth > del.max_depth → ACTDelegationError ---
    bad_depth_mandate = ACTMandate(
        alg="EdDSA", kid="iss-key",
        iss="agent-issuer", sub="agent-subject",
        aud="agent-subject",
        iat=base_time, exp=base_time + 900,
        jti=str(uuid.uuid4()),
        task=TaskClaim(purpose="bad_depth"),
        cap=[Capability(action="read.data")],
        delegation=Delegation(depth=3, max_depth=2, chain=[
            DelegationEntry(delegator="a", jti="j1", sig="sig1"),
            DelegationEntry(delegator="b", jti="j2", sig="sig2"),
            DelegationEntry(delegator="c", jti="j3", sig="sig3"),
        ]),
    )
    sig_b6 = crypto_sign(iss_priv, bad_depth_mandate.signing_input())
    compact_b6 = encode_jws(bad_depth_mandate, sig_b6)

    vectors.append(TestVector(
        id="B.6",
        description="del.depth > del.max_depth → ACTDelegationError",
        valid=False,
        expected_exception=ACTDelegationError,
        compact=compact_b6,
    ))

    # --- B.7: cap escalation in delegated ACT → ACTPrivilegeEscalationError ---
    vectors.append(TestVector(
        id="B.7",
        description="cap escalation in delegated ACT → ACTPrivilegeEscalationError",
        valid=False,
        expected_exception=ACTPrivilegeEscalationError,
    ))

    # --- B.8: exec_act not in cap → ACTCapabilityError ---
    bad_exec_mandate = ACTMandate(
        alg="EdDSA", kid="iss-key",
        iss="agent-issuer", sub="agent-subject",
        aud="agent-subject",
        iat=base_time, exp=base_time + 900,
        jti=str(uuid.uuid4()),
        task=TaskClaim(purpose="bad_exec"),
        cap=[Capability(action="read.data")],
        delegation=Delegation(depth=0, max_depth=1, chain=[]),
    )
    sig_b8m = crypto_sign(iss_priv, bad_exec_mandate.signing_input())

    # Manually construct Phase 2 with wrong exec_act
    bad_exec_record = ACTRecord(
        alg="EdDSA", kid="sub-key",
        iss="agent-issuer", sub="agent-subject",
        aud="agent-subject",
        iat=base_time, exp=base_time + 900,
        jti=bad_exec_mandate.jti,
        task=TaskClaim(purpose="bad_exec"),
        cap=[Capability(action="read.data")],
        exec_act="delete.everything",
        par=[], exec_ts=base_time + 100, status="completed",
    )
    sig_b8 = crypto_sign(sub_priv, bad_exec_record.signing_input())
    compact_b8 = encode_jws(bad_exec_record, sig_b8)

    vectors.append(TestVector(
        id="B.8",
        description="exec_act not in cap → ACTCapabilityError",
        valid=False,
        expected_exception=ACTCapabilityError,
        compact=compact_b8,
    ))

    # --- B.9: DAG cycle (par references own jti) → ACTDAGError ---
    cycle_jti = str(uuid.uuid4())
    cycle_record = ACTRecord(
        alg="EdDSA", kid="sub-key",
        iss="agent-issuer", sub="agent-subject",
        aud="agent-subject",
        iat=base_time, exp=base_time + 900,
        jti=cycle_jti,
        task=TaskClaim(purpose="cycle_test"),
        cap=[Capability(action="read.data")],
        exec_act="read.data",
        par=[cycle_jti],
        exec_ts=base_time + 100, status="completed",
    )
    sig_b9 = crypto_sign(sub_priv, cycle_record.signing_input())
    compact_b9 = encode_jws(cycle_record, sig_b9)

    vectors.append(TestVector(
        id="B.9",
        description="DAG cycle (par references own jti) → ACTDAGError",
        valid=False,
        expected_exception=ACTDAGError,
        compact=compact_b9,
    ))

    # --- B.10: Missing parent jti in DAG → ACTDAGError ---
    missing_parent_record = ACTRecord(
        alg="EdDSA", kid="sub-key",
        iss="agent-issuer", sub="agent-subject",
        aud="agent-subject",
        iat=base_time, exp=base_time + 900,
        jti=str(uuid.uuid4()),
        task=TaskClaim(purpose="missing_parent"),
        cap=[Capability(action="read.data")],
        exec_act="read.data",
        par=["nonexistent-parent-jti"],
        exec_ts=base_time + 100, status="completed",
    )
    sig_b10 = crypto_sign(sub_priv, missing_parent_record.signing_input())
    compact_b10 = encode_jws(missing_parent_record, sig_b10)

    vectors.append(TestVector(
        id="B.10",
        description="Missing parent jti in DAG → ACTDAGError",
        valid=False,
        expected_exception=ACTDAGError,
        compact=compact_b10,
    ))

    # --- B.11: Tampered payload (bit flip) → ACTSignatureError ---
    # Take a valid compact and flip a byte in the payload
    parts = compact_b1.split(".")
    payload_bytes = bytearray(parts[1].encode("ascii"))
    # Flip a character in the payload
    flip_idx = len(payload_bytes) // 2
    payload_bytes[flip_idx] = (payload_bytes[flip_idx] + 1) % 128
    if payload_bytes[flip_idx] == 0:
        payload_bytes[flip_idx] = 65  # 'A'
    tampered_compact = f"{parts[0]}.{payload_bytes.decode('ascii')}.{parts[2]}"

    vectors.append(TestVector(
        id="B.11",
        description="Tampered payload (bit flip in claims) → ACTSignatureError",
        valid=False,
        expected_exception=ACTSignatureError,
        compact=tampered_compact,
    ))

    # --- B.12: Expired token → ACTExpiredError ---
    expired_mandate = ACTMandate(
        alg="EdDSA", kid="iss-key",
        iss="agent-issuer", sub="agent-subject",
        aud="agent-subject",
        iat=base_time - 3600,
        exp=base_time - 2700,  # expired 45 minutes ago
        jti=str(uuid.uuid4()),
        task=TaskClaim(purpose="expired_test"),
        cap=[Capability(action="read.data")],
    )
    sig_b12 = crypto_sign(iss_priv, expired_mandate.signing_input())
    compact_b12 = encode_jws(expired_mandate, sig_b12)

    vectors.append(TestVector(
        id="B.12",
        description="Expired token → ACTExpiredError",
        valid=False,
        expected_exception=ACTExpiredError,
        compact=compact_b12,
    ))

    # --- B.13: Wrong audience → ACTAudienceMismatchError ---
    wrong_aud_mandate = ACTMandate(
        alg="EdDSA", kid="iss-key",
        iss="agent-issuer", sub="wrong-agent",
        aud="wrong-agent",
        iat=base_time, exp=base_time + 900,
        jti=str(uuid.uuid4()),
        task=TaskClaim(purpose="wrong_aud_test"),
        cap=[Capability(action="read.data")],
    )
    sig_b13 = crypto_sign(iss_priv, wrong_aud_mandate.signing_input())
    compact_b13 = encode_jws(wrong_aud_mandate, sig_b13)

    vectors.append(TestVector(
        id="B.13",
        description="Wrong audience → ACTAudienceMismatchError",
        valid=False,
        expected_exception=ACTAudienceMismatchError,
        compact=compact_b13,
    ))

    # --- B.14: Phase 2 re-signed by iss key instead of sub → ACTSignatureError ---
    record_b14 = ACTRecord.from_mandate(
        mandate_b1,
        kid="sub-key",  # claims to be sub's key
        exec_act="read.data",
        par=[], exec_ts=base_time + 300, status="completed",
    )
    # But signed with ISS's private key (wrong signer)
    sig_b14 = crypto_sign(iss_priv, record_b14.signing_input())
    compact_b14 = encode_jws(record_b14, sig_b14)

    vectors.append(TestVector(
        id="B.14",
        description="Phase 2 re-signed by iss key instead of sub → ACTSignatureError",
        valid=False,
        expected_exception=ACTSignatureError,
        compact=compact_b14,
    ))

    # --- B.15: Algorithm "none" → ACTValidationError ---
    # Manually construct a JWS with alg: none
    import json
    import base64

    none_header = base64.urlsafe_b64encode(
        json.dumps({"alg": "none", "typ": "act+jwt", "kid": "k"}, separators=(",", ":")).encode()
    ).rstrip(b"=").decode()
    none_payload = base64.urlsafe_b64encode(
        json.dumps({"iss": "a", "sub": "b"}, separators=(",", ":")).encode()
    ).rstrip(b"=").decode()
    compact_b15 = f"{none_header}.{none_payload}."

    vectors.append(TestVector(
        id="B.15",
        description='Algorithm "none" → ACTValidationError',
        valid=False,
        expected_exception=ACTValidationError,
        compact=compact_b15,
    ))

    context = {
        "iss_priv": iss_priv,
        "iss_pub": iss_pub,
        "sub_priv": sub_priv,
        "sub_pub": sub_pub,
        "agent_c_priv": agent_c_priv,
        "agent_c_pub": agent_c_pub,
        "registry": registry,
        "resolver": resolver,
        "base_time": base_time,
        "compacts": compacts,
        "parent1_record": parent1_record,
        "parent2_record": parent2_record,
        "mandate_b1": mandate_b1,
    }

    return vectors, context


def validate_vectors() -> bool:
    """Run all test vectors and validate results.

    Returns True if all vectors pass.

    Reference: ACT Appendix B.
    """
    vectors, ctx = generate_vectors()
    resolver = ctx["resolver"]
    base_time = ctx["base_time"]

    verifier = ACTVerifier(
        key_resolver=resolver,
        verifier_id="agent-subject",
        trusted_issuers={"agent-issuer"},
    )

    passed = 0
    failed = 0

    for v in vectors:
        try:
            if v.id == "B.7":
                # Special case: test cap escalation during delegation creation
                try:
                    from .delegation import verify_capability_subset
                    verify_capability_subset(
                        [Capability(action="read.data", constraints={"max_records": 10})],
                        [Capability(action="read.data", constraints={"max_records": 100})],
                    )
                    print(f"  FAIL {v.id}: Expected {v.expected_exception.__name__}")
                    failed += 1
                except ACTPrivilegeEscalationError:
                    print(f"  PASS {v.id}: {v.description}")
                    passed += 1
                continue

            if v.valid:
                # Valid vectors: should parse and verify without error
                header, claims, sig, si = decode_jws(v.compact)
                kid = header["kid"]
                pub = resolver.resolve(kid, header=header)
                crypto_verify(pub, sig, si)
                print(f"  PASS {v.id}: {v.description}")
                passed += 1
            else:
                # Invalid vectors: should raise the expected exception
                try:
                    if v.expected_exception == ACTDelegationError:
                        header, claims, sig, si = decode_jws(v.compact)
                        kid = header["kid"]
                        pub = resolver.resolve(kid, header=header)
                        crypto_verify(pub, sig, si)
                        # Parse and check delegation
                        from .token import ACTMandate as _M
                        m = _M.from_claims(header, claims)
                        from .delegation import verify_delegation_chain
                        verify_delegation_chain(m, lambda d: resolver.resolve(d))
                        print(f"  FAIL {v.id}: Expected {v.expected_exception.__name__}")
                        failed += 1
                    elif v.expected_exception == ACTCapabilityError:
                        header, claims, sig, si = decode_jws(v.compact)
                        kid = header["kid"]
                        pub = resolver.resolve(kid, header=header)
                        crypto_verify(pub, sig, si)
                        r = ACTRecord.from_claims(header, claims)
                        cap_actions = {c.action for c in r.cap}
                        if r.exec_act not in cap_actions:
                            raise ACTCapabilityError("exec_act mismatch")
                        print(f"  FAIL {v.id}: Expected {v.expected_exception.__name__}")
                        failed += 1
                    elif v.expected_exception == ACTDAGError:
                        header, claims, sig, si = decode_jws(v.compact)
                        kid = header["kid"]
                        pub = resolver.resolve(kid, header=header)
                        crypto_verify(pub, sig, si)
                        r = ACTRecord.from_claims(header, claims)
                        ledger = ACTLedger()
                        validate_dag(r, ledger)
                        print(f"  FAIL {v.id}: Expected {v.expected_exception.__name__}")
                        failed += 1
                    elif v.expected_exception == ACTExpiredError:
                        verifier.verify_mandate(v.compact, check_sub=False)
                        print(f"  FAIL {v.id}: Expected {v.expected_exception.__name__}")
                        failed += 1
                    elif v.expected_exception == ACTAudienceMismatchError:
                        verifier.verify_mandate(
                            v.compact,
                            now=base_time + 100,
                            check_sub=False,
                        )
                        print(f"  FAIL {v.id}: Expected {v.expected_exception.__name__}")
                        failed += 1
                    elif v.expected_exception == ACTSignatureError:
                        header, claims, sig, si = decode_jws(v.compact)
                        kid = header["kid"]
                        pub = resolver.resolve(kid, header=header)
                        crypto_verify(pub, sig, si)
                        print(f"  FAIL {v.id}: Expected {v.expected_exception.__name__}")
                        failed += 1
                    elif v.expected_exception == ACTValidationError:
                        decode_jws(v.compact)
                        print(f"  FAIL {v.id}: Expected {v.expected_exception.__name__}")
                        failed += 1
                    else:
                        print(f"  SKIP {v.id}: Unknown expected exception type")
                        failed += 1
                except Exception as e:
                    if isinstance(e, v.expected_exception):
                        print(f"  PASS {v.id}: {v.description}")
                        passed += 1
                    else:
                        print(f"  FAIL {v.id}: Expected {v.expected_exception.__name__}, "
                              f"got {type(e).__name__}: {e}")
                        failed += 1

        except Exception as e:
            print(f"  FAIL {v.id}: Unexpected error: {type(e).__name__}: {e}")
            failed += 1

    print(f"\nResults: {passed} passed, {failed} failed out of {len(vectors)}")
    return failed == 0
