import asyncio
import os
import platform
import queue
import signal
import sys
import threading
import time
from typing import Any, Iterator, Sequence, Union, overload
from unittest.mock import AsyncMock, Mock

import nest_asyncio  # type: ignore[import-untyped]
import pytest

from phoenix.client.utils.executors import (
    AsyncExecutor,
    ExecutionStatus,
    SyncExecutor,
    get_executor_on_sync_context,
)

# AsyncExecutor tests


async def test_async_executor_executes() -> None:
    async def dummy_fn(payload: int) -> int:
        return payload - 1

    executor = AsyncExecutor(dummy_fn, concurrency=10, max_retries=0)
    inputs = [1, 2, 3, 4, 5]
    outputs, _ = await executor.execute(inputs)
    assert outputs == [0, 1, 2, 3, 4]


async def test_async_executor_executes_many_tasks() -> None:
    async def dummy_fn(payload: int) -> int:
        return payload

    executor = AsyncExecutor(dummy_fn, concurrency=10, max_retries=0)
    inputs = [x for x in range(100)]
    outputs, _ = await executor.execute(inputs)
    assert outputs == inputs


def test_async_executor_runs_synchronously() -> None:
    async def dummy_fn(payload: int) -> int:
        return payload - 2

    executor = AsyncExecutor(dummy_fn, concurrency=10, max_retries=0)
    inputs = [1, 2, 3, 4, 5]
    outputs, _ = executor.run(inputs)
    assert outputs == [-1, 0, 1, 2, 3]


async def test_async_executor_execute_exits_early_on_error() -> None:
    async def dummy_fn(payload: int) -> int:
        if payload == 3:
            raise ValueError("test error")
        return payload - 1

    executor = AsyncExecutor(
        dummy_fn, concurrency=1, max_retries=0, exit_on_error=True, fallback_return_value=52
    )
    inputs = [1, 2, 3, 4, 5]
    outputs, _ = await executor.execute(inputs)
    assert outputs == [0, 1, 52, 52, 52]


def test_async_executor_run_exits_early_on_error() -> None:
    async def dummy_fn(payload: int) -> int:
        if payload == 3:
            raise ValueError("test error")
        return payload - 1

    executor = AsyncExecutor(
        dummy_fn, concurrency=1, max_retries=0, exit_on_error=True, fallback_return_value=52
    )
    inputs = [1, 2, 3, 4, 5]
    outputs, statuses = executor.run(inputs)
    exceptions = [status.exceptions for status in statuses]
    status_types = [status.status for status in statuses]
    assert outputs == [0, 1, 52, 52, 52]
    assert [len(excs) if excs else 0 for excs in exceptions] == [
        0,
        0,
        1,
        0,
        0,
    ], "one exception raised, then exits"
    assert status_types == [
        ExecutionStatus.COMPLETED,
        ExecutionStatus.COMPLETED,
        ExecutionStatus.FAILED,
        ExecutionStatus.DID_NOT_RUN,
        ExecutionStatus.DID_NOT_RUN,
    ]
    assert all(isinstance(exc, ValueError) for exc in exceptions[2])


async def test_async_executor_can_continue_on_error() -> None:
    async def dummy_fn(payload: int) -> int:
        if payload == 3:
            raise ValueError("test error")
        return payload - 1

    executor = AsyncExecutor(
        dummy_fn, concurrency=1, max_retries=1, exit_on_error=False, fallback_return_value=52
    )
    inputs = [1, 2, 3, 4, 5]
    outputs, statuses = await executor.execute(inputs)
    exceptions = [status.exceptions for status in statuses]
    status_types = [status.status for status in statuses]
    execution_times = [status.execution_seconds for status in statuses]
    assert outputs == [0, 1, 52, 3, 4], "failed tasks use the fallback value"
    assert [len(excs) if excs else 0 for excs in exceptions] == [
        0,
        0,
        2,
        0,
        0,
    ], "two exceptions due to retries"
    assert status_types == [
        ExecutionStatus.COMPLETED,
        ExecutionStatus.COMPLETED,
        ExecutionStatus.FAILED,
        ExecutionStatus.COMPLETED,
        ExecutionStatus.COMPLETED,
    ]
    assert len(execution_times) == 5
    assert all(isinstance(runtime, float) for runtime in execution_times)
    assert all(isinstance(exc, ValueError) for exc in exceptions[2])


async def test_async_executor_marks_completed_with_retries_status() -> None:
    retry_counter = 0

    async def dummy_fn(payload: int) -> int:
        if payload == 3:
            nonlocal retry_counter
            if retry_counter < 2:
                retry_counter += 1
                raise ValueError("test error")
        return payload - 1

    executor = AsyncExecutor(
        dummy_fn, concurrency=1, max_retries=3, exit_on_error=False, fallback_return_value=52
    )
    inputs = [1, 2, 3, 4, 5]
    outputs, execution_details = await executor.execute(inputs)
    assert outputs == [0, 1, 2, 3, 4], "input 3 should only fail twice"
    assert [status.status for status in execution_details] == [
        ExecutionStatus.COMPLETED,
        ExecutionStatus.COMPLETED,
        ExecutionStatus.COMPLETED_WITH_RETRIES,
        ExecutionStatus.COMPLETED,
        ExecutionStatus.COMPLETED,
    ]


class InterruptingIterator(Sequence[int]):
    def __init__(self, interruption_index: int, max_elements: int):
        self.interruption_index = interruption_index
        self.max_elements = max_elements
        self.current = 0

    def __len__(self) -> int:
        return self.max_elements

    @overload
    def __getitem__(self, index: int) -> int: ...

    @overload
    def __getitem__(self, index: slice) -> Sequence[int]: ...

    def __getitem__(self, index: Union[int, slice]) -> Union[int, Sequence[int]]:
        if isinstance(index, slice):
            start, stop, step = index.indices(self.max_elements)
            return [i for i in range(start, stop, step)]
        if index < 0:
            index = self.max_elements + index
        if index < 0 or index >= self.max_elements:
            raise IndexError("Index out of range")
        return index

    def __iter__(self) -> Iterator[int]:
        return self

    def __next__(self) -> int:
        if self.current < self.max_elements:
            if self.current == self.interruption_index:
                # Trigger interruption signal
                os.kill(os.getpid(), signal.SIGUSR1)  # type: ignore[attr-defined, unused-ignore]
                time.sleep(0.1)

            res = self.current
            self.current += 1
            return res
        else:
            raise StopIteration


@pytest.mark.skipif(platform.system() == "Windows", reason="SIGUSR1 not supported on Windows")
async def test_async_executor_sigint_handling() -> None:
    async def async_fn(x: int) -> int:
        await asyncio.sleep(0.01)
        return x

    result_length = 1000
    sigint_index = 50
    executor = AsyncExecutor(
        async_fn,
        concurrency=5,
        max_retries=0,
        fallback_return_value="test",
        termination_signal=signal.SIGUSR1,  # type: ignore[attr-defined, unused-ignore]
    )
    task = asyncio.create_task(executor.execute(InterruptingIterator(sigint_index, result_length)))

    results, _ = await task
    assert len(results) == result_length
    assert results.count("test") > 100, "most inputs should not have been processed"


async def test_async_executor_retries() -> None:
    mock_generate = AsyncMock(side_effect=RuntimeError("Test exception"))
    executor = AsyncExecutor(mock_generate, max_retries=3)

    await executor.execute([1])  # by default the executor does not raise on generation errors

    assert mock_generate.call_count == 4, "1 initial call + 3 retries"


# SyncExecutor tests


def test_sync_executor_runs_many_tasks() -> None:
    def dummy_fn(payload: int) -> int:
        return payload

    executor = SyncExecutor(dummy_fn, max_retries=0)
    inputs = [x for x in range(1000)]
    outputs, _ = executor.run(inputs)
    assert outputs == inputs


