import asyncio
from typing import Awaitable, Callable, cast, Coroutine, Union
from .models import (
    CompletionFunctionOutputs,
    CompletionRequest,
    RiskEvaluationOutputs,
    RiskEvaluationRequest,
)
from .types import (
    CompletionFnTelemetryContext,
    RiskEvalTelemetryContext,
)

# Import tool context management
try:
    from snowglobe.tools.context import ToolCallContext

    TOOLS_AVAILABLE = True
except ImportError:
    ToolCallContext = None
    TOOLS_AVAILABLE = False


async def run_completion_fn(
    completion_fn: Union[
        Callable[[CompletionRequest], CompletionFunctionOutputs],
        Callable[[CompletionRequest], Awaitable[CompletionFunctionOutputs]],
    ],
    app_id: str,
    completion_request: CompletionRequest,
    telemetry_context: CompletionFnTelemetryContext,  # noqa
    enable_tool_mocking: bool = True,
) -> CompletionFunctionOutputs:
    context_tokens = None
    # Inject tool call context if tools module is available
    if TOOLS_AVAILABLE:
        # Extract Snowglobe metadata from request
        conversation_id = None
        test_id = None
        # Use the most recent message with snowglobe_data (current test)
        for msg in reversed(completion_request.messages or []):
            if msg.snowglobe_data:
                test_id = msg.snowglobe_data.test_id
                # Prefer the latest conversation_id if present
                conversation_id = msg.snowglobe_data.conversation_id or conversation_id
                break

        # Set context for tool execution
        if conversation_id and test_id and ToolCallContext is not None:
            context_tokens = ToolCallContext.set_context(
                conversation_id=conversation_id,
                test_id=test_id,
                application_id=app_id,
                mock_enabled=enable_tool_mocking,
            )

    try:
        # Execute completion function
        if asyncio.iscoroutinefunction(completion_fn):
            try:
                asyncio.get_running_loop()
                awaitable_completion_output = completion_fn(completion_request)
                completion_output = await awaitable_completion_output
            except RuntimeError:
                awaitable_completion_output: Coroutine[
                    None, None, CompletionFunctionOutputs
                ] = completion_fn(completion_request)
                completion_output = asyncio.run(awaitable_completion_output)
        else:
            sync_completion_fn = cast(
                Callable[[CompletionRequest], CompletionFunctionOutputs], completion_fn
            )
            completion_output = await asyncio.to_thread(
                sync_completion_fn,
                completion_request,
            )

        return completion_output
    finally:
        # Clean up context
        if (
            TOOLS_AVAILABLE
            and context_tokens is not None
            and ToolCallContext is not None
        ):
            ToolCallContext.clear_context(context_tokens)


async def run_risk_evaluation_fn(
    risk_evaluation_fn: Union[
        Callable[[RiskEvaluationRequest], RiskEvaluationOutputs],
        Callable[[RiskEvaluationRequest], Awaitable[RiskEvaluationOutputs]],
    ],
    risk_evaluation_request: RiskEvaluationRequest,
    telemetry_context: RiskEvalTelemetryContext,  # noqa
) -> RiskEvaluationOutputs:
    if asyncio.iscoroutinefunction(risk_evaluation_fn):
        try:
            asyncio.get_running_loop()
            awaitable_risk_eval = risk_evaluation_fn(risk_evaluation_request)
            risk_evaluation = await awaitable_risk_eval
        except RuntimeError:
            awaitable_risk_eval: Coroutine[
                None, None, RiskEvaluationOutputs
            ] = risk_evaluation_fn(risk_evaluation_request)
            risk_evaluation = asyncio.run(awaitable_risk_eval)
    else:
        # pyright doesn't understand that the above check means the function is synchronous here...
        sync_risk_evaluation_fn = cast(
            Callable[[RiskEvaluationRequest], RiskEvaluationOutputs], risk_evaluation_fn
        )
        risk_evaluation = await asyncio.to_thread(
            sync_risk_evaluation_fn,
            risk_evaluation_request,
        )
    return risk_evaluation
