import asyncio

from aiohttp import ClientSession, ClientTimeout
from smithy_core.retries import RetryStrategyOptions
from smithy_http.aio.aiohttp import AIOHTTPClient
from smithy_http.aio.crt import AWSCRTHTTPClient
from aws_sdk_bedrock_runtime.client import BedrockRuntimeClient
from aws_sdk_bedrock_runtime.models import ListAsyncInvokesInput
from aws_sdk_bedrock_runtime.config import Config
from smithy_aws_core.identity.environment import EnvironmentCredentialsResolver

async def main():
    config = Config(
        region="us-west-2",
        aws_credentials_identity_resolver=EnvironmentCredentialsResolver(),
        endpoint_uri="https://192.0.0.1",
        transport=AWSCRTHTTPClient(),
        #Preventing retries makes this fail faster
        retry_strategy=RetryStrategyOptions(max_attempts=1)
    )
    crt_client = BedrockRuntimeClient(config=config)
    try:
        await crt_client.list_async_invokes(ListAsyncInvokesInput())
    except Exception as e:
        print(f"CRT exception: {e}")

    aio_session = ClientSession(timeout=ClientTimeout(total=0.5))
    #Hackey to set an internal variable, but aiohttp has a pretty slow default conn timeout.  This is just for testing purposes.
    transport = AIOHTTPClient(_session=aio_session)
    config = Config(
        region="us-west-2",
        aws_credentials_identity_resolver=EnvironmentCredentialsResolver(),
        endpoint_uri="https://192.0.0.1",
        transport=transport,
        retry_strategy=RetryStrategyOptions(max_attempts=1)
    )
    aiohttp_client = BedrockRuntimeClient(config=config)
    try:
        await aiohttp_client.list_async_invokes(ListAsyncInvokesInput())
    except Exception as e:
        print(f"aiohttp total_timeout exception: {e}")

    aio_session2 = ClientSession(timeout=ClientTimeout(sock_connect=0.5))
    transport2 = AIOHTTPClient(_session=aio_session2)
    config2 = Config(
        region="us-west-2",
        aws_credentials_identity_resolver=EnvironmentCredentialsResolver(),
        endpoint_uri="https://192.0.0.1",
        transport=transport2,
        retry_strategy=RetryStrategyOptions(max_attempts=1)
    )
    aiohttp_client2 = BedrockRuntimeClient(config=config2)
    try:
        await aiohttp_client2.list_async_invokes(ListAsyncInvokesInput())
    except Exception as e:
        print(f"aiohttp sock_connect exception: {e}")


if __name__ == "__main__":
    asyncio.run(main())

# Output before change:
# CRT exception: Client timeout occurred:
# aiohttp total_timeout exception: Client timeout occurred:
# aiohttp sock_connect exception: Client timeout occurred: Connection timeout to host https://192.0.0.1/async-invoke?sortBy=SubmissionTime&sortOrder=Descending

# Output after change:
# CRT exception: CRT AWS_IO_SOCKET_TIMEOUT: socket operation timed out.
# aiohttp total_timeout exception: A timeout error occurred.
# aiohttp sock_connect exception: Connection timeout to host https://192.0.0.1/async-invoke?sortBy=SubmissionTime&sortOrder=Descending