# Copyright (c) 2020 CloudZero - ALL RIGHTS RESERVED - PROPRIETARY AND CONFIDENTIAL
# Unauthorized copying of this file and/or project, via any medium is strictly prohibited.
# Direct all questions to legal@cloudzero.com

import asyncio
import functools
from collections.abc import Callable, Iterable
from typing import Any, Optional

from aws_xray_sdk.core import xray_recorder
from aws_xray_sdk.core.models.dummy_entities import DummySegment


def _capture_xray_entity() -> Optional[Any]:
    """Return the X-Ray entity (segment or subsegment) active in the calling thread,
    or None if no entity is active or the entity is a sentinel (DummySegment)
    indicating the SDK is disabled.
    """
    try:
        entity = xray_recorder.get_trace_entity()
    except Exception:
        return None
    if entity is None:
        return None
    if isinstance(entity, DummySegment):
        return None
    return entity


def _xray_propagating(fn: Callable, entity: Optional[Any]) -> Callable:
    """Wrap `fn` so that, when invoked in a worker thread, the captured X-Ray
    `entity` is planted as that thread's active entity for the duration of the
    call. If `entity` is None, the wrapper is a no-op pass-through.

    Concurrent workers share the same parent entity. The X-Ray SDK's Entity
    type does not lock its `subsegments` list, so safe concurrent subsegment
    creation relies on CPython's GIL making `list.append` effectively atomic.
    """
    if entity is None:
        return fn

    @functools.wraps(fn)
    def wrapper(*args, **kwargs):
        xray_recorder.set_trace_entity(entity)
        try:
            return fn(*args, **kwargs)
        finally:
            xray_recorder.clear_trace_entities()

    return wrapper


def async_map(function: Callable, arguments: Iterable[Iterable], max_concurrent: Optional[int] = None) -> Iterable:
    """Asynchronously map the synchronous `function` over a list of
    `arguments` using asyncio.

    The X-Ray trace entity active in the calling thread is propagated into each
    worker thread so that subsegments created inside `function` are parented
    correctly. Propagation is a no-op when no X-Ray entity is active or tracing
    is disabled.
    """

    if max_concurrent is not None and max_concurrent <= 0:
        raise ValueError('max_concurrent must be at least 1')

    entity = _capture_xray_entity()
    function_list = [_xray_propagating(functools.partial(function, *x), entity) for x in arguments]

    async def main():
        loop = asyncio.get_running_loop()
        if max_concurrent is None:
            tasks = [loop.run_in_executor(None, f) for f in function_list]
            return await asyncio.gather(*tasks)

        semaphore = asyncio.Semaphore(max_concurrent)

        async def run_with_semaphore(f):
            async with semaphore:
                return await loop.run_in_executor(None, f)

        tasks = [run_with_semaphore(f) for f in function_list]
        return await asyncio.gather(*tasks)

    return asyncio.run(main())
