import asyncio
from typing import Any
from unittest.mock import patch

import pytest

from pynenc.call import Call
from pynenc.exceptions import InvocationError, RetryError
from pynenc.invocation import (
    DistributedInvocation,
    InvocationStatus,
    InvocationStatusRecord,
)
from pynenc.runner import RunnerContext
from pynenc_tests.conftest import MockPynenc
from pynenc_tests.util import capture_logs

app = MockPynenc()


@app.task
def add(x: int, y: int) -> int:
    return x + y


def test_distributed_invocation_instantiation() -> None:
    invocation = add(1, 2)
    assert isinstance(invocation, DistributedInvocation)
    assert isinstance(invocation.call, Call)
    assert invocation.parent_invocation_id is None


def test_distributed_invocation_to_and_from_dto() -> None:
    # Patch orchestrator to return a real value for get_call_invocation_ids
    app.orchestrator.get_call_invocation_ids.return_value = []
    invocation: DistributedInvocation = add(1, 2)  # type: ignore
    assert invocation == DistributedInvocation.from_dto(
        invocation.to_dto(), invocation.call
    )


def test_get_final_result_exception_not_final() -> None:
    invocation: DistributedInvocation = add(1, 2)  # type: ignore
    app.orchestrator.get_invocation_status_record.return_value = InvocationStatusRecord(
        status=InvocationStatus.REGISTERED
    )
    # Test that it will raise an exception if the invocation is not finished
    with pytest.raises(InvocationError):
        invocation.get_final_result()


@app.task(max_retries=1)
def retry() -> int:
    raise RetryError()


def test_max_retries() -> None:
    app.orchestrator.get_invocation_retries.return_value = 1
    invocation: DistributedInvocation = retry()  # type: ignore
    app.orchestrator.get_invocation_status_record.return_value = InvocationStatusRecord(
        status=InvocationStatus.RUNNING
    )

    with capture_logs(app.logger) as log_buffer:
        with pytest.raises(RetryError):
            invocation.run(RunnerContext.from_runner(app.runner))  # type: ignore
        assert "MAX-RETRY" in log_buffer.getvalue()


def test_reroute_on_running_control() -> None:
    """
    Tests that invocations are rerouted when not authorized to run by concurrency control.
    """
    with (
        patch.object(
            app.orchestrator,
            "is_authorize_to_run_by_concurrency_control",
            return_value=False,
        ) as mock_is_authorized,
        patch.object(
            app.orchestrator, "reroute_invocations"
        ) as mock_reroute_invocations,
    ):
        invocation: DistributedInvocation = add(1, 2)  # type: ignore
        runner_ctx = RunnerContext.from_runner(app.runner)  # type: ignore
        invocation.run(runner_ctx)  # type: ignore
        mock_is_authorized.assert_called_once()
        mock_reroute_invocations.assert_called_once_with(
            {invocation.invocation_id}, runner_ctx
        )


app2 = MockPynenc()


@app2.task
def add2(x: int, y: int) -> int:
    return x + y


@pytest.mark.asyncio
async def test_distributed_async_result_wait_loop(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    mock_base_app = add2.app
    # Create an invocation using the add task (which returns a DistributedInvocation)
    invocation: DistributedInvocation = add2(1, 2)  # type: ignore

    # Simulate status changes: first RUNNING then SUCCESS.
    statuses = [InvocationStatus.RUNNING, InvocationStatus.SUCCESS]

    def status_side_effect(inv: Any) -> InvocationStatusRecord:
        status = statuses.pop(0) if statuses else InvocationStatus.SUCCESS
        return InvocationStatusRecord(status=status)

    mock_base_app.orchestrator.get_invocation_status_record.side_effect = (  # type: ignore
        status_side_effect
    )

    # Patch get_final_result to return a fixed value (e.g. 3)
    monkeypatch.setattr(type(invocation), "get_final_result", lambda self: 3)

    # Patch runner.async_waiting_for_results to simulate a short wait
    async def fake_async_waiting(
        parent: Any, invs: Any, runner_args: Any = None
    ) -> None:
        await asyncio.sleep(0.01)

    monkeypatch.setattr(
        mock_base_app.runner, "async_waiting_for_results", fake_async_waiting
    )

    # Now, calling async_result() should loop until status becomes final, then return 3.
    result = await invocation.async_result()
    assert result == 3


def test_distributed_invocation_get_and_set_state() -> None:
    """
    Test that __getstate__ and __setstate__ correctly serialize and deserialize
    the distributed invocation.
    """
    # Arrange: Create a sample invocation
    invocation: DistributedInvocation = add(1, 2)  # type: ignore

    # Act: Serialize the invocation state
    serialized_state = invocation.__getstate__()

    # Act: Deserialize the state into a new DistributedInvocation
    deserialized_invocation = DistributedInvocation.__new__(DistributedInvocation)
    deserialized_invocation.__setstate__(serialized_state)

    # Assert: Verify that the deserialized object matches the original
    assert deserialized_invocation.invocation_id == invocation.invocation_id
    assert deserialized_invocation.call == invocation.call
    assert (
        deserialized_invocation.parent_invocation_id == invocation.parent_invocation_id
    )
    assert deserialized_invocation.workflow == invocation.workflow
    assert deserialized_invocation._cached_status == invocation._cached_status
    assert deserialized_invocation._cached_status_time == invocation._cached_status_time
    assert deserialized_invocation._stored_in_backend == invocation._stored_in_backend
    assert deserialized_invocation._num_retries == invocation._num_retries
    assert deserialized_invocation._result == invocation._result
