import time
import traceback

import wrapt

from aws_xray_sdk.core.recorder import AWSXRayRecorder


class AsyncAWSXRayRecorder(AWSXRayRecorder):
    def capture_async(self, name=None):
        """
        A decorator that records enclosed function in a subsegment.
        It only works with asynchronous functions.

        params str name: The name of the subsegment. If not specified
        the function name will be used.
        """

        @wrapt.decorator
        async def wrapper(wrapped, instance, args, kwargs):
            func_name = name
            if not func_name:
                func_name = wrapped.__name__

            result = await self.record_subsegment_async(
                wrapped, instance, args, kwargs,
                name=func_name,
                namespace='local',
                meta_processor=None,
            )

            return result

        return wrapper

    async def record_subsegment_async(self, wrapped, instance, args, kwargs, name,
                                namespace, meta_processor):

        subsegment = self.begin_subsegment(name, namespace)

        exception = None
        stack = None
        return_value = None

        try:
            return_value = await wrapped(*args, **kwargs)
            return return_value
        except Exception as e:
            exception = e
            stack = traceback.extract_stack(limit=self._max_trace_back)
            raise
        finally:
            end_time = time.time()
            if callable(meta_processor):
                meta_processor(
                    wrapped=wrapped,
                    instance=instance,
                    args=args,
                    kwargs=kwargs,
                    return_value=return_value,
                    exception=exception,
                    subsegment=subsegment,
                    stack=stack,
                )
            elif exception:
                if subsegment:
                    subsegment.add_exception(exception, stack)

            self.end_subsegment(end_time)
