# Code generated by smithy-python-codegen DO NOT EDIT.

import asyncio
from asyncio import Future, iscoroutine, sleep
from copy import copy, deepcopy
from dataclasses import replace
import logging
from typing import Any, Awaitable, Callable, cast

from smithy_core import URI
from smithy_core.deserializers import DeserializeableShape
from smithy_core.endpoints import EndpointResolverParams
from smithy_core.exceptions import SmithyRetryException
from smithy_core.interceptors import (
    InputContext,
    Interceptor,
    InterceptorChain,
    OutputContext,
    RequestContext,
    ResponseContext,
)
from smithy_core.interfaces.exceptions import HasFault
from smithy_core.interfaces.retries import RetryErrorInfo, RetryErrorType
from smithy_core.schemas import APIOperation
from smithy_core.serializers import SerializeableShape
from smithy_core.types import TypedProperties
from smithy_http.aio.interfaces import HTTPRequest, HTTPResponse
from smithy_http.interfaces import HTTPRequestConfiguration
from smithy_http.plugins import user_agent_plugin

from .config import Config, Plugin
from .deserialize import _deserialize_query_agent
from .models import QUERY_AGENT, QueryAgentInput, QueryAgentOutput, ServiceError
from .serialize import _serialize_query_agent
from smithy_http import Field



logger = logging.getLogger(__name__)

