import asyncio
from contextlib import asynccontextmanager
from typing import Any, AsyncGenerator, Callable, List, Type

import structlog
from temporalio.testing import WorkflowEnvironment
from temporalio.worker import Worker

from mistralai_workflows.core.activity import activity
from mistralai_workflows.core.execution.sticky_session.get_sticky_worker_session import (
    GET_STICKY_WORKER_SESSION_ACTIVITY_NAME,
)
from mistralai_workflows.core.execution.sticky_session.sticky_worker_session import StickyWorkerSession

logger = structlog.get_logger(__name__)


@asynccontextmanager
async def create_test_worker(
    env: WorkflowEnvironment,
    workflows: List[Type],
    activities: List[Callable] | None = None,
    task_queue: str = "test-task-queue",
) -> AsyncGenerator[Worker, None]:
    # Create the get_sticky_worker_session activity that returns a fixed task queue
    # For tests, we just use the same task queue - no need to replicate production's multi-worker setup
    @activity(name=GET_STICKY_WORKER_SESSION_ACTIVITY_NAME, _skip_registering=True)
    async def get_sticky_worker_session() -> StickyWorkerSession:
        return StickyWorkerSession(task_queue=task_queue)

    all_activities = list(activities or [])
    all_activities.append(get_sticky_worker_session)

    worker = Worker(
        env.client,
        task_queue=task_queue,
        workflows=workflows,
        activities=all_activities,
    )

    async with worker:
        yield worker


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()
