"""
production inbox system demonstration.

shows complete setup with:
- cryptographic identity
- redis queue
- local and remote agents
- signature verification
- distributed tracing
"""

import asyncio
import logging
from typing import Any, Dict

from fastapi import FastAPI
import httpx
import uvicorn

from synqed.agent_email.addressing import AgentId
from synqed.agent_email.registry.api import get_registry
from synqed.agent_email.registry.models import AgentRegistryEntry
from synqed.agent_email.inbox import (
    LocalAgentRuntime,
    register_agent_runtime,
    router,
    generate_keypair,
    sign_message,
)
from synqed.agent_email.inbox.startup import create_lifespan


# configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger(__name__)


class EchoAgent(LocalAgentRuntime):
    """simple echo agent for local processing."""
    
    def __init__(self, name: str):
        self.name = name
    
    async def handle_a2a_envelope(
        self,
        sender: str,
        recipient: str,
        envelope: Dict[str, Any],
    ) -> Dict[str, Any] | None:
        """echo back the message with a response."""
        content = envelope.get("content", "")
        logger.info(f"{self.name} received: '{content}' from {sender}")
        
        # return response envelope
        return {
            "thread_id": envelope.get("thread_id"),
            "role": "assistant",
            "content": f"Echo from {self.name}: {content}",
        }


async def setup_agents():
    """
    setup demo agents with keypairs.
    
    creates:
    - alice: local agent (has runtime)
    - bob: local agent (has runtime)
    - charlie: remote agent (no runtime, will forward to remote_url)
    """
    logger.info("setting up agents...")
    
    registry = get_registry()
    
    # generate keypairs for all agents
    alice_private, alice_public = generate_keypair()
    bob_private, bob_public = generate_keypair()
    charlie_private, charlie_public = generate_keypair()
    
    # register alice (local agent)
    registry.register(AgentRegistryEntry(
        agent_id="agent://demo/alice",
        email_like="alice@demo",
        inbox_url="http://localhost:8000/v1/a2a/inbox",
        public_key=alice_public,
        capabilities=["a2a/1.0", "echo"],
        metadata={"description": "Alice the Echo Agent"},
    ))
    
    # register alice's local runtime
    register_agent_runtime("agent://demo/alice", EchoAgent("Alice"))
    
    # register bob (local agent)
    registry.register(AgentRegistryEntry(
        agent_id="agent://demo/bob",
        email_like="bob@demo",
        inbox_url="http://localhost:8000/v1/a2a/inbox",
        public_key=bob_public,
        capabilities=["a2a/1.0", "echo"],
        metadata={"description": "Bob the Echo Agent"},
    ))
    
    # register bob's local runtime
    register_agent_runtime("agent://demo/bob", EchoAgent("Bob"))
    
    # register charlie (remote agent - no local runtime)
    # in real deployment, this would point to charlie's actual server
    registry.register(AgentRegistryEntry(
        agent_id="agent://demo/charlie",
        email_like="charlie@demo",
        inbox_url="http://localhost:8001/v1/a2a/inbox",  # different host
        public_key=charlie_public,
        capabilities=["a2a/1.0", "remote"],
        metadata={"description": "Charlie the Remote Agent"},
    ))
    
    logger.info("agents registered:")
    for entry in registry.list_all():
        logger.info(f"  - {entry.agent_id} ({entry.email_like})")
    
    # store keypairs for sending messages
    return {
        "alice": (alice_private, alice_public),
        "bob": (bob_private, bob_public),
        "charlie": (charlie_private, charlie_public),
    }


async def send_message(
    sender_id: str,
    sender_private_key: str,
    recipient_id: str,
    content: str,
    thread_id: str = "demo-thread",
    trace_id: str | None = None,
) -> Dict[str, Any]:
    """
    send signed message to another agent.
    
    args:
        sender_id: sender canonical uri
        sender_private_key: sender's private key for signing
        recipient_id: recipient canonical uri
        content: message content
        thread_id: conversation thread id
        trace_id: optional trace id
        
    returns:
        response dict
    """
    # prepare message
    message = {
        "thread_id": thread_id,
        "role": "user",
        "content": content,
    }
    
    # sign message
    signature = sign_message(
        private_key_b64=sender_private_key,
        sender=sender_id,
        recipient=recipient_id,
        message=message,
        thread_id=thread_id,
    )
    
    # prepare envelope
    envelope = {
        "sender": sender_id,
        "recipient": recipient_id,
        "message": message,
        "signature": signature,
    }
    
    if trace_id:
        envelope["trace_id"] = trace_id
    
    # send to inbox
    async with httpx.AsyncClient() as client:
        response = await client.post(
            "http://localhost:8000/v1/a2a/inbox",
            json=envelope,
            timeout=30.0,
        )
        response.raise_for_status()
        return response.json()