class Vibe:
    """
    Client for Vibe

    :param config: Optional configuration for the client. Here you can set things like the
        endpoint for HTTP services or auth credentials.

    :param plugins: A list of callables that modify the configuration dynamically. These
        can be used to set defaults, for example.
    """
    def __init__(self, config: Config | None = None, plugins: list[Plugin] | None = None):
        self._config = config or Config()

        client_plugins: list[Plugin] = [
            user_agent_plugin,
        ]
        if plugins:
            client_plugins.extend(plugins)

        for plugin in client_plugins:
            plugin(self._config)

    async def query_agent(self, input: QueryAgentInput, plugins: list[Plugin] | None = None) -> QueryAgentOutput:
        """
        Invokes the QueryAgent operation.

        :param input: The operation's input.

        :param plugins: A list of callables that modify the configuration dynamically.
            Changes made by these plugins only apply for the duration of the operation
            execution and will not affect any other operation invocations.

        """
        operation_plugins: list[Plugin] = [

        ]
        if plugins:
            operation_plugins.extend(plugins)

        return await self._execute_operation(
            input=input,
            plugins=operation_plugins,
            serialize=_serialize_query_agent,
            deserialize=_deserialize_query_agent,
            config=self._config,
            operation=QUERY_AGENT,
        )

    def _classify_error(
        self,
        *,
        error: Exception,
        context: ResponseContext[Any, HTTPRequest, HTTPResponse | None]
    ) -> RetryErrorInfo:
        logger.debug("Classifying error: %s", error)

        if not isinstance(error, HasFault) and not context.transport_response:
            return RetryErrorInfo(error_type=RetryErrorType.TRANSIENT)

        if context.transport_response:
            if context.transport_response.status in [429, 503]:
                retry_after = None
                retry_header = context.transport_response.fields["retry-after"]
                if retry_header and retry_header.values:
                    retry_after = float(retry_header.values[0])
                return RetryErrorInfo(error_type=RetryErrorType.THROTTLING, retry_after_hint=retry_after)

            if context.transport_response.status >= 500:
                return RetryErrorInfo(error_type=RetryErrorType.SERVER_ERROR)

        error_type = RetryErrorType.CLIENT_ERROR
        if isinstance(error, HasFault) and error.fault == "server":
            error_type = RetryErrorType.SERVER_ERROR

        return RetryErrorInfo(error_type=error_type)

    async def _execute_operation[Input: SerializeableShape, Output: DeserializeableShape](
        self,
        input: Input,
        plugins: list[Plugin],
        serialize: Callable[[Input, Config], Awaitable[HTTPRequest]],
        deserialize: Callable[[HTTPResponse, Config], Awaitable[Output]],
        config: Config,
        operation: APIOperation[Input, Output],
        request_future: Future[RequestContext[Any, HTTPRequest]] | None = None,
        response_future: Future[HTTPResponse] | None = None,
    ) -> Output:
        try:
            return await self._handle_execution(
                input, plugins, serialize, deserialize, config, operation,
                request_future, response_future,
            )
        except Exception as e:
            if request_future is not None and not request_future.done():
                request_future.set_exception(ServiceError(e))
            if response_future is not None and not response_future.done():
                response_future.set_exception(ServiceError(e))

            # Make sure every exception that we throw is an instance of ServiceError so
            # customers can reliably catch everything we throw.
            if not isinstance(e, ServiceError):
                raise ServiceError(e) from e
            raise

    async def _handle_execution[Input: SerializeableShape, Output: DeserializeableShape](
        self,
        input: Input,
        plugins: list[Plugin],
        serialize: Callable[[Input, Config], Awaitable[HTTPRequest]],
        deserialize: Callable[[HTTPResponse, Config], Awaitable[Output]],
        config: Config,
        operation: APIOperation[Input, Output],
        request_future: Future[RequestContext[Any, HTTPRequest]] | None,
        response_future: Future[HTTPResponse] | None,
    ) -> Output:
        operation_name = operation.schema.id.name
        logger.debug('Making request for operation "%s" with parameters: %s', operation_name, input)
        config = deepcopy(config)
        for plugin in plugins:
            plugin(config)

        input_context = InputContext(request=input, properties=TypedProperties({"config": config}))
        transport_request: HTTPRequest | None = None
        output_context: OutputContext[Input, Output, HTTPRequest | None, HTTPResponse | None] | None = None

        client_interceptors = cast(
            list[Interceptor[Input, Output, HTTPRequest, HTTPResponse]], list(config.interceptors)
        )
        interceptor_chain = InterceptorChain(client_interceptors)

        try:
            # Step 1: Invoke read_before_execution
            interceptor_chain.read_before_execution(input_context)

            # Step 2: Invoke the modify_before_serialization hooks
            input_context = replace(
                input_context,
                request=interceptor_chain.modify_before_serialization(input_context)
            )

            # Step 3: Invoke the read_before_serialization hooks
            interceptor_chain.read_before_serialization(input_context)

            # Step 4: Serialize the request
            logger.debug("Serializing request for: %s", input_context.request)
            transport_request = await serialize(input_context.request, config)
            request_context = RequestContext(
                request=input_context.request,
                transport_request=transport_request,
                properties=input_context.properties,
            )
            logger.debug("Serialization complete. Transport request: %s", request_context.transport_request)

            # Step 5: Invoke read_after_serialization
            interceptor_chain.read_after_serialization(request_context)

            # Step 6: Invoke modify_before_retry_loop
            request_context = replace(
                request_context,
                transport_request=interceptor_chain.modify_before_retry_loop(request_context)
            )

            # Step 7: Acquire the retry token.
            retry_strategy = config.retry_strategy
            retry_token = retry_strategy.acquire_initial_retry_token()

            while True:
                # Make an attempt
                output_context = await self._handle_attempt(
                    deserialize,
                    interceptor_chain,
                    replace(
                      request_context,
                      transport_request = copy(request_context.transport_request)
                    ),
                    config,
                    operation,
                    request_future,
                )

                if isinstance(output_context.response, Exception):
                    # Step 7u: Reacquire retry token if the attempt failed
                    try:
                        retry_token = retry_strategy.refresh_retry_token_for_retry(
                            token_to_renew=retry_token,
                            error_info=self._classify_error(
                                error=output_context.response,
                                context=output_context,
                            )
                        )
                    except SmithyRetryException:
                        raise output_context.response
                    logger.debug(
                        "Retry needed. Attempting request #%s in %.4f seconds.",
                        retry_token.retry_count + 1,
                        retry_token.retry_delay
                    )
                    await sleep(retry_token.retry_delay)
                    current_body = output_context.transport_request.body
                    if (seek := getattr(current_body, "seek", None)) is not None:
                        if iscoroutine((result := seek(0))):
                            await result
                else:
                    # Step 8: Invoke record_success
                    retry_strategy.record_success(token=retry_token)
                    if response_future is not None:
                        response_future.set_result(
                            output_context.transport_response  # type: ignore
                        )
                    break
        except Exception as e:
            if output_context is not None:
                logger.exception("Exception occurred while handling: %s", output_context.response)
                output_context = replace(output_context, response=e)
            else:
                output_context = OutputContext(
                    request=input_context.request,
                    response=e,
                    transport_request=transport_request,
                    transport_response=None,
                    properties=input_context.properties
                )

        return await self._finalize_execution(interceptor_chain, output_context)

    async def _handle_attempt[Input: SerializeableShape, Output: DeserializeableShape](
        self,
        deserialize: Callable[[HTTPResponse, Config], Awaitable[Output]],
        interceptor: Interceptor[Input, Output, HTTPRequest, HTTPResponse],
        context: RequestContext[Input, HTTPRequest],
        config: Config,
        operation: APIOperation[Input, Output],
        request_future: Future[RequestContext[Input, HTTPRequest]] | None,
    ) -> OutputContext[Input, Output, HTTPRequest, HTTPResponse | None]:
        transport_response: HTTPResponse | None = None
        try:
            # Step 7a: Invoke read_before_attempt
            interceptor.read_before_attempt(context)

            # Step 7f: Invoke endpoint_resolver.resolve_endpoint
            endpoint_resolver_parameters = EndpointResolverParams(
                operation=operation,
                input=context.request,
                context=context.properties
            )
            logger.debug("Calling endpoint resolver with parameters: %s", endpoint_resolver_parameters)
            endpoint = await config.endpoint_resolver.resolve_endpoint(
                endpoint_resolver_parameters
            )
            logger.debug("Endpoint resolver result: %s", endpoint)
            if not endpoint.uri.path:
                path = ""
            elif endpoint.uri.path.endswith("/"):
                path = endpoint.uri.path[:-1]
            else:
                path = endpoint.uri.path
            if context.transport_request.destination.path:
                path += context.transport_request.destination.path
            context.transport_request.destination = URI(
                scheme=endpoint.uri.scheme,
                host=context.transport_request.destination.host + endpoint.uri.host,
                path=path,
                port=endpoint.uri.port,
                query=context.transport_request.destination.query,
            )

            if (headers := endpoint.properties.get("headers")) is not None:
                context.transport_request.fields.extend(headers)

            # Step 7g: Invoke modify_before_signing
            context = replace(
                context,
                transport_request=interceptor.modify_before_signing(context)
            )

            # Step 7h: Invoke read_before_signing
            interceptor.read_before_signing(context)

            # Step 7j: Invoke read_after_signing
            interceptor.read_after_signing(context)

            # Step 7k: Invoke modify_before_transmit
            context = replace(
                context,
                transport_request=interceptor.modify_before_transmit(context)
            )

            # Step 7l: Invoke read_before_transmit
            interceptor.read_before_transmit(context)

            # Step 7m: Invoke http_client.send
            request_config = config.http_request_config or HTTPRequestConfiguration()
            logger.debug("HTTP request config: %s", request_config)
            logger.debug("Sending HTTP request: %s", context.transport_request)

            if request_future is not None:
                response_task = asyncio.create_task(config.http_client.send(
                    request=context.transport_request,
                    request_config=request_config,
                ))
                request_future.set_result(context)
                transport_response = await response_task
            else:
                transport_response = await config.http_client.send(
                    request=context.transport_request,
                    request_config=request_config,
                )

            response_context = ResponseContext(
                request=context.request,
                transport_request=context.transport_request,
                transport_response=transport_response,
                properties=context.properties
            )
            logger.debug("Received HTTP response: %s", response_context.transport_response)

            # Step 7n: Invoke read_after_transmit
            interceptor.read_after_transmit(response_context)

            # Step 7o: Invoke modify_before_deserialization
            response_context = replace(
                response_context,
                transport_response=interceptor.modify_before_deserialization(response_context)
            )

            # Step 7p: Invoke read_before_deserialization
            interceptor.read_before_deserialization(response_context)

            # Step 7q: deserialize
            logger.debug("Deserializing transport response: %s", response_context.transport_response)
            output = await deserialize(
                response_context.transport_response, config
            )
            output_context = OutputContext(
                request=response_context.request,
                response=output,
                transport_request=response_context.transport_request,
                transport_response=response_context.transport_response,
                properties=response_context.properties
            )
            logger.debug("Deserialization complete. Response: %s", output_context.response)

            # Step 7r: Invoke read_after_deserialization
            interceptor.read_after_deserialization(output_context)
        except Exception as e:
            output_context: OutputContext[Input, Output, HTTPRequest, HTTPResponse] = OutputContext(
                request=context.request,
                response=e,  # type: ignore
                transport_request=context.transport_request,
                transport_response=transport_response,
                properties=context.properties
            )

        return await self._finalize_attempt(interceptor, output_context)

    async def _finalize_attempt[Input: SerializeableShape, Output: DeserializeableShape](
        self,
        interceptor: Interceptor[Input, Output, HTTPRequest, HTTPResponse],
        context: OutputContext[Input, Output, HTTPRequest, HTTPResponse | None],
    ) -> OutputContext[Input, Output, HTTPRequest, HTTPResponse | None]:
        # Step 7s: Invoke modify_before_attempt_completion
        try:
            context = replace(
                context,
                response=interceptor.modify_before_attempt_completion(context)
            )
        except Exception as e:
            logger.exception("Exception occurred while handling: %s", context.response)
            context = replace(context, response=e)

        # Step 7t: Invoke read_after_attempt
        try:
            interceptor.read_after_attempt(context)
        except Exception as e:
            context = replace(context, response=e)

        return context

    async def _finalize_execution[Input: SerializeableShape, Output: DeserializeableShape](
        self,
        interceptor: Interceptor[Input, Output, HTTPRequest, HTTPResponse],
        context: OutputContext[Input, Output, HTTPRequest | None, HTTPResponse | None],
    ) -> Output:
        try:
            # Step 9: Invoke modify_before_completion
            context = replace(
                context,
                response=interceptor.modify_before_completion(context)
            )

            # Step 10: Invoke trace_probe.dispatch_events
            try:
                pass
            except Exception as e:
                # log and ignore exceptions
                logger.exception("Exception occurred while dispatching trace events: %s", e)
                pass
        except Exception as e:
            logger.exception("Exception occurred while handling: %s", context.response)
            context = replace(context, response=e)

        # Step 11: Invoke read_after_execution
        try:
            interceptor.read_after_execution(context)
        except Exception as e:
            context = replace(context, response=e)

        # Step 12: Return / throw
        if isinstance(context.response, Exception):
            raise context.response

        # We may want to add some aspects of this context to the output types so we can
        # return it to the end-users.
        return context.response



