import asyncio
from typing import Any, Type

from temporalio.testing import WorkflowEnvironment

from mistralai.workflows.core.activity import get_all_temporal_activities
from mistralai.workflows.testing import (
    create_test_worker,
    create_test_worker_with_events,
    wait_for_pending_inputs,
)

__all__ = [
    "create_test_worker",
    "create_test_worker_with_events",
    "execute_workflow_in_test_env",
    "wait_for_pending_inputs",
]


async def execute_workflow_in_test_env(
    env: WorkflowEnvironment,
    workflow_class: Type,
    workflow_input: Any,
    workflow_id: str | None = None,
    task_queue: str = "test-task-queue",
) -> Any:
    from mistralai.workflows.core.definition.workflow_definition import get_workflow_definition

    workflow_def: Any = get_workflow_definition(workflow_class)
    if not workflow_def:
        raise ValueError(f"Workflow {workflow_class} is not properly decorated")

    handle = await env.client.start_workflow(
        workflow_def.name,
        workflow_input,
        id=workflow_id or f"test-workflow-{asyncio.current_task().get_name()}",  # type: ignore
        task_queue=task_queue,
    )

    return await handle.result()


def get_temporal_activities_by_names(names: list[str]) -> list:
    """Get temporal activities from the registry by name."""
    all_activities = get_all_temporal_activities()
    return [act for act in all_activities if act.__name__ in names]