async def demo_local_to_local(keypairs):
    """demonstrate alice sending to bob (both local)."""
    logger.info("\n=== DEMO 1: Local to Local ===")
    
    alice_private, _ = keypairs["alice"]
    
    result = await send_message(
        sender_id="agent://demo/alice",
        sender_private_key=alice_private,
        recipient_id="agent://demo/bob",
        content="Hello Bob, this is Alice!",
        thread_id="alice-bob-thread",
    )
    
    logger.info(f"Message sent: {result['message_id']}")
    logger.info(f"Trace ID: {result['trace_id']}")
    logger.info(f"Status: {result['status']}")
    
    # give worker time to process
    await asyncio.sleep(2)


async def demo_signature_verification(keypairs):
    """demonstrate signature verification (should fail with wrong key)."""
    logger.info("\n=== DEMO 2: Signature Verification ===")
    
    # try to send message with bob's key but claim to be alice (should fail)
    bob_private, _ = keypairs["bob"]
    
    try:
        result = await send_message(
            sender_id="agent://demo/alice",  # claim to be alice
            sender_private_key=bob_private,    # but use bob's key
            recipient_id="agent://demo/bob",
            content="Impersonation attempt!",
            thread_id="fake-thread",
        )
        
        logger.info(f"Result: {result}")
        
        if result["status"] == "rejected":
            logger.info("✓ Signature verification correctly rejected invalid signature")
        else:
            logger.error("✗ Signature verification failed to catch invalid signature!")
    
    except Exception as e:
        logger.error(f"Error: {e}")


async def demo_rate_limiting(keypairs):
    """demonstrate rate limiting."""
    logger.info("\n=== DEMO 3: Rate Limiting ===")
    
    alice_private, _ = keypairs["alice"]
    
    # send many messages rapidly
    success_count = 0
    rate_limited_count = 0
    
    for i in range(120):  # exceed 100 per minute limit
        try:
            result = await send_message(
                sender_id="agent://demo/alice",
                sender_private_key=alice_private,
                recipient_id="agent://demo/bob",
                content=f"Message {i}",
                thread_id=f"rate-test-{i}",
            )
            
            if result["status"] == "accepted":
                success_count += 1
        
        except httpx.HTTPStatusError as e:
            if e.response.status_code == 429:
                rate_limited_count += 1
            else:
                raise
    
    logger.info(f"Sent: {success_count}, Rate limited: {rate_limited_count}")
    logger.info(f"✓ Rate limiting working (limited after ~100 requests)")


async def demo_trace_propagation(keypairs):
    """demonstrate trace propagation."""
    logger.info("\n=== DEMO 4: Trace Propagation ===")
    
    alice_private, _ = keypairs["alice"]
    
    # first message
    result1 = await send_message(
        sender_id="agent://demo/alice",
        sender_private_key=alice_private,
        recipient_id="agent://demo/bob",
        content="Starting a traced conversation",
        thread_id="traced-thread",
    )
    
    parent_trace = result1["trace_id"]
    logger.info(f"Parent trace: {parent_trace}")
    
    # second message with parent trace
    result2 = await send_message(
        sender_id="agent://demo/alice",
        sender_private_key=alice_private,
        recipient_id="agent://demo/bob",
        content="Continuing the conversation",
        thread_id="traced-thread",
        trace_id=parent_trace,
    )
    
    logger.info(f"Child trace: {result2['trace_id']}")
    logger.info("✓ Trace IDs linked for distributed tracing")


async def run_demo():
    """run all demonstrations."""
    logger.info("Starting production inbox demo...")
    
    # setup agents
    keypairs = await setup_agents()
    
    # wait for workers to start
    await asyncio.sleep(2)
    
    # run demos
    await demo_local_to_local(keypairs)
    await demo_signature_verification(keypairs)
    await demo_trace_propagation(keypairs)
    
    # rate limiting demo last (sends many messages)
    # await demo_rate_limiting(keypairs)  # uncomment to test
    
    logger.info("\n=== Demo Complete ===")
    logger.info("Check logs above to see:")
    logger.info("  - Message processing by local agents")
    logger.info("  - Signature verification")
    logger.info("  - Distributed tracing")
    logger.info("  - Rate limiting (if enabled)")


# create fastapi app
app = FastAPI(
    title="Synqed Production Inbox Demo",
    lifespan=create_lifespan("redis://localhost:6379"),
)

# include inbox router
app.include_router(router)


@app.on_event("startup")
async def startup_event():
    """run demo after startup."""
    # give system time to initialize
    await asyncio.sleep(3)
    
    # run demo
    try:
        await run_demo()
    except Exception as e:
        logger.error(f"Demo failed: {e}", exc_info=True)


if __name__ == "__main__":
    """
    run demo server.
    
    requirements:
    1. redis running on localhost:6379
       docker run -d -p 6379:6379 redis:7
    
    2. install dependencies:
       pip install cryptography redis httpx
    
    3. run:
       python examples/production_inbox_demo.py
    """
    print("=" * 60)
    print("Production Inbox Demo")
    print("=" * 60)
    print()
    print("Requirements:")
    print("  1. Redis: docker run -d -p 6379:6379 redis:7")
    print("  2. Dependencies: pip install cryptography redis httpx")
    print()
    print("Starting server on http://localhost:8000")
    print("=" * 60)
    print()
    
    uvicorn.run(
        app,
        host="0.0.0.0",
        port=8000,
        log_level="info",
    )