def test_sync_executor_runs_once_per_task() -> None:
    dummy_fn = Mock(side_effect=lambda x: x)  # pyright: ignore[reportUnknownLambdaType, reportUnknownArgumentType]

    executor = SyncExecutor(dummy_fn, max_retries=3)
    inputs = [x for x in range(10)]
    executor.run(inputs)
    assert dummy_fn.call_count == 10


def test_sync_executor_runs() -> None:
    def dummy_fn(payload: int) -> int:
        return payload - 2

    executor = SyncExecutor(dummy_fn, max_retries=0)
    inputs = [1, 2, 3, 4, 5]
    outputs, _ = executor.run(inputs)
    assert outputs == [-1, 0, 1, 2, 3]


def test_sync_executor_run_exits_early_on_error() -> None:
    def dummy_fn(payload: int) -> int:
        if payload == 3:
            raise ValueError("test error")
        return payload - 1

    executor = SyncExecutor(dummy_fn, exit_on_error=True, fallback_return_value=52, max_retries=0)
    inputs = [1, 2, 3, 4, 5]
    outputs, execution_details = executor.run(inputs)
    exceptions = [status.exceptions for status in execution_details]
    status_types = [status.status for status in execution_details]
    assert outputs == [0, 1, 52, 52, 52]
    assert [len(excs) if excs else 0 for excs in exceptions] == [
        0,
        0,
        1,
        0,
        0,
    ], "one exception raised, then exits"
    assert status_types == [
        ExecutionStatus.COMPLETED,
        ExecutionStatus.COMPLETED,
        ExecutionStatus.FAILED,
        ExecutionStatus.DID_NOT_RUN,
        ExecutionStatus.DID_NOT_RUN,
    ]
    assert all(isinstance(exc, ValueError) for exc in exceptions[2])


def test_sync_executor_can_continue_on_error() -> None:
    def dummy_fn(payload: int) -> int:
        if payload == 3:
            raise ValueError("test error")
        return payload - 1

    executor = SyncExecutor(dummy_fn, exit_on_error=False, fallback_return_value=52, max_retries=1)
    inputs = [1, 2, 3, 4, 5]
    outputs, execution_details = executor.run(inputs)
    exceptions = [status.exceptions for status in execution_details]
    status_types = [status.status for status in execution_details]
    execution_times = [status.execution_seconds for status in execution_details]
    assert outputs == [0, 1, 52, 3, 4]
    assert [len(excs) if excs else 0 for excs in exceptions] == [
        0,
        0,
        2,
        0,
        0,
    ], "two exceptions due to retries"
    assert status_types == [
        ExecutionStatus.COMPLETED,
        ExecutionStatus.COMPLETED,
        ExecutionStatus.FAILED,
        ExecutionStatus.COMPLETED,
        ExecutionStatus.COMPLETED,
    ]
    assert len(execution_times) == 5
    assert all(isinstance(runtime, float) for runtime in execution_times)
    assert all(isinstance(exc, ValueError) for exc in exceptions[2])


def test_sync_executor_marks_completed_with_retries_status() -> None:
    retry_counter = 0

    def dummy_fn(payload: int) -> int:
        if payload == 3:
            nonlocal retry_counter
            if retry_counter < 2:
                retry_counter += 1
                raise ValueError("test error")
        return payload - 1

    executor = SyncExecutor(dummy_fn, max_retries=3, exit_on_error=False, fallback_return_value=52)
    inputs = [1, 2, 3, 4, 5]
    outputs, execution_details = executor.run(inputs)
    assert outputs == [0, 1, 2, 3, 4], "input 3 should only fail twice"
    assert [status.status for status in execution_details] == [
        ExecutionStatus.COMPLETED,
        ExecutionStatus.COMPLETED,
        ExecutionStatus.COMPLETED_WITH_RETRIES,
        ExecutionStatus.COMPLETED,
        ExecutionStatus.COMPLETED,
    ]


