"""
Bartholomew Agent-to-Agent (A2A) Cryptographic Telemetry Protocol (v2.3)
========================================================================
Enables trustless, verifiable multi-agent swarms across process & cloud boundaries.

Architecture:
  [Agent A (Planner)] ────(Signed A2A Envelope)────> [Agent B (Executor)]
           │                                                  │
  [Signs Task Payload]                              [Verifies Ed25519 Seal]
  [BTP Attestation]                                 [Executes with Invariant Guard]

Guarantees:
  1. Non-Repudiation: Every inter-agent task delegation is signed with Ed25519.
  2. Scope Enforcement: Agent A cannot delegate capabilities beyond its own policy bounds.
  3. Replay Protection: Time-bound nonces prevent replayed agent instructions.
"""

import sys
import os
import time
import json
import secrets
from typing import Dict, Any, Tuple, Optional, List

sys.path.insert(0, os.path.abspath("."))
from src.trust_protocol import BartholomewTrustAuthority, IndependentTrustVerifier, rfc8785_canonicalize
from src.polyglot_ast_validator import PolyglotASTValidator


class AgentToAgentProtocol:
    """
    Cryptographic envelope generator and validator for multi-agent swarms.
    BTP v3.1: Enforces Sovereign Digital Passports & Capability Least-Privilege.
    """

    @classmethod
    def create_signed_handoff(cls, 
                              sender_authority: BartholomewTrustAuthority,
                              originating_agent: str,
                              target_agent: str,
                              task_action: str,
                              task_payload: Dict[str, Any],
                              capability_scope: Optional[List[str]] = None,
                              sender_passport: Optional[Any] = None,
                              ttl_seconds: int = 60) -> Dict[str, Any]:
        """
        Agent A creates an RFC 8785 canonical signed handoff envelope for Agent B.
        If sender_passport is provided, verifies authorization and attaches signed passport.
        """
        now = time.time()
        nonce = secrets.token_hex(16)
        scope = capability_scope or ["READ_ONLY", "AST_STRICT", "NO_RAW_SHELL"]

        # Validate task payload with AST
        if "command" in task_payload:
            is_safe, msg, _ = PolyglotASTValidator.validate_code(task_payload["command"])
            if not is_safe:
                raise ValueError(f"Cannot delegate unsafe command to Agent '{target_agent}': {msg}")

        passport_data = None
        if sender_passport is not None:
            if hasattr(sender_passport, "verify_signature"):
                is_valid, msg = sender_passport.verify_signature()
                if not is_valid:
                    raise ValueError(f"Cannot delegate with invalid passport: {msg}")
                for cap in scope:
                    if not sender_passport.has_capability(cap):
                        raise ValueError(f"Privilege escalation blocked: Passport does not authorize capability '{cap}'")
                passport_data = sender_passport.to_dict()
            elif isinstance(sender_passport, dict):
                passport_data = sender_passport

        envelope_body = {
            "protocol": "BTP/A2A/3.1",
            "envelope_nonce": nonce,
            "issued_at_unix": now,
            "expires_at_unix": now + ttl_seconds,
            "sender_agent_id": originating_agent,
            "sender_pubkey": sender_authority.public_key_hex,
            "recipient_agent_id": target_agent,
            "task_action": task_action,
            "task_payload": task_payload,
            "granted_scope": scope
        }
        if passport_data:
            envelope_body["sender_passport"] = passport_data

        canonical_bytes = rfc8785_canonicalize(envelope_body)
        signature = sender_authority.private_key.sign(canonical_bytes).hex()

        return {
            "a2a_envelope": envelope_body,
            "signature": signature
        }

    @classmethod
    def verify_incoming_handoff(cls, 
                                signed_packet: Dict[str, Any],
                                expected_recipient: str,
                                trusted_sender_pubkey: Optional[str] = None,
                                required_capability: Optional[str] = None) -> Tuple[bool, str, Dict[str, Any]]:
        """
        Agent B verifies incoming A2A envelope before executing delegated task.
        Validates envelope signature, expiration, recipient match, and sovereign passport.
        """
        try:
            envelope = signed_packet["a2a_envelope"]
            sig_hex = signed_packet["signature"]

            # 1. Recipient check
            if envelope.get("recipient_agent_id") != expected_recipient:
                return False, f"Recipient mismatch: Envelope targeted to '{envelope.get('recipient_agent_id')}', but received by '{expected_recipient}'", {}

            # 2. Expiration check
            now = time.time()
            if now > envelope.get("expires_at_unix", 0):
                return False, "A2A Envelope expired", {}

            # 3. Public key verification
            pubkey_hex = trusted_sender_pubkey or envelope.get("sender_pubkey")
            canonical_bytes = rfc8785_canonicalize(envelope)

            from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey
            public_key = Ed25519PublicKey.from_public_bytes(bytes.fromhex(pubkey_hex))
            public_key.verify(bytes.fromhex(sig_hex), canonical_bytes)

            # 4. Sovereign Passport Verification (if attached)
            if "sender_passport" in envelope:
                from src.agent_passport import SovereignAgentPassport
                p_dict = envelope["sender_passport"]
                passport = SovereignAgentPassport.from_dict(p_dict)
                p_valid, p_msg = passport.verify_signature()
                if not p_valid:
                    return False, f"A2A Delegated Passport Invalid: {p_msg}", {}

                # Check required capability against passport
                if required_capability and not passport.has_capability(required_capability):
                    return False, f"A2A Passport Missing Required Capability: '{required_capability}'", {}

            return True, "A2A Cryptographic Handoff Verified Clean", envelope

        except Exception as e:
            return False, f"A2A Signature Verification Failed: {str(e)}", {}