import asyncio
import os
import threading
from typing import Any, cast
from acelery.errors import MaxPendingInterruptsError
from acelery.executable import AIOExecutable
from acelery.result import AIOResult
import pytest
from pytest_mock import MockType, MockerFixture
import logging
import contextlib
from contextlib import contextmanager
from typing import Literal
from collections.abc import Iterator
from celery import Celery, Task
from celery.contrib.testing.worker import TestWorkController, setup_app_for_worker  # type: ignore[import-untyped]
from celery.utils.nodenames import anon_nodename  # type: ignore[import-untyped]
from celery.worker import WorkController
from celery.worker import state as worker_state  # type: ignore[import-untyped]
from acelery.task import async_task


logger = logging.getLogger(__name__)
PoolTypes = Literal["solo"] | Literal["prefork"]


app = Celery("tests")
app.conf.update(
    broker_url=os.getenv("CELERY_BROKER_URL", "amqp://guest:guest@localhost:5672//"),
    result_backend=os.getenv("CELERY_RESULT_BACKEND", "redis://localhost:6379/"),
)


@app.task()
@async_task()
async def hello_world():
    return "hello world"


@app.task()
@async_task()
async def residue_coroutine_check():
    t = asyncio.create_task(asyncio.sleep(5), name="sleep-task")  # noqa
    return "hello world"


@app.task(bind=True)
@async_task()
async def retying_task(task: Task):
    if task.request.retries == 0:
        task.retry(exc=Exception("Retrying task"))
    return "hello world"


@app.task(track_started=True)
@async_task()
async def sleeping_task(sleep_time: float) -> str:
    await asyncio.sleep(sleep_time)
    return "OK"


@app.task(track_started=False)
@async_task()
async def untracked_started_sleeping_task(sleep_time: float) -> str:
    await asyncio.sleep(sleep_time)
    return "OK"


@contextmanager
def start_test_worker(
    app: Celery,
    *,
    pool: PoolTypes = "solo",
    timeout: float = 30,
    queues: list[str] | None = None,
) -> Iterator[WorkController]:
    """
    Start a test Celery worker in a daemon thread.

    Unlike celery's built-in ``start_worker``, this implementation adds a
    timeout to the startup wait (celery's ``ensure_started`` blocks forever)
    and guarantees ``worker_state.should_terminate`` is reset on exit so that
    consecutive workers within the same process can start cleanly.

    :param app: Celery app to start the worker for.
    :param pool: Pool type to use ("solo" or "prefork").
    :param timeout: Timeout in seconds for startup and shutdown.
    :param queues: Optional list of queues for the worker to consume from.
    :return: WorkController instance.
    """
    name = app.main  # type: ignore[attr-defined]
    logger.info("Starting Celery worker for %s with pool type %s", f"{name!r}", pool)
    logger.debug("Available tasks: %s", list(app.tasks.keys()))

    setup_app_for_worker(app, "ERROR", None)

    # Purge leftover tasks from previous tests so the new worker doesn't
    # waste time processing stale work (e.g. a 5s sleep task that was
    # requeued when the previous worker's transport was force-closed).
    with app.connection_for_write() as conn:
        default_queue = app.conf.task_default_queue or "celery"
        with contextlib.suppress(Exception):
            conn.default_channel.queue_purge(default_queue)

    worker = TestWorkController(
        app=app,
        concurrency=1,
        hostname=anon_nodename(),
        pool=pool,
        loglevel="ERROR",
        ready_callback=None,
        without_heartbeat=True,
        without_mingle=True,
        without_gossip=True,
        queues=queues,
    )

    t = threading.Thread(target=worker.start, daemon=True)
    t.start()

    # Celery's ensure_started() has no timeout — use the Event directly
    if not worker._on_started.wait(timeout=timeout):
        worker.stop()  # type: ignore[attr-defined]
        t.join(timeout=5)
        raise RuntimeError(f"Worker failed to start within {timeout}s")

    check_worker_ping(app, worker.hostname, timeout=timeout)  # type: ignore[attr-defined]
    try:
        yield worker
    finally:
        logger.info(f"Stopping Celery worker for {name}")
        # Signal the worker thread to self-terminate and force-close the
        # underlying AMQP transport socket.  Both worker.stop() and the
        # worker's own graceful shutdown path call basic_cancel →
        # drain_events(timeout=None), which blocks on a socket read
        # indefinitely.  Closing the transport crashes that read and lets
        # the thread exit.
        worker_state.should_terminate = 0
        with contextlib.suppress(Exception):
            worker.consumer.connection.transport.close()  # type: ignore[union-attr]
        t.join(timeout)
        # Always reset so the next worker in this process can start
        worker_state.should_terminate = None


def check_worker_ping(app: Celery, worker_hostname: str, timeout: float = 1) -> None:
    resp = app.control.ping([worker_hostname], timeout=timeout)
    if resp is None or len(resp) != 1:
        raise RuntimeError(f"Failed to ping worker, received invalid response: {resp}")

    value = resp[0].get(worker_hostname, {}).get("ok")
    if value != "pong":
        raise RuntimeError(f"Failed to ping worker, received invalid response: {resp}")


@pytest.fixture(scope="function")
def celery_test_worker():
    with start_test_worker(app, pool="solo") as worker:
        yield worker


@pytest.mark.usefixtures("celery_test_worker")
async def test_asyncio_task_get_timeout():
    with pytest.raises(TimeoutError):
        aio_result = await AIOExecutable(
            sleeping_task.si(5), return_type=str
        ).apply_async()
        result = await aio_result.get(timeout=1)
        assert result == "OK"