@pytest.mark.skipif(platform.system() == "Windows", reason="SIGUSR1 not supported on Windows")
def test_sync_executor_sigint_handling() -> None:
    def sync_fn(x: int) -> int:
        time.sleep(0.01)
        return x

    result_length = 1000
    sigint_index = 50
    executor = SyncExecutor(
        sync_fn,
        max_retries=0,
        fallback_return_value="test",
        termination_signal=signal.SIGUSR1,  # type: ignore[attr-defined, unused-ignore]
    )
    results, _ = executor.run(InterruptingIterator(sigint_index, result_length))
    assert len(results) == result_length
    assert results.count("test") > 100, "most inputs should not have been processed"


def test_sync_executor_defaults_sigint_handling() -> None:
    def sync_fn(x: int) -> Any:
        return signal.getsignal(signal.SIGINT)

    executor = SyncExecutor(
        sync_fn,
        max_retries=0,
        fallback_return_value="test",
    )
    res, _ = executor.run(["test"])
    assert res[0] != signal.default_int_handler


def test_sync_executor_bypasses_sigint_handling_if_none() -> None:
    def sync_fn(x: int) -> Any:
        return signal.getsignal(signal.SIGINT)

    executor = SyncExecutor(
        sync_fn,
        max_retries=0,
        fallback_return_value="test",
        termination_signal=None,
    )
    res, _ = executor.run(["test"])
    assert res[0] == signal.default_int_handler


def test_sync_executor_retries() -> None:
    mock_generate = Mock(side_effect=RuntimeError("Test exception"))
    executor = SyncExecutor(mock_generate, max_retries=3)

    executor.run([1])  # by default the executor does not raise on generation errors

    assert mock_generate.call_count == 4, "1 initial call + 3 retries"


# test executor factory


@pytest.mark.xfail(reason="This test has started failing, marking for further investigation")
async def test_executor_factory_returns_sync_in_async_context() -> None:
    def sync_fn(x: Any) -> Any:
        return x

    async def async_fn(x: Any) -> Any:
        return x

    async def executor_in_async_context() -> Any:
        return get_executor_on_sync_context(sync_fn, async_fn)

    executor = await executor_in_async_context()
    assert isinstance(executor, SyncExecutor)


@pytest.mark.skipif(
    sys.version_info >= (3, 14),
    reason="nest_asyncio is incompatible with Python 3.14's C-level asyncio.current_task",
)
async def test_executor_factory_returns_async_in_patched_async_context() -> None:
    nest_asyncio.apply()  # pyright: ignore

    def sync_fn(x: Any) -> Any:
        return x

    async def async_fn(x: Any) -> Any:
        return x

    async def executor_in_async_context() -> Any:
        return get_executor_on_sync_context(sync_fn, async_fn)

    executor = await executor_in_async_context()
    assert isinstance(executor, AsyncExecutor)


def test_executor_factory_returns_async_in_sync_context() -> None:
    def sync_fn(x: Any) -> Any:
        return x

    async def async_fn(x: Any) -> Any:
        return x

    def executor_in_sync_context() -> Any:
        return get_executor_on_sync_context(sync_fn, async_fn)

    executor = executor_in_sync_context()
    assert isinstance(executor, AsyncExecutor)


def test_executor_factory_returns_sync_in_sync_context_if_asked() -> None:
    def sync_fn(x: Any) -> Any:
        return x

    async def async_fn(x: Any) -> Any:
        return x

    def executor_in_sync_context() -> Any:
        return get_executor_on_sync_context(
            sync_fn,
            async_fn,
            run_sync=True,  # request a sync_executor
        )

    executor = executor_in_sync_context()
    assert isinstance(executor, SyncExecutor)


def test_executor_factory_returns_sync_in_threads() -> None:
    def sync_fn(x: Any) -> Any:
        return x

    async def async_fn(x: Any) -> Any:
        return x

    exception_log: queue.Queue[Exception] = queue.Queue()

    def run_test() -> None:
        try:
            executor = get_executor_on_sync_context(
                sync_fn,
                async_fn,
                run_sync=True,  # request a sync_executor
            )
            assert isinstance(executor, SyncExecutor)
            assert executor.termination_signal is None
        except Exception as e:
            exception_log.put(e)

    test_thread = threading.Thread(target=run_test)
    test_thread.start()
    test_thread.join()
    if not exception_log.empty():
        raise exception_log.get()


