"""Tests for workflow execution functions in tasks.py."""

from __future__ import annotations
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch
from uuid import uuid4
import pytest
from pydantic import BaseModel


@pytest.fixture
def mock_run() -> MagicMock:
    """Create a mock pending run."""
    run = MagicMock()
    run.id = uuid4()
    run.workflow_version_id = uuid4()
    run.status = MagicMock()
    run.status.value = "pending"
    run.input_payload = {"test": "data"}
    run.runnable_config = None
    return run


@pytest.fixture
def mock_version() -> MagicMock:
    """Create a mock workflow version."""
    version = MagicMock()
    version.id = uuid4()
    version.workflow_id = uuid4()
    version.graph = {"nodes": [], "edges": []}
    version.runnable_config = None
    return version


class TestExtractOutput:
    """Tests for _extract_output function."""

    def test_extracts_dict_output(self) -> None:
        """Test that dict final state is wrapped correctly."""
        from orcheo_backend.worker.tasks import _extract_output

        final_state = {"key": "value", "count": 42}
        result = _extract_output(final_state)

        assert result == {"final_state": final_state}

    def test_extracts_pydantic_model_output(self) -> None:
        """Test that Pydantic model is dumped correctly."""
        from orcheo_backend.worker.tasks import _extract_output

        class SampleModel(BaseModel):
            name: str
            value: int

        model = SampleModel(name="test", value=123)
        result = _extract_output(model)

        assert result == {"final_state": {"name": "test", "value": 123}}

    def test_returns_none_for_unknown_type(self) -> None:
        """Test that unknown types return None."""
        from orcheo_backend.worker.tasks import _extract_output

        result = _extract_output("just a string")
        assert result is None

    def test_returns_none_for_none_input(self) -> None:
        """Test that None input returns None."""
        from orcheo_backend.worker.tasks import _extract_output

        result = _extract_output(None)
        assert result is None

    def test_extracts_object_with_model_dump(self) -> None:
        """Test extraction of non-Pydantic object with model_dump method."""
        from orcheo_backend.worker.tasks import _extract_output

        class CustomObject:
            def model_dump(self) -> dict[str, Any]:
                return {"custom": "data"}

        obj = CustomObject()
        result = _extract_output(obj)

        assert result == {"final_state": {"custom": "data"}}


class TestHandleExecutionFailure:
    """Tests for _handle_execution_failure function."""

    @pytest.mark.asyncio
    async def test_marks_run_as_failed(self, mock_run: MagicMock) -> None:
        """Test that run is marked as failed."""
        from orcheo_backend.worker.tasks import _handle_execution_failure

        mock_repo = AsyncMock()
        exception = ValueError("Workflow crashed")

        with patch(
            "orcheo_backend.app.dependencies.get_repository", return_value=mock_repo
        ):
            with patch("orcheo_backend.worker.tasks.logger"):
                result = await _handle_execution_failure(mock_run, exception)

        assert result["status"] == "failed"
        assert "Workflow crashed" in result["error"]
        mock_repo.mark_run_failed.assert_called_once_with(
            mock_run.id, actor="worker", error="Workflow crashed"
        )

    @pytest.mark.asyncio
    async def test_logs_exception_when_mark_failed_fails(
        self, mock_run: MagicMock
    ) -> None:
        """Test that exception is logged if mark_run_failed fails."""
        from orcheo_backend.worker.tasks import _handle_execution_failure

        mock_repo = AsyncMock()
        mock_repo.mark_run_failed = AsyncMock(
            side_effect=RuntimeError("Database error")
        )
        exception = ValueError("Original error")

        with patch(
            "orcheo_backend.app.dependencies.get_repository", return_value=mock_repo
        ):
            with patch("orcheo_backend.worker.tasks.logger") as mock_logger:
                result = await _handle_execution_failure(mock_run, exception)

        # Should still return the original error
        assert result["status"] == "failed"
        assert "Original error" in result["error"]
        # Should log both the original exception and the mark_run_failed failure
        assert mock_logger.exception.call_count == 2

    @pytest.mark.asyncio
    async def test_returns_error_message_from_exception(
        self, mock_run: MagicMock
    ) -> None:
        """Test that error message comes from the exception."""
        from orcheo_backend.worker.tasks import _handle_execution_failure

        mock_repo = AsyncMock()
        exception = RuntimeError("Specific runtime error")

        with patch(
            "orcheo_backend.app.dependencies.get_repository", return_value=mock_repo
        ):
            with patch("orcheo_backend.worker.tasks.logger"):
                result = await _handle_execution_failure(mock_run, exception)

        assert result["error"] == "Specific runtime error"


class TestExecuteWorkflow:
    """Tests for _execute_workflow function."""

    @pytest.mark.asyncio
    async def test_successful_execution(
        self, mock_run: MagicMock, mock_version: MagicMock
    ) -> None:
        """Test successful workflow execution."""
        from orcheo_backend.worker.tasks import _execute_workflow

        mock_repo = AsyncMock()
        mock_repo.get_version = AsyncMock(return_value=mock_version)
        mock_repo.mark_run_succeeded = AsyncMock()

        mock_graph = MagicMock()
        mock_compiled = MagicMock()
        mock_compiled.ainvoke = AsyncMock(return_value={"result": "success"})
        mock_graph.compile = MagicMock(return_value=mock_compiled)

        mock_checkpointer = MagicMock()
        mock_checkpointer.__aenter__ = AsyncMock(return_value=mock_checkpointer)
        mock_checkpointer.__aexit__ = AsyncMock(return_value=None)

        with patch(
            "orcheo_backend.app.dependencies.get_repository", return_value=mock_repo
        ):
            with patch(
                "orcheo_backend.app.dependencies.get_vault",
                return_value=MagicMock(),
            ):
                with patch("orcheo.graph.builder.build_graph", return_value=mock_graph):
                    with patch(
                        "orcheo.persistence.create_checkpointer",
                        return_value=mock_checkpointer,
                    ):
                        with patch(
                            "orcheo_backend.app.workflow_execution._build_initial_state",
                            return_value={},
                        ):
                            with patch("orcheo.config.get_settings"):
                                with patch(
                                    "orcheo.runtime.runnable_config.merge_runnable_configs"
                                ) as mock_merge:
                                    mock_config = MagicMock()
                                    mock_config.to_runnable_config = MagicMock(
                                        return_value={}
                                    )
                                    mock_config.to_state_config = MagicMock(
                                        return_value={}
                                    )
                                    mock_merge.return_value = mock_config

                                    with patch(
                                        "orcheo.runtime.credentials.credential_resolution"
                                    ):
                                        result = await _execute_workflow(mock_run)

        assert result["status"] == "succeeded"
        mock_repo.mark_run_succeeded.assert_called_once()

    @pytest.mark.asyncio
    async def test_handles_execution_exception(
        self, mock_run: MagicMock, mock_version: MagicMock
    ) -> None:
        """Test that execution exceptions are handled."""
        from orcheo_backend.worker.tasks import _execute_workflow

        mock_repo = AsyncMock()
        mock_repo.get_version = AsyncMock(side_effect=RuntimeError("Version not found"))
        mock_repo.mark_run_failed = AsyncMock()

        with patch(
            "orcheo_backend.app.dependencies.get_repository", return_value=mock_repo
        ):
            with patch(
                "orcheo_backend.app.dependencies.get_vault",
                return_value=MagicMock(),
            ):
                with patch("orcheo.config.get_settings"):
                    with patch("orcheo_backend.worker.tasks.logger"):
                        result = await _execute_workflow(mock_run)

        assert result["status"] == "failed"
        assert "Version not found" in result["error"]