class ApiKeyInterceptor(Interceptor):
    def __init__(self, api_key: str):
        self.api_key = api_key

    def modify_before_signing(self, context: RequestContext) -> HTTPRequest:
        """Add the x-api-key header to the request before it's signed."""

        request = context.transport_request
        # Add the x-api-key header directly
        request.fields.set_field(Field(name="x-api-key", values=[str(self.api_key)]))
        return context.transport_request

class VibeClient(Vibe):
    """
    An extension of the Vibe client with built-in API key authentication.
    """

    def __init__(self, api_key: str, endpoint: str = "https://us-east-1.vibe.api.astartech.ai/", **kwargs):
        """
        Initialize a Vibe client with API key.

        Args:
            api_key: The API key for authentication
            endpoint: The API endpoint URI (default: https://us-east-1.vibe.api.astartech.ai/)
            **kwargs: Additional arguments to pass to the Vibe constructor

        Raises:
            ValueError: If api_key is None or empty
        """
        # Check that api_key is not None or empty
        if not api_key:
            raise ValueError("api_key cannot be None or empty")

        # Create config with API key interceptor
        config = Config(
            endpoint_uri=endpoint,
            interceptors=[ApiKeyInterceptor(api_key)]
        )

        # Initialize the parent Vibe class with our config
        super().__init__(config=config, **kwargs)
