# 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.

"""UserMessageTraceContext — root span for user interactions.

Creates a klira.user.message root span and sets Klira context values
(user_id, conversation_id, message_id) for child span propagation.

Usage:
    with UserMessageTraceContext(user_id="u1", conversation_id="c1"):
        # All klira spans created here are children of the root span
        result = my_workflow("Hello")
"""

from __future__ import annotations

import uuid
import logging
from types import TracebackType

from opentelemetry import context, trace

from klira.sdk.tracing.propagation import set_klira_context

logger = logging.getLogger(__name__)


class UserMessageTraceContext:
    """Context manager that creates a klira.user.message root span.

    On exit, the root span is ended and traces are flushed.
    """

    INPUT_TRUNCATION_LIMIT: int = 10000

    def __init__(
        self,
        user_id: str | None = None,
        conversation_id: str | None = None,
        message_id: str | None = None,
        message_text: str | None = None,
    ) -> None:
        self._user_id = user_id or "anonymous"
        self._conversation_id = conversation_id or str(uuid.uuid4())
        self._message_id = message_id or str(uuid.uuid4())
        self._message_text = message_text
        self._tracer = trace.get_tracer("klira")
        self._span: trace.Span | None = None
        self._token: context.Token | None = None  # type: ignore[type-arg]

    def __enter__(self) -> UserMessageTraceContext:
        # Set Klira context values
        ctx = set_klira_context(
            user_id=self._user_id,
            conversation_id=self._conversation_id,
            message_id=self._message_id,
        )
        self._token = context.attach(ctx)

        # Create root span
        attrs: dict[str, str] = {
            "klira.entity_type": "user_message",
            "klira.user_id": self._user_id,
            "klira.conversation_id": self._conversation_id,
            "klira.message_id": self._message_id,
        }
        try:
            from klira.sdk import Klira

            if Klira.is_initialized():
                config = Klira.get_config()
                if config.evals_run is not None:
                    attrs["klira.evals.evals_run"] = config.evals_run
        except Exception:
            logger.debug("Failed to inject klira.evals.evals_run", exc_info=True)

        self._span = self._tracer.start_span(
            "klira.user.message",
            attributes=attrs,
        )

        if self._message_text is not None:
            text = self._message_text
            if len(text) > self.INPUT_TRUNCATION_LIMIT:
                text = text[: self.INPUT_TRUNCATION_LIMIT]
            self._span.set_attribute("klira.input", text)

        # Attach span to context
        ctx_with_span = trace.set_span_in_context(self._span, context.get_current())
        self._span_token = context.attach(ctx_with_span)

        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> None:
        if self._span is not None:
            if exc_val is not None:
                self._span.set_status(
                    trace.StatusCode.ERROR,
                    description=str(exc_val),
                )
                self._span.record_exception(exc_val)
            self._span.end()

        # Detach contexts in reverse order
        if hasattr(self, "_span_token") and self._span_token is not None:
            context.detach(self._span_token)
        if self._token is not None:
            context.detach(self._token)

    # Async context manager support — delegates to sync __enter__/__exit__
    # since OTel context operations are not truly async.
    async def __aenter__(self) -> UserMessageTraceContext:
        return self.__enter__()

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> None:
        self.__exit__(exc_type, exc_val, exc_tb)

    @property
    def span(self) -> trace.Span | None:
        """The root klira.user.message span."""
        return self._span

    @property
    def user_id(self) -> str:
        return self._user_id

    @property
    def conversation_id(self) -> str:
        return self._conversation_id

    @property
    def message_id(self) -> str:
        return self._message_id