class TestExecuteRunAsync:
    """Tests for _execute_run_async function."""

    @pytest.mark.asyncio
    async def test_returns_error_when_load_fails(self) -> None:
        """Test that load failure returns error."""
        from orcheo_backend.worker.tasks import _execute_run_async

        with patch(
            "orcheo_backend.worker.tasks._load_and_validate_run",
            return_value=(None, {"status": "failed", "error": "Not found"}),
        ):
            result = await _execute_run_async(str(uuid4()))

        assert result["status"] == "failed"

    @pytest.mark.asyncio
    async def test_returns_error_when_mark_started_fails(
        self, mock_run: MagicMock
    ) -> None:
        """Test that mark_started failure returns error."""
        from orcheo_backend.worker.tasks import _execute_run_async

        with patch(
            "orcheo_backend.worker.tasks._load_and_validate_run",
            return_value=(mock_run, None),
        ):
            with patch(
                "orcheo_backend.worker.tasks._mark_run_started",
                return_value={"status": "skipped", "reason": "Already started"},
            ):
                result = await _execute_run_async(str(mock_run.id))

        assert result["status"] == "skipped"

    @pytest.mark.asyncio
    async def test_calls_execute_workflow_on_success(self, mock_run: MagicMock) -> None:
        """Test that _execute_workflow is called when validation succeeds."""
        from orcheo_backend.worker.tasks import _execute_run_async

        with patch(
            "orcheo_backend.worker.tasks._load_and_validate_run",
            return_value=(mock_run, None),
        ):
            with patch(
                "orcheo_backend.worker.tasks._mark_run_started",
                return_value=None,
            ):
                with patch(
                    "orcheo_backend.worker.tasks._execute_workflow",
                    return_value={"status": "succeeded"},
                ) as mock_execute:
                    result = await _execute_run_async(str(mock_run.id))

        assert result["status"] == "succeeded"
        mock_execute.assert_called_once_with(mock_run)


class TestDispatchCronTriggersAsync:
    """Tests for _dispatch_cron_triggers_async function."""

    @pytest.mark.asyncio
    async def test_dispatches_due_runs(self) -> None:
        """Test that due cron runs are dispatched."""
        from orcheo_backend.worker.tasks import _dispatch_cron_triggers_async

        mock_run1 = MagicMock()
        mock_run1.id = uuid4()
        mock_run2 = MagicMock()
        mock_run2.id = uuid4()

        mock_repo = AsyncMock()
        mock_repo.dispatch_due_cron_runs = AsyncMock(
            return_value=[mock_run1, mock_run2]
        )

        with patch(
            "orcheo_backend.app.dependencies.get_repository", return_value=mock_repo
        ):
            result = await _dispatch_cron_triggers_async()

        assert len(result) == 2
        assert str(mock_run1.id) in result
        assert str(mock_run2.id) in result

    @pytest.mark.asyncio
    async def test_returns_empty_list_when_no_runs_due(self) -> None:
        """Test that empty list is returned when no runs are due."""
        from orcheo_backend.worker.tasks import _dispatch_cron_triggers_async

        mock_repo = AsyncMock()
        mock_repo.dispatch_due_cron_runs = AsyncMock(return_value=[])

        with patch(
            "orcheo_backend.app.dependencies.get_repository", return_value=mock_repo
        ):
            result = await _dispatch_cron_triggers_async()

        assert result == []

    @pytest.mark.asyncio
    async def test_returns_run_ids(self) -> None:
        """Test that dispatched runs are returned."""
        from orcheo_backend.worker.tasks import _dispatch_cron_triggers_async

        mock_run = MagicMock()
        mock_run.id = uuid4()

        mock_repo = AsyncMock()
        mock_repo.dispatch_due_cron_runs = AsyncMock(return_value=[mock_run])

        with patch(
            "orcheo_backend.app.dependencies.get_repository", return_value=mock_repo
        ):
            result = await _dispatch_cron_triggers_async()

        assert result == [str(mock_run.id)]
