from contextlib import ExitStack

import pytest
from temporalio.testing import WorkflowEnvironment

import mistralai.workflows as workflows
from mistralai.workflows.core.config import config

from .utils import create_test_worker


@workflows.workflow.define(name="nondeterministic")
class NonDeterministicWorkflow:
    import os

    _pid = os.getpid()  # nondeterministic call at class body level

    @workflows.workflow.entrypoint
    async def run(self) -> str:
        return ""


@workflows.workflow.define(name="nondeterministic-enforced", enforce_determinism=True)
class NonDeterministicWorkflowEnforced:
    import os

    _pid = os.getpid()

    @workflows.workflow.entrypoint
    async def run(self) -> str:
        return ""


@workflows.workflow.define(name="nondeterministic-not-enforced", enforce_determinism=False)
class NonDeterministicWorkflowNotEnforced:
    import os

    _pid = os.getpid()

    @workflows.workflow.entrypoint
    async def run(self) -> str:
        return ""


class TestEnforceDeterminism:
    @pytest.mark.asyncio
    async def test_nondeterministic_workflow(self, temporal_env: WorkflowEnvironment) -> None:
        """Determinism enforcement follows the config default."""
        with ExitStack() as context:
            if config.worker.default_enforce_determinism:
                context.enter_context(pytest.raises(RuntimeError))
            async with create_test_worker(
                temporal_env,
                workflows=[NonDeterministicWorkflow],
                activities=[],
            ):
                pass

    @pytest.mark.asyncio
    async def test_nondeterministic_workflow_decorator_enforced(self, temporal_env: WorkflowEnvironment) -> None:
        """enforce_determinism=True on decorator always enforces, regardless of config."""
        with pytest.raises(RuntimeError):
            async with create_test_worker(
                temporal_env,
                workflows=[NonDeterministicWorkflowEnforced],
                activities=[],
            ):
                pass

    @pytest.mark.asyncio
    async def test_nondeterministic_workflow_decorator_not_enforced(self, temporal_env: WorkflowEnvironment) -> None:
        """enforce_determinism=False on decorator never enforces, regardless of config."""
        async with create_test_worker(
            temporal_env,
            workflows=[NonDeterministicWorkflowNotEnforced],
            activities=[],
        ):
            pass
