# 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 threading
import time

import hypothesis.strategies as st
import pytest
from aws_xray_sdk import global_sdk_config
from aws_xray_sdk.core import xray_recorder
from aws_xray_sdk.core.context import Context as XRayContext
from hypothesis import given, settings

from czc.concurrency import _capture_xray_entity, async_map


@pytest.mark.unit
@given(st.lists(st.integers()))
@settings(max_examples=100, deadline=1000.00)
def test_async_map(inputs):
    outputs = async_map(lambda x: x, [[x] for x in inputs])
    assert outputs == inputs


@pytest.mark.unit
@given(st.lists(st.integers(), min_size=6, max_size=6))
@settings(max_examples=1, deadline=3000.00)
def test_async_map_slow_task(inputs):
    def slow_task(x):
        time.sleep(1)
        return x

    outputs = async_map(slow_task, [[x] for x in inputs])
    assert outputs == inputs


@pytest.mark.unit
@pytest.mark.parametrize(
    'tasks_run,max_concurrent',
    [
        pytest.param(6, 1, id='single_concurrent_task'),
        pytest.param(6, 2, id='dual_concurrent_tasks'),
        pytest.param(6, 10, id='large_limit_larger_than_concurrency_limit'),
        pytest.param(25, 10, id='large_limit_smaller_than_concurrency_limit'),
    ],
)
def test_max_concurrent_limits_tasks(tasks_run, max_concurrent):
    state = {'running': 0, 'tasks': [], 'snapshot': []}
    lock = threading.Lock()

    def track_concurrent_task(x):
        with lock:
            state['running'] += 1
            state['snapshot'].append(state['running'])
        time.sleep(0.5)
        with lock:
            state['running'] -= 1
            state['tasks'].append(x)

        return x

    inputs = [[x] for x in range(tasks_run)]
    async_map(track_concurrent_task, inputs, max_concurrent=max_concurrent)

    assert state['running'] == 0
    assert set(state['tasks']) == set(range(tasks_run))
    assert max(state['snapshot']) <= max_concurrent


@pytest.mark.unit
@given(st.integers(max_value=0))
def test_max_concurrency_cannot_be_less_than_1(max_concurrent):
    def task(x):
        time.sleep(0.1)
        return x

    inputs = [[x] for x in range(6)]
    with pytest.raises(ValueError):
        async_map(task, inputs, max_concurrent=max_concurrent)


@pytest.fixture
def fresh_xray_context(monkeypatch):
    """Reset xray_recorder to a clean LocalContext for each test, then restore.

    The unit-test suite sets AWS_XRAY_SDK_ENABLED=false (see pytest_unit.ini)
    which would make begin_segment return DummySegment unconditionally. The
    SDK ignores set_sdk_enabled(True) when that env var is present, so we
    delete it for the test and re-enable. Sampling is also disabled so the
    sampler doesn't reach for an unavailable X-Ray daemon.
    """
    monkeypatch.delenv('AWS_XRAY_SDK_ENABLED', raising=False)
    original_context = xray_recorder.context
    original_sampling = xray_recorder.sampling
    original_sdk_enabled = global_sdk_config.sdk_enabled()
    global_sdk_config.set_sdk_enabled(True)
    xray_recorder.configure(context=XRayContext(), sampling=False)
    yield xray_recorder
    xray_recorder.configure(context=original_context, sampling=original_sampling)
    global_sdk_config.set_sdk_enabled(original_sdk_enabled)


@pytest.mark.unit
def test_capture_xray_entity_returns_none_when_no_active_entity(fresh_xray_context):
    assert _capture_xray_entity() is None


@pytest.mark.unit
def test_capture_xray_entity_returns_active_segment(fresh_xray_context):
    segment = xray_recorder.begin_segment('test_segment')
    try:
        captured = _capture_xray_entity()
        assert captured is segment
    finally:
        xray_recorder.end_segment()


@pytest.mark.unit
def test_async_map_propagates_xray_parent(fresh_xray_context):
    """A subsegment created inside an async_map task must be a child of the
    caller's active segment (or subsegment), not orphaned to the root."""
    captured_parent_ids = []

    def task(x):
        sub = xray_recorder.begin_subsegment(f'task_{x}')
        assert sub is not None
        try:
            captured_parent_ids.append(sub.parent_id)
            return x
        finally:
            xray_recorder.end_subsegment()

    segment = xray_recorder.begin_segment('caller')
    try:
        async_map(task, [[1], [2], [3]])
    finally:
        xray_recorder.end_segment()

    assert captured_parent_ids == [segment.id, segment.id, segment.id]


@pytest.mark.unit
def test_async_map_propagates_xray_parent_with_max_concurrent(fresh_xray_context):
    """Same propagation guarantee when max_concurrent is set (semaphore-gated path)."""
    captured_parent_ids = []

    def task(x):
        sub = xray_recorder.begin_subsegment(f'task_{x}')
        assert sub is not None
        try:
            captured_parent_ids.append(sub.parent_id)
            return x
        finally:
            xray_recorder.end_subsegment()

    segment = xray_recorder.begin_segment('caller')
    try:
        async_map(task, [[1], [2], [3], [4]], max_concurrent=2)
    finally:
        xray_recorder.end_segment()

    assert captured_parent_ids == [segment.id] * 4


@pytest.mark.unit
def test_async_map_exception_does_not_corrupt_caller_entity(fresh_xray_context):
    """A worker raising must not leak state into the caller's thread-local."""

    def boom(x):
        if x == 0:
            raise RuntimeError('boom')
        return x

    segment = xray_recorder.begin_segment('caller')
    try:
        with pytest.raises(RuntimeError, match='boom'):
            async_map(boom, [[0]])
        assert xray_recorder.get_trace_entity() is segment
    finally:
        xray_recorder.end_segment()