@pytest.mark.usefixtures("celery_test_worker")
async def test_asyncio_task_get_success():
    aio_result = await AIOExecutable(sleeping_task.si(0), return_type=str).apply_async()
    result = await aio_result.get(timeout=5)
    assert result == "OK"


async def _await_spy_called(mocker: MockType, *, timeout: float | None = None) -> None:
    async def _await_spy_called_impl():
        while not mocker.called:
            await asyncio.sleep(0.01)

    if timeout:
        await asyncio.wait_for(_await_spy_called_impl(), timeout=timeout)
    else:
        await _await_spy_called_impl()


@pytest.mark.usefixtures("celery_test_worker")
async def test_asyncio_task_trigger_or_join_task_success_no_concurrent():
    _, result = await AIOResult.trigger_or_join_task(
        sleeping_task.si(0),
        timeout=1,
        return_type=str,
    )
    assert result == "OK"


@pytest.mark.usefixtures("celery_test_worker")
async def test_asyncio_task_trigger_or_join_task_success_concurrent_untracked_started(
    mocker: MockerFixture,
):
    send_task_spy = mocker.spy(app, "send_task")
    workflow = untracked_started_sleeping_task.si(2)
    task1: asyncio.Task[tuple[AIOResult[str], str]] = asyncio.create_task(
        AIOResult.trigger_or_join_task(
            workflow, timeout=5, return_type=str, track_started=False
        )
    )

    await _await_spy_called(send_task_spy, timeout=0.1)
    # wait for the task to be started
    await asyncio.sleep(0.1)
    task2: asyncio.Task[tuple[AIOResult[str], str]] = asyncio.create_task(
        AIOResult.trigger_or_join_task(
            workflow, timeout=5, return_type=str, track_started=False
        )
    )
    results_tuples = await asyncio.gather(task1, task2)
    results = list(map(lambda result: result[1], results_tuples))
    send_task_spy.assert_called_once()
    assert len(results) == 2
    assert all(result == "OK" for result in results)


@pytest.mark.usefixtures("celery_test_worker")
async def test_asyncio_task_trigger_or_join_task_success_concurrent_tracked_started(
    mocker: MockerFixture,
):
    send_task_spy = mocker.spy(app, "send_task")
    workflow = sleeping_task.si(2)
    task1: asyncio.Task[tuple[AIOResult[str], str]] = asyncio.create_task(
        AIOResult.trigger_or_join_task(
            workflow, timeout=5, return_type=str, track_started=True
        )
    )
    await _await_spy_called(send_task_spy, timeout=0.1)
    # wait for the task to be started
    await asyncio.sleep(0.1)
    task2: asyncio.Task[tuple[AIOResult[str], str]] = asyncio.create_task(
        AIOResult.trigger_or_join_task(
            workflow, timeout=5, return_type=str, track_started=True
        )
    )
    results_tuples = await asyncio.gather(task1, task2)
    results = list(map(lambda result: result[1], results_tuples))
    send_task_spy.assert_called_once()
    assert len(results) == 2
    assert all(result == "OK" for result in results)


@pytest.mark.usefixtures("celery_test_worker")
async def test_asyncio_task_trigger_or_join_task_success_at_same_time(
    mocker: MockerFixture,
):
    send_task_spy = mocker.spy(app, "send_task")
    workflow = sleeping_task.si(1)
    gathered_results: Any = await asyncio.gather(
        AIOResult.trigger_or_join_task(
            workflow,
            timeout=10,
            interrupt_every=0.5,
            retry_on_timeout=False,
            return_type=str,
            max_pending_interrupts=10,
        ),
        AIOResult.trigger_or_join_task(
            workflow,
            timeout=10,
            interrupt_every=0.5,
            retry_on_timeout=False,
            return_type=str,
            max_pending_interrupts=10,
        ),
        return_exceptions=True,
    )

    results_tuples = cast(
        list[tuple[AIOResult[str], str] | Exception],
        list(gathered_results),
    )
    # there is a chance for only one trigger if one of the tasks was fast enough to publish
    # the task before the other was able to get the unique running task key from
    # the redis backend
    assert send_task_spy.call_count in (1, 2)
    assert len(results_tuples) == 2
    exceptions: list[Exception] = cast(
        list[Exception],
        list(filter(lambda result: isinstance(result, Exception), results_tuples)),
    )
    if exceptions:
        raise exceptions[0]
    valid_result_tuples = cast(list[tuple[AIOResult[str], str]], results_tuples)
    results = list(map(lambda x: x[1], valid_result_tuples))
    assert all(result == "OK" for result in results)


@pytest.mark.usefixtures("celery_test_worker")
async def test_asyncio_task_trigger_or_join_task_retries_on_timeout(
    mocker: MockerFixture,
):
    send_task_mock = mocker.spy(app, "send_task")
    with pytest.raises(TimeoutError):
        await AIOResult.trigger_or_join_task(
            sleeping_task.si(2),
            timeout=1,
            max_retries_on_timeout=3,
            return_type=str,
        )
    assert send_task_mock.call_count == 3


async def test_asyncio_task_trigger_or_join_task_retries_on_max_pending_interrupts(
    mocker: MockerFixture,
):
    send_task_mock = mocker.spy(app, "send_task")
    with pytest.raises(MaxPendingInterruptsError):
        await AIOResult.trigger_or_join_task(
            sleeping_task.si(0),
            interrupt_every=0.5,
            max_pending_interrupts=3,
            # timeout to be greater than the interrupt_every * max_pending_interrupts
            # to make sure we get the max pending interrupts here
            timeout=5,
            return_type=str,
        )
    assert send_task_mock.call_count == 3
