#!/usr/bin/env python3
"""
Basic AIPM Handshake Example

Demonstrates two AI agents (OpenAI-based and LangGraph-based) performing
a complete handshake using the AIPM protocol.
"""

import json
from aipm import (
    AIPMAgent,
    AgentIdentity,
    Capabilities,
    TrustScore,
    MessageType,
)


def print_message(label: str, message):
    """Pretty print message"""
    print(f"\n{'='*60}")
    print(f"  {label}")
    print(f"{'='*60}")
    print(f"Type: {message.type.value}")
    print(f"From: {message.sender.agent_id}")
    print(f"To: {message.receiver.agent_id}")
    print(f"Payload: {json.dumps(message.payload, indent=2)}")


def main():
    """Run basic handshake example"""

    print("\n" + "="*60)
    print("  AIPM BASIC HANDSHAKE EXAMPLE")
    print("="*60)

    # Create OpenAI-based agent
    print("\n[1] Creating OpenAI-based Agent...")
    openai_identity = AgentIdentity(
        agent_id="agent-openai-001",
        organization_id="openai",
        name="OpenAI Assistant",
        version="1.0.0",
        capabilities=Capabilities(
            skills=["text-generation", "code-review", "summarization"],
            models=["gpt-4", "gpt-3.5-turbo"],
            tools=["code-interpreter", "web-browser"],
            max_context=128000,
            memory_support=True,
            languages=["en", "es", "fr", "de"],
        ),
        trust_score=TrustScore(
            reliability=0.99,
            accuracy=0.95,
            avg_latency_ms=250.0,
            success_rate=0.98,
            total_interactions=10000,
        ),
    )
    openai_agent = AIPMAgent(openai_identity)
    print(f"✓ Created {openai_agent}")

    # Create LangGraph-based agent
    print("\n[2] Creating LangGraph-based Agent...")
    langgraph_identity = AgentIdentity(
        agent_id="agent-langgraph-001",
        organization_id="langchain",
        name="LangGraph Orchestrator",
        version="1.0.0",
        capabilities=Capabilities(
            skills=["workflow-orchestration", "multi-agent-coordination", "data-processing"],
            models=["claude-3-opus", "claude-3-sonnet"],
            tools=["database", "api-caller", "document-processor"],
            max_context=200000,
            memory_support=True,
            languages=["en", "zh", "ja"],
        ),
        trust_score=TrustScore(
            reliability=0.97,
            accuracy=0.93,
            avg_latency_ms=300.0,
            success_rate=0.96,
            total_interactions=5000,
        ),
    )
    langgraph_agent = AIPMAgent(langgraph_identity)
    print(f"✓ Created {langgraph_agent}")

    # Start handshake
    print("\n" + "="*60)
    print("  HANDSHAKE PROTOCOL")
    print("="*60)

    # Step 1: OpenAI agent initiates with HELLO
    print("\n[Step 1] OpenAI Agent → LangGraph Agent: HELLO")
    hello_msg = openai_agent.initiate_handshake(langgraph_identity.to_reference())
    print_message("HELLO Message", hello_msg)

    # Step 2: LangGraph agent responds with CAPABILITY_EXCHANGE
    print("\n[Step 2] LangGraph Agent → OpenAI Agent: CAPABILITY_EXCHANGE")
    capability_msg = langgraph_agent.process_message(hello_msg)
    print_message("CAPABILITY_EXCHANGE Message", capability_msg)

    # Step 3: OpenAI agent continues with AUTHENTICATION
    print("\n[Step 3] OpenAI Agent → LangGraph Agent: AUTHENTICATION")
    auth_msg = openai_agent.process_message(capability_msg)
    print_message("AUTHENTICATION Message", auth_msg)

    # Step 4: LangGraph agent exchanges PUBLIC_KEY
    print("\n[Step 4] LangGraph Agent → OpenAI Agent: PUBLIC_KEY_EXCHANGE")
    key_msg = langgraph_agent.process_message(auth_msg)
    print_message("PUBLIC_KEY_EXCHANGE Message", key_msg)

    # Step 5: OpenAI agent verifies TRUST
    print("\n[Step 5] OpenAI Agent → LangGraph Agent: TRUST_VERIFICATION")
    trust_msg = openai_agent.process_message(key_msg)
    print_message("TRUST_VERIFICATION Message", trust_msg)

    # Step 6: LangGraph agent sends READY
    print("\n[Step 6] LangGraph Agent → OpenAI Agent: READY")
    ready_msg = langgraph_agent.process_message(trust_msg)
    print_message("READY Message", ready_msg)

    # Step 7: OpenAI agent confirms READY
    print("\n[Step 7] OpenAI Agent confirms READY")
    openai_agent.process_message(ready_msg)

    # Verify handshake completion
    print("\n" + "="*60)
    print("  HANDSHAKE COMPLETE")
    print("="*60)

    openai_ready = openai_agent.is_ready(langgraph_identity.to_reference())
    langgraph_ready = langgraph_agent.is_ready(openai_identity.to_reference())

    print(f"\nOpenAI Agent Ready: {openai_ready} ✓")
    print(f"LangGraph Agent Ready: {langgraph_ready} ✓")

    # Exchange identity information
    print("\n" + "="*60)
    print("  PEER IDENTITY EXCHANGE")
    print("="*60)

    openai_peer = openai_agent.get_peer_identity(langgraph_identity.to_reference())
    langgraph_peer = langgraph_agent.get_peer_identity(openai_identity.to_reference())

    print(f"\nOpenAI knows LangGraph as:")
    print(f"  - Name: {openai_peer.name}")
    print(f"  - Skills: {', '.join(openai_peer.capabilities.skills)}")
    print(f"  - Models: {', '.join(openai_peer.capabilities.models)}")
    print(f"  - Trust Score: {openai_peer.trust_score.reliability:.2f}")

    print(f"\nLangGraph knows OpenAI as:")
    print(f"  - Name: {langgraph_peer.name}")
    print(f"  - Skills: {', '.join(langgraph_peer.capabilities.skills)}")
    print(f"  - Models: {', '.join(langgraph_peer.capabilities.models)}")
    print(f"  - Trust Score: {langgraph_peer.trust_score.reliability:.2f}")

    # Create a task request
    print("\n" + "="*60)
    print("  TASK REQUEST EXAMPLE")
    print("="*60)

    task_msg = openai_agent.create_task_request(
        langgraph_identity.to_reference(),
        task_description="Process customer feedback data and generate insights",
        priority="high",
        dataset_size=1000,
        deadline="2026-07-10T00:00:00Z",
    )
    print_message("Task Request from OpenAI to LangGraph", task_msg)

    print("\n" + "="*60)
    print("  ✓ EXAMPLE COMPLETE")
    print("="*60)
    print("\nTwo agents from different vendors successfully:")
    print("  1. Completed secure handshake")
    print("  2. Exchanged capabilities")
    print("  3. Verified trust")
    print("  4. Ready for task delegation")
    print("\nThis is the foundation for cross-vendor agent interoperability! 🚀")
    print()


if __name__ == "__main__":
    main()