import logging
from typing import Any
from uuid import uuid4

import httpx

from a2a.client import A2ACardResolver, A2AClient
from a2a.types import (
    AgentCard,
    MessageSendParams,
    SendStreamingMessageRequest,
)

PUBLIC_AGENT_CARD_PATH = "/.well-known/agent.json"
EXTENDED_AGENT_CARD_PATH = "/agent/authenticatedExtendedCard"
# USER_INPUT = 'how much is 10 USD in INR?'
USER_INPUT = "My name is John Doe, email: john@example.com, phone: (555) 123-4567"


async def main() -> None:
    # Configure logging to show INFO level messages
    logging.basicConfig(level=logging.INFO)
    logger = logging.getLogger(__name__)  # Get a logger instance

    # --8<-- [start:A2ACardResolver]

    # base_url = 'http://localhost:9000'
    base_url = "http://a2a-contact-extractor.localtest.me:8080"

    async with httpx.AsyncClient() as httpx_client:
        # Initialize A2ACardResolver
        resolver = A2ACardResolver(
            httpx_client=httpx_client,
            base_url=base_url,
            # agent_card_path uses default, extended_agent_card_path also uses default
        )
        # --8<-- [end:A2ACardResolver]

        # Fetch Public Agent Card and Initialize Client
        _final_agent_card_to_use: AgentCard | None = None

        try:
            logger.info(f"Attempting to fetch public agent card from: {base_url}{PUBLIC_AGENT_CARD_PATH}")
            _public_card = await resolver.get_agent_card()  # Fetches from default public path
            logger.info("Successfully fetched public agent card:")
            logger.info(_public_card.model_dump_json(indent=2, exclude_none=True))
            _final_agent_card_to_use = _public_card
            logger.info("\nUsing PUBLIC agent card for client initialization (default).")

            if _public_card.supportsAuthenticatedExtendedCard:
                try:
                    logger.info(
                        "\nPublic card supports authenticated extended card. "
                        "Attempting to fetch from: "
                        f"{base_url}{EXTENDED_AGENT_CARD_PATH}"
                    )
                    auth_headers_dict = {"Authorization": "Bearer dummy-token-for-extended-card"}
                    _extended_card = await resolver.get_agent_card(
                        relative_card_path=EXTENDED_AGENT_CARD_PATH,
                        http_kwargs={"headers": auth_headers_dict},
                    )
                    logger.info("Successfully fetched authenticated extended agent card:")
                    logger.info(_extended_card.model_dump_json(indent=2, exclude_none=True))
                    _final_agent_card_to_use = _extended_card  # Update to use the extended card
                    logger.info("\nUsing AUTHENTICATED EXTENDED agent card for client initialization.")
                except Exception as e_extended:
                    logger.warning(
                        f"Failed to fetch extended agent card: {e_extended}. Will proceed with public card.",
                        exc_info=True,
                    )
            elif _public_card:  # supportsAuthenticatedExtendedCard is False or None
                logger.info("\nPublic card does not indicate support for an extended card. Using public card.")

        except Exception as e:
            logger.error(f"Critical error fetching public agent card: {e}", exc_info=True)
            raise RuntimeError("Failed to fetch the public agent card. Cannot continue.") from e

        # --8<-- [start:send_message]
        client = A2AClient(
            # httpx_client=httpx_client, agent_card=final_agent_card_to_use
            httpx_client=httpx_client,
            url=base_url,
        )
        logger.info("A2AClient initialized.")

        send_message_payload: dict[str, Any] = {
            "message": {
                "role": "user",
                "parts": [{"kind": "text", "text": USER_INPUT}],
                "messageId": uuid4().hex,
            },
        }
        # request = SendMessageRequest(
        #     id=str(uuid4()), params=MessageSendParams(**send_message_payload)
        # )

        # response = await client.send_message(request)
        # print(response.model_dump(mode='json', exclude_none=True))
        # --8<-- [end:send_message]

        # --8<-- [start:send_message_streaming]

        streaming_request = SendStreamingMessageRequest(
            id=str(uuid4()), params=MessageSendParams(**send_message_payload)
        )

        stream_response = client.send_message_streaming(streaming_request)

        async for chunk in stream_response:
            print(chunk.model_dump(mode="json", exclude_none=True))
        # --8<-- [end:send_message_streaming]


if __name__ == "__main__":
    import asyncio

    asyncio.run(main())
