# Copyright 2024 Klira AI
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Context propagation utilities for Klira tracing.

Provides helpers for Klira-specific context values (user_id, conversation_id,
message_id) and cross-thread context propagation.
"""

from __future__ import annotations

from typing import Any, Callable, TypeVar

from opentelemetry import context

T = TypeVar("T")

_KLIRA_USER_ID_KEY = context.create_key("klira.user_id")
_KLIRA_CONVERSATION_ID_KEY = context.create_key("klira.conversation_id")
_KLIRA_MESSAGE_ID_KEY = context.create_key("klira.message_id")


def set_klira_context(
    user_id: str | None = None,
    conversation_id: str | None = None,
    message_id: str | None = None,
) -> context.Context:
    """Set Klira context values in the current OTel context.

    Returns:
        The new context with Klira values set.
    """
    ctx = context.get_current()
    if user_id is not None:
        ctx = context.set_value(_KLIRA_USER_ID_KEY, user_id, ctx)
    if conversation_id is not None:
        ctx = context.set_value(_KLIRA_CONVERSATION_ID_KEY, conversation_id, ctx)
    if message_id is not None:
        ctx = context.set_value(_KLIRA_MESSAGE_ID_KEY, message_id, ctx)
    return ctx


def get_klira_user_id() -> str | None:
    """Get the current Klira user_id from context."""
    return context.get_value(_KLIRA_USER_ID_KEY)  # type: ignore[return-value]


def get_klira_conversation_id() -> str | None:
    """Get the current Klira conversation_id from context."""
    return context.get_value(_KLIRA_CONVERSATION_ID_KEY)  # type: ignore[return-value]


def get_klira_message_id() -> str | None:
    """Get the current Klira message_id from context."""
    return context.get_value(_KLIRA_MESSAGE_ID_KEY)  # type: ignore[return-value]


def is_in_klira_trace() -> bool:
    """Check if we're currently inside a Klira trace context."""
    return get_klira_user_id() is not None


def propagate_context(fn: Callable[..., T]) -> Callable[..., T]:
    """Wrap a callable so it carries the current OTel context into a new thread.

    Use this when passing Klira-decorated functions to ``run_in_executor()``
    or ``ThreadPoolExecutor``. Without it, the OTel context (and therefore
    span parent-child relationships) is lost across thread boundaries.

    Example::

        from klira.sdk.tracing import propagate_context

        result = await loop.run_in_executor(
            None,
            propagate_context(lambda: graph.invoke(state)),
        )
    """
    captured_ctx = context.get_current()

    def _wrapper(*args: Any, **kwargs: Any) -> T:
        token = context.attach(captured_ctx)
        try:
            return fn(*args, **kwargs)
        finally:
            context.detach(token)

    return _wrapper  # type: ignore[return-value]
