# Agent-to-Agent (A2A) Protocol Handshake Simulator
import json

class SpecializedAgentB:
    """Agent B: Database Specialist Agent with secure task delegation."""
    def __init__(self, agent_id: str, auth_token: str):
        self.agent_id = agent_id
        self.auth_token = auth_token
        self._db = {
            "query_customer_ltv": {"status": "success", "result": "$12,450.00 USD"},
            "query_server_status": {"status": "success", "result": "uptime: 99.98%"}
        }

    def verify_handshake(self, sender_id: str, token: str) -> bool:
        """Verify the authentication handshake from requester Agent A."""
        print(f"[Agent B] Handshake received from {sender_id}. Verifying token...")
        return token == "SUPERVISOR_KEY_99"

    def handle_request(self, payload_str: str) -> str:
        """Process incoming tasks from trusted Agent A and return structured JSON."""
        try:
            payload = json.loads(payload_str)
            action = payload.get("action")
            args = payload.get("args", {})
            
            print(f"[Agent B] Received task payload: Action='{action}', Args={args}")
            
            if action in self._db:
                result = self._db[action]
                return json.dumps({
                    "status": "completed",
                    "sender": self.agent_id,
                    "response": result
                })
            else:
                return json.dumps({
                    "status": "error",
                    "sender": self.agent_id,
                    "error_message": f"Action '{action}' is not supported by Agent B."
                })
        except Exception as e:
            return json.dumps({"status": "error", "error_message": str(e)})

class SupervisorAgentA:
    """Agent A: Supervisor Agent that orchestrates and delegates tasks to peers."""
    def __init__(self, agent_id: str):
        self.agent_id = agent_id

    def execute_delegated_task(self, peer_agent: SpecializedAgentB, action: str) -> None:
        print(f"[Agent A] Starting A2A connection with {peer_agent.agent_id}...")
        
        # 1. Handshake & Authenticate
        auth_success = peer_agent.verify_handshake(self.agent_id, "SUPERVISOR_KEY_99")
        if not auth_success:
            print("[Agent A Error] Handshake failed! Aborting task execution.")
            return

        print("[Agent A] Handshake successful! Delegating task...")

        # 2. Formulate Structured Task Payload
        task_payload = {
            "action": action,
            "args": {"query_limit": 1}
        }
        payload_str = json.dumps(task_payload)

        # 3. Transmit task and wait for observation response
        raw_response = peer_agent.handle_request(payload_str)
        
        # 4. Parse A2A structured observation
        response = json.loads(raw_response)
        print(f"[Agent A] Received response from {peer_agent.agent_id}:")
        print(json.dumps(response, indent=4))
        
        if response.get("status") == "completed":
            print(f"[Agent A Success] Task resolved. Result: {response['response']['result']}")
        else:
            print(f"[Agent A Error] Task execution failed: {response.get('error_message')}")

# Run A2A Protocol loop simulation
if __name__ == "__main__":
    print("="*60)
    print("Agent-to-Agent Protocol Handshake & Delegation Loop")
    print("="*60)
    
    # Initialize agents
    agent_a = SupervisorAgentA(agent_id="Agent-Supervisor-A")
    agent_b = SpecializedAgentB(agent_id="Agent-DB-Specialist-B", auth_token="SUPERVISOR_KEY_99")

    # Delegate task to Agent B
    agent_a.execute_delegated_task(peer_agent=agent_b, action="query_customer_ltv")
    print("="*60)