import inspect
from contextlib import asynccontextmanager
from typing import AsyncGenerator, Callable, List, Optional, Set, Type

from temporalio.testing import WorkflowEnvironment
from temporalio.worker import Worker
from temporalio.worker.workflow_sandbox import SandboxedWorkflowRunner

from mistralai.workflows.core.activity import activity
from mistralai.workflows.core.dependencies.dependency_injector import DependencyInjector, DependsCls
from mistralai.workflows.core.events.event_activities import (
    _emit_task_completed,
    _emit_task_failed,
    _emit_task_in_progress,
    _emit_task_started,
    _emit_waiting_for_input_completed,
    _emit_waiting_for_input_failed,
    _emit_waiting_for_input_started,
    _emit_workflow_canceled,
    _emit_workflow_completed,
    _emit_workflow_failed,
    _emit_workflow_started,
)
from mistralai.workflows.core.events.event_interceptor import EventInterceptor
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
from mistralai.workflows.core.sandbox import SANDBOX_RESTRICTIONS


@asynccontextmanager
async def _scoped_dependencies(injector: DependencyInjector, needed_deps: Set[Callable]) -> AsyncGenerator[None, None]:
    """Establish a DI context scoped to only the given dependency callables.

    Temporarily replaces the global registrations with the subset matching
    ``needed_deps``, resolves them via ``with_dependencies()``, and restores
    the original registrations on exit.
    """
    snapshot = injector._dependencies_ctx.copy()
    try:
        injector._dependencies_ctx = {dep: ctx for dep, ctx in snapshot.items() if dep in needed_deps}
        async with injector.with_dependencies():
            yield
    finally:
        injector._dependencies_ctx = snapshot
        injector._dependencies_values.set(None)


def _collect_activity_dependencies(activities: List[Callable]) -> Set[Callable]:
    """Walk the __wrapped__ chain of each activity to find Depends() declarations."""
    deps: Set[Callable] = set()
    for act in activities:
        func: Optional[Callable] = act
        while func is not None:
            for param in inspect.signature(func).parameters.values():
                if isinstance(param.default, DependsCls):
                    deps.add(param.default.dependency)
            func = getattr(func, "__wrapped__", None)
    return deps


@asynccontextmanager
async def create_test_worker(
    env: WorkflowEnvironment,
    workflows: List[Type],
    activities: List[Callable] | None = None,
    task_queue: str = "test-task-queue",
    interceptors: List | None = None,
) -> AsyncGenerator[Worker, None]:
    """Create a test worker with basic workflow event activities.

    This helper sets up a Temporal worker for testing with the necessary
    event-emitting activities pre-registered.
    """

    @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)
    all_activities.extend(
        [
            _emit_waiting_for_input_started,
            _emit_waiting_for_input_completed,
            _emit_waiting_for_input_failed,
            _emit_workflow_started,
            _emit_workflow_completed,
            _emit_workflow_canceled,
            _emit_workflow_failed,
        ]
    )

    worker = Worker(
        env.client,
        task_queue=task_queue,
        workflows=workflows,
        activities=all_activities,
        interceptors=interceptors or [],
        workflow_runner=SandboxedWorkflowRunner(
            restrictions=SANDBOX_RESTRICTIONS,
        ),
    )

    # Establish a DI context like the production worker does (worker.py:361).
    # Scope it to only the dependencies required by the activities under test,
    # so unrelated deps registered by other modules (e.g. webhook plugin)
    # in the same process don't get resolved and cause connection failures.
    needed_deps = _collect_activity_dependencies(all_activities)
    injector = DependencyInjector.get_singleton_instance()
    async with _scoped_dependencies(injector, needed_deps), worker:
        yield worker


@asynccontextmanager
async def create_test_worker_with_events(
    env: WorkflowEnvironment,
    workflows: List[Type],
    activities: List[Callable] | None = None,
    task_queue: str = "test-task-queue",
) -> AsyncGenerator[Worker, None]:
    """Create a test worker with full event support including task events.

    This helper sets up a Temporal worker for testing with both workflow
    and task event-emitting activities, plus the EventInterceptor for
    capturing custom task events.
    """

    @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)
    all_activities.extend(
        [
            _emit_waiting_for_input_started,
            _emit_waiting_for_input_completed,
            _emit_waiting_for_input_failed,
            _emit_workflow_started,
            _emit_workflow_completed,
            _emit_workflow_canceled,
            _emit_workflow_failed,
            _emit_task_started,
            _emit_task_in_progress,
            _emit_task_completed,
            _emit_task_failed,
        ]
    )

    worker = Worker(
        env.client,
        task_queue=task_queue,
        workflows=workflows,
        activities=all_activities,
        interceptors=[EventInterceptor()],
        workflow_runner=SandboxedWorkflowRunner(
            restrictions=SANDBOX_RESTRICTIONS,
        ),
    )

    needed_deps = _collect_activity_dependencies(all_activities)
    injector = DependencyInjector.get_singleton_instance()
    async with _scoped_dependencies(injector, needed_deps), worker:
        yield worker