@pytest.mark.skipif(
    sys.version_info >= (3, 14),
    reason="nest_asyncio is incompatible with Python 3.14's C-level asyncio.current_task",
)
async def test_executor_factory_returns_sync_in_threads_even_if_async_context() -> None:
    def sync_fn(x: Any) -> Any:
        return x

    async def async_fn(x: Any) -> Any:
        return x

    exception_log: queue.Queue[Exception] = queue.Queue()

    async def run_test() -> None:
        nest_asyncio.apply()  # pyright: ignore
        try:
            executor = get_executor_on_sync_context(
                sync_fn,
                async_fn,
            )
            assert isinstance(executor, SyncExecutor)
            assert executor.termination_signal is None
        except Exception as e:
            exception_log.put(e)

    def async_task(loop: asyncio.AbstractEventLoop) -> None:
        asyncio.set_event_loop(loop)
        loop.run_until_complete(run_test())

    loop = asyncio.new_event_loop()
    test_thread = threading.Thread(target=async_task, args=(loop,))
    test_thread.start()
    test_thread.join()

    if not exception_log.empty():
        raise exception_log.get()


@pytest.mark.skipif(
    sys.version_info >= (3, 14),
    reason="nest_asyncio is incompatible with Python 3.14's C-level asyncio.current_task",
)
def test_executor_factory_returns_async_not_in_thread_if_async_context() -> None:
    def sync_fn(x: Any) -> Any:
        return x

    async def async_fn(x: Any) -> Any:
        return x

    exception_log: queue.Queue[Exception] = queue.Queue()

    async def run_test() -> None:
        nest_asyncio.apply()  # pyright: ignore
        try:
            executor = get_executor_on_sync_context(
                sync_fn,
                async_fn,
            )
            assert isinstance(executor, AsyncExecutor)
            assert executor.termination_signal is not None
        except Exception as e:
            exception_log.put(e)

    def async_task() -> None:
        asyncio.run(run_test())

    async_task()

    if not exception_log.empty():
        raise exception_log.get()


@pytest.mark.skipif(platform.system() == "Windows", reason="SIGUSR1 not supported on Windows")
def test_sync_executor_run_works_in_background_thread() -> None:
    def sync_fn(x: int) -> int:
        return x * 2

    outputs = []
    errors: list[Exception] = []

    def run_in_background() -> None:
        try:
            executor = SyncExecutor(sync_fn, termination_signal=signal.SIGUSR1)  # type: ignore[attr-defined, unused-ignore]
            result, _ = executor.run(range(3))
            outputs.extend(result)  # pyright: ignore[reportUnknownMemberType]
        except Exception as e:
            errors.append(e)

    test_thread = threading.Thread(target=run_in_background)
    test_thread.start()
    test_thread.join()

    assert not errors, f"run() failed in background thread: {errors}"
    assert outputs == [0, 2, 4], f"Expected [0, 2, 4], got {outputs}"


@pytest.mark.skipif(platform.system() == "Windows", reason="SIGUSR1 not supported on Windows")
def test_async_executor_run_works_in_background_thread() -> None:
    async def async_fn(x: int) -> int:
        return x * 3

    outputs = []
    errors: list[Exception] = []

    def run_in_background() -> None:
        try:
            executor = AsyncExecutor(async_fn, termination_signal=signal.SIGUSR1)  # type: ignore[attr-defined, unused-ignore]
            result, _ = executor.run(range(3))
            outputs.extend(result)  # pyright: ignore[reportUnknownMemberType]
        except Exception as e:
            errors.append(e)

    test_thread = threading.Thread(target=run_in_background)
    test_thread.start()
    test_thread.join()

    assert not errors, f"run() failed in background thread: {errors}"
    assert outputs == [0, 3, 6], f"Expected [0, 3, 6], got {outputs}"
