from __future__ import annotations

import asyncio
from concurrent.futures import ThreadPoolExecutor
from time import sleep
from typing import Any

import pytest

from cmvr_edge_ai.workers import (
    PersistentProcessWorker,
    ProcessWorkerError,
    RemoteWorkerError,
)


def _worker_handler(payload: Any) -> Any:
    if payload == "explode":
        raise RuntimeError("intentional remote failure")
    if payload == "slow":
        sleep(0.3)
        return "late-result"
    return payload * 2


def _worker_factory():  # type: ignore[no-untyped-def]
    """Top-level factory so the spawn start method can import and pickle it."""

    return _worker_handler


async def _keep_restricted_event_loop_responsive() -> None:
    """Provide wakeups where sandboxed self-pipe notifications are unavailable."""

    while True:
        await asyncio.sleep(0.01)


def test_spawn_worker_returns_results_and_surfaces_remote_errors() -> None:
    async def exercise():  # type: ignore[no-untyped-def]
        loop = asyncio.get_running_loop()
        host_executor = ThreadPoolExecutor(max_workers=2)
        loop.set_default_executor(host_executor)
        ticker = asyncio.create_task(_keep_restricted_event_loop_responsive())
        worker = PersistentProcessWorker(
            _worker_factory,
            start_method="spawn",
            queue_capacity=1,
        )
        try:
            await worker.start()
            assert await worker.submit(21, timeout_s=10) == 42
            with pytest.raises(RemoteWorkerError) as captured:
                await worker.submit("explode", timeout_s=10)
            assert captured.value.remote_type == "RuntimeError"
            assert "intentional remote failure" in str(captured.value)
        finally:
            await worker.stop(timeout_s=5)
            ticker.cancel()
            await asyncio.gather(ticker, return_exceptions=True)
            host_executor.shutdown(wait=True, cancel_futures=True)
            loop._default_executor = None  # type: ignore[attr-defined]
        return worker

    worker = asyncio.run(exercise())

    assert worker.is_alive is False


def test_timed_out_request_poisons_worker_and_rejects_later_submits() -> None:
    async def exercise():  # type: ignore[no-untyped-def]
        loop = asyncio.get_running_loop()
        host_executor = ThreadPoolExecutor(max_workers=2)
        loop.set_default_executor(host_executor)
        ticker = asyncio.create_task(_keep_restricted_event_loop_responsive())
        worker = PersistentProcessWorker(
            _worker_factory,
            start_method="spawn",
            queue_capacity=1,
        )
        try:
            await worker.start()
            with pytest.raises(TimeoutError, match="timed out"):
                await worker.submit("slow", timeout_s=0.05)
            assert worker.is_alive is True
            with pytest.raises(ProcessWorkerError, match="worker is poisoned"):
                await worker.submit(22, timeout_s=1)
        finally:
            await worker.stop(timeout_s=2)
            ticker.cancel()
            await asyncio.gather(ticker, return_exceptions=True)
            host_executor.shutdown(wait=True, cancel_futures=True)
            loop._default_executor = None  # type: ignore[attr-defined]
        return worker

    worker = asyncio.run(exercise())

    assert worker.is_alive is False