"""
Context management for tool execution.

This module provides async-safe context propagation for Snowglobe metadata
(conversation_id, test_id, mock_enabled) used by tool decorators.
"""

from contextvars import ContextVar
from typing import Dict, Optional

_conversation_id: ContextVar[Optional[str]] = ContextVar(
    "snowglobe_tool_conversation_id",
    default=None,
)
_test_id: ContextVar[Optional[str]] = ContextVar(
    "snowglobe_tool_test_id",
    default=None,
)
_application_id: ContextVar[Optional[str]] = ContextVar(
    "snowglobe_tool_application_id",
    default=None,
)
_mock_enabled: ContextVar[bool] = ContextVar(
    "snowglobe_tool_mock_enabled",
    default=False,
)


class ToolCallContext:
    """
    Thread-local storage for tool call context.

    This allows the runner to inject Snowglobe-specific context (conversation_id, test_id, mock_enabled)
    that can be retrieved by tool decorators during execution.
    """

    @classmethod
    def set_context(
        cls,
        conversation_id: str,
        test_id: str,
        application_id: str,
        mock_enabled: bool = False,
    ) -> Dict[str, object]:
        """
        Set the context for the current thread.

        Args:
            conversation_id: The Snowglobe conversation ID
            test_id: The Snowglobe test ID
            application_id: The Snowglobe application ID
            mock_enabled: Whether tool mocking is enabled
        """
        return {
            "conversation_id": _conversation_id.set(conversation_id),
            "test_id": _test_id.set(test_id),
            "application_id": _application_id.set(application_id),
            "mock_enabled": _mock_enabled.set(mock_enabled),
        }

    @classmethod
    def get_context(cls) -> dict:
        """
        Get the context for the current thread.

        Returns:
            Dictionary containing:
                - conversation_id: str or None
                - test_id: str or None
                - application_id: str or None
                - mock_enabled: bool (default False)
        """
        return {
            "conversation_id": _conversation_id.get(),
            "test_id": _test_id.get(),
            "application_id": _application_id.get(),
            "mock_enabled": _mock_enabled.get(),
        }

    @classmethod
    def clear_context(cls, tokens: Optional[Dict[str, object]] = None) -> None:
        """Clear or reset the context for the current async context."""
        if tokens:
            conversation_token = tokens.get("conversation_id")
            if conversation_token is not None:
                _conversation_id.reset(conversation_token)
            test_token = tokens.get("test_id")
            if test_token is not None:
                _test_id.reset(test_token)
            mock_token = tokens.get("mock_enabled")
            if mock_token is not None:
                _mock_enabled.reset(mock_token)
            application_token = tokens.get("application_id")
            if application_token is not None:
                _application_id.reset(application_token)
            return

        _conversation_id.set(None)
        _test_id.set(None)
        _mock_enabled.set(False)
        _application_id.set(None)

    @classmethod
    def get_conversation_id(cls) -> Optional[str]:
        """Get the conversation ID from context."""
        return _conversation_id.get()

    @classmethod
    def get_test_id(cls) -> Optional[str]:
        """Get the test ID from context."""
        return _test_id.get()

    @classmethod
    def is_mock_enabled(cls) -> bool:
        """Check if tool mocking is enabled."""
        return _mock_enabled.get()

    @classmethod
    def get_application_id(cls) -> Optional[str]:
        """Get the application ID from context."""
        return _application_id.get()
