"""
OpenTelemetry-native tracing for Galtea SDK.

This module provides a simple, OTel-native tracing API similar to Patronus AI:
- @traced() decorator for automatic function tracing
- start_trace() context manager for manual trace creation

All traces are automatically exported to Galtea API via OpenTelemetry processors.
No manual context management required.
"""

import inspect
import json
import logging
from contextlib import contextmanager
from functools import wraps
from typing import Any, Callable, Dict, Optional, TypeVar

from opentelemetry import context as otel_context
from opentelemetry import trace as otel_trace
from opentelemetry.context import create_key
from opentelemetry.trace import SpanKind, Status, StatusCode

from galtea.domain.models.trace import TraceType
from galtea.infrastructure.telemetry.attributes import (
    GALTEA_ATTR_INFERENCE_RESULT_ID,
    GALTEA_ATTR_SESSION_ID,
    GALTEA_ATTR_TRACE_DESCRIPTION,
    GALTEA_ATTR_TRACE_ERROR,
    GALTEA_ATTR_TRACE_INPUT,
    GALTEA_ATTR_TRACE_METADATA,
    GALTEA_ATTR_TRACE_OUTPUT,
    GALTEA_ATTR_TRACE_TYPE,
)

F = TypeVar("F", bound=Callable[..., Any])

_INFERENCE_ID_KEY = create_key("galtea_inference_id")
_SESSION_ID_KEY = create_key("galtea_session_id")

_logger = logging.getLogger(__name__)


def set_context(
    inference_result_id: str,
) -> object:
    """Set Galtea context values for trace correlation.

    Call this before traced functions to associate spans with an inference result.

    Args:
            inference_result_id: The inference result ID to associate traces with.

    Returns:
            A context token that can be used with clear_context().

    Example:
            token = set_context(inference_result_id="inf_123")
            try:
                    my_traced_function()
            finally:
                    clear_context(token)
    """
    if not inference_result_id:
        raise ValueError("Inference result ID cannot be empty")

    ctx = otel_context.get_current()
    ctx = otel_context.set_value(_INFERENCE_ID_KEY, inference_result_id, ctx)
    return otel_context.attach(ctx)


def clear_context(token: object, flush: bool = True) -> None:
    """Clear a previously attached context.

    Args:
            token: The token returned by `set_context()`.
            flush: If True (default), export any pending traces associated with the current
                inference result to the telemetry exporter before detaching the context.
    """
    # Get the inference result ID before detaching
    inference_result_id = get_inference_result_id() if flush else None

    # Detach the context
    otel_context.detach(token)

    # Flush traces for this inference result
    if flush and inference_result_id:
        from galtea.infrastructure.telemetry.provider import flush_inference

        flush_inference(inference_result_id)


def get_inference_result_id() -> Optional[str]:
    """Get the current inference result ID from context."""
    return otel_context.get_value(_INFERENCE_ID_KEY)


def get_session_id() -> Optional[str]:
    """Get the current session ID from context."""
    return otel_context.get_value(_SESSION_ID_KEY)


class GalteaSpan:
    """Wrapper around OpenTelemetry Span with Galtea-specific helper methods."""

    def __init__(self, span: otel_trace.Span):
        self._span = span

    def update(
        self,
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        type: Optional[str] = None,
    ) -> None:
        """Update span attributes.

        Args:
                input: The input data to record. Will be JSON-serialized.
                output: The output data to record. Will be JSON-serialized.
                metadata: The metadata to record. Will be JSON-serialized.
                type: TraceType value (e.g., TOOL, GENERATION, AGENT, CHAIN, RETRIEVER).
        """
        if not self._span.is_recording():
            _logger.warning("update called on a non-recording span")
            return

        attrs = _build_span_attributes(type=type, input_data=input, output_data=output, metadata=metadata)
        for key, value in attrs.items():
            self._span.set_attribute(key, value)

    def set_attribute(self, key: str, value: Any) -> None:
        """Set a custom attribute on the span."""
        self._span.set_attribute(key, _serialize_to_json(value))

    def record_exception(self, exception: BaseException) -> None:
        """Record an exception on the span with Galtea error attributes."""
        _record_span_error(self._span, exception)

    @property
    def span(self) -> otel_trace.Span:
        """Access the underlying OpenTelemetry Span."""
        return self._span


@contextmanager
def start_trace(
    name: str,
    type: Optional[str] = None,
    description: Optional[str] = None,
    input: Optional[Any] = None,
    metadata: Optional[Any] = None,
    attributes: Optional[Dict[str, Any]] = None,
):
    """Context manager for creating a trace.

    Use this for fine-grained tracing of specific code blocks.

    Args:
            name: Name of the trace.
            type: TraceType value (e.g., TOOL, GENERATION, AGENT, CHAIN, RETRIEVER).
            description: Description of the trace (e.g., what this operation does).
            input: Input data for the trace.
            metadata: Metadata for the trace.
            attributes: Custom OpenTelemetry attributes to add to the span.

    Yields:
            A GalteaSpan object with update() and other helper methods.

    Example:
            with start_trace("database_query", type="TOOL", description="Query user data") as s:
                    result = db.query("SELECT * FROM users")
                    s.update(output=result)

            # Or with dynamic attributes:
            with start_trace("llm_call", input="my prompt") as s:
                    response = llm.generate(prompt)
                    s.update(output={"response": response.text})
    """
    tracer = otel_trace.get_tracer("galtea")

    span_attributes = _build_span_attributes(
        type=type,
        description=description,
        input_data=input,
        metadata=metadata,
    )

    # Add raw OTel attributes (not Galtea semantic fields)
    if attributes:
        span_attributes.update(attributes)

    with tracer.start_as_current_span(
        name=name,
        kind=SpanKind.INTERNAL,
        attributes=span_attributes,
    ) as span:
        galtea_span = GalteaSpan(span)
        try:
            yield galtea_span
        except Exception as e:
            _record_span_error(span, e)
            raise


def trace(
    name: Optional[str] = None,
    type: Optional[str] = TraceType.SPAN,
    log_args: bool = True,
    log_results: bool = True,
    include_docstring: bool = False,
    attributes: Optional[Dict[str, Any]] = None,
    node_type: Optional[str] = None,
) -> Callable[[F], F]:
    """Decorator to automatically trace a function.

    All traced functions automatically create OpenTelemetry spans that are
    exported to Galtea. Exceptions are always
    recorded in traces for debugging purposes.

    Args:
            name: Custom span name. Defaults to function name.
            type: TraceType value (e.g., TOOL, GENERATION, AGENT, CHAIN, RETRIEVER).
                More info in: https://docs.galtea.ai/concepts/product/version/session/trace#trace-types
            log_args: Whether to log function arguments as input.
            log_results: Whether to log return value as output.
            include_docstring: Whether to include the function's docstring as the trace description.
            attributes: Custom attributes to add to the span.
            node_type: Node type for the trace (e.g., CUSTOM, LLM, TOOL, CHAIN). This parameter is deprecated and will be removed, use `type` instead.

    Returns:
            Decorated function with automatic tracing.

    Example:
            @trace
            def my_function():
                    pass

            @trace(type="TOOL", include_docstring=True)
            def fetch_user(user_id: str) -> dict:
                    # Docstring: "Fetches user data from the database."
                    return db.get_user(user_id)

            @trace(name="gpt4_completion", type="GENERATION", attributes={"model": "gpt-4"})
            def generate_response(prompt: str) -> str:
                    return openai.chat.completions.create(...)
    """
    if node_type is not None:
        import warnings

        warnings.warn(
            "The 'node_type' parameter is deprecated and will be removed in a future version. "
            "Use the 'type' parameter with TraceType enum instead.",
            DeprecationWarning,
            stacklevel=2,
        )

        # Convert node_type to type. The `node_type` parameter takes precedence when type is SPAN for backward compatibility.
        if type == TraceType.SPAN:
            if node_type == "LLM":
                type = TraceType.GENERATION
            elif node_type == "CUSTOM":
                type = TraceType.SPAN
            else:
                type = node_type

    # Handle @trace without parentheses (function passed directly as name)
    if callable(name):
        func = name
        # Call trace() with default parameters to get the decorator, then apply it
        return trace(
            name=None,
            type=type,
            log_args=log_args,
            log_results=log_results,
            include_docstring=include_docstring,
            attributes=attributes,
        )(func)

    if type:
        valid_types = {t.value for t in TraceType}
        if type.upper() not in valid_types:
            raise ValueError(
                f"Invalid type '{type}'. Must be one of: {', '.join(valid_types)}. "
                f"Use TraceType enum from galtea.domain.models.trace for type safety."
            )

    def decorator(func: F) -> F:
        span_name = name or func.__name__
        sig = inspect.signature(func)
        tracer = otel_trace.get_tracer("galtea")
        func_description = inspect.getdoc(func) if include_docstring else None

        def _build_input_data(args: tuple, kwargs: Dict[str, Any]) -> Optional[Dict[str, Any]]:
            if not log_args:
                return None
            try:
                bound = sig.bind(*args, **kwargs)
                bound.apply_defaults()
                input_data = dict(bound.arguments)
                input_data.pop("self", None)
                input_data.pop("cls", None)

                # Convert non-serializable objects to serializable format
                # Keep None values for inputs (they represent explicit parameters)
                serializable_data = {}
                for key, value in input_data.items():
                    serializable_data[key] = _to_serializable(value)

                return serializable_data if serializable_data else None
            except TypeError:
                # Fallback: convert everything to strings
                safe_args = [str(arg) for arg in args]
                safe_kwargs = {k: str(v) for k, v in kwargs.items()}
                return {"args": safe_args, "kwargs": safe_kwargs} if args or kwargs else None

        if inspect.iscoroutinefunction(func):

            @wraps(func)
            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
                input_data = _build_input_data(args, kwargs)
                span_attributes = _build_span_attributes(
                    type=type,
                    description=func_description,
                    input_data=input_data,
                    metadata=attributes,
                )

                with tracer.start_as_current_span(
                    name=span_name,
                    kind=SpanKind.INTERNAL,
                    attributes=span_attributes,
                ) as span:
                    try:
                        result = await func(*args, **kwargs)

                        if log_results and result is not None:
                            _set_output_on_span(span, result)

                        return result
                    except Exception as e:
                        _record_span_error(span, e)
                        raise

            return async_wrapper  # type: ignore[return-value]

        else:

            @wraps(func)
            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
                input_data = _build_input_data(args, kwargs)
                span_attributes = _build_span_attributes(
                    type=type,
                    description=func_description,
                    input_data=input_data,
                    metadata=attributes,
                )

                with tracer.start_as_current_span(
                    name=span_name,
                    kind=SpanKind.INTERNAL,
                    attributes=span_attributes,
                ) as span:
                    try:
                        result = func(*args, **kwargs)

                        if log_results and result is not None:
                            _set_output_on_span(span, result)

                        return result
                    except Exception as e:
                        _record_span_error(span, e)
                        raise

            return sync_wrapper  # type: ignore[return-value]

    return decorator


def _record_span_error(span: otel_trace.Span, exception: BaseException) -> None:
    """Record an exception on a span with Galtea error attributes.

    Sets span status to ERROR, adds Galtea error attribute to the span, and records the exception.

    Args:
        span: The OpenTelemetry span to record the error on.
        exception: The exception to record.
    """
    span.set_status(Status(StatusCode.ERROR, str(exception)))
    span.set_attribute(GALTEA_ATTR_TRACE_ERROR, str(exception))
    span.record_exception(exception)


def _serialize_to_json(value: Any) -> str:
    """Serialize a value to JSON string for span attributes.

    Handles Pydantic models by converting them to dicts first.
    """
    try:
        serializable_value = _to_serializable(value)
        return json.dumps(serializable_value)
    except (TypeError, ValueError):
        return json.dumps(str(value))


def _to_serializable(value: Any) -> Any:
    """Convert a value to a JSON-serializable format.

    Recursively handles Pydantic models, dicts, and lists.
    """
    # Handle Pydantic v2 models
    if hasattr(value, "model_dump"):
        return _to_serializable(value.model_dump())
    # Handle Pydantic v1 models (check for __fields__ to identify Pydantic models)
    if hasattr(value, "__fields__") and hasattr(value, "dict"):
        return _to_serializable(value.dict())
    # Handle dicts recursively
    if isinstance(value, dict):
        return {k: _to_serializable(v) for k, v in value.items()}
    # Handle lists recursively
    if isinstance(value, list):
        return [_to_serializable(item) for item in value]
    return value


def _set_output_on_span(span: otel_trace.Span, output: Any) -> None:
    """Set output attributes on a span using unified serialization.

    This is the single entry point for setting output on spans,
    ensuring consistent serialization across all code paths.
    """
    output_json = _serialize_to_json(output)
    span.set_attribute(GALTEA_ATTR_TRACE_OUTPUT, output_json)


def _build_span_attributes(
    type: Optional[str] = None,
    description: Optional[str] = None,
    input_data: Optional[Any] = None,
    output_data: Optional[Any] = None,
    error: Optional[str] = None,
    metadata: Optional[Any] = None,
) -> Dict[str, Any]:
    """Build span attributes for Galtea."""
    attrs: Dict[str, Any] = {}

    # Add context values
    inference_id = get_inference_result_id()
    session_id = get_session_id()

    if inference_id:
        attrs[GALTEA_ATTR_INFERENCE_RESULT_ID] = inference_id

    if session_id:
        attrs[GALTEA_ATTR_SESSION_ID] = session_id

    # Add trace type
    if type:
        valid_types = {t.value for t in TraceType}
        if type.upper() not in valid_types:
            raise ValueError(
                f"Invalid type '{type}'. Must be one of: {', '.join(valid_types)}. "
                f"Use TraceType enum from galtea.domain.models.trace for type safety."
            )
        attrs[GALTEA_ATTR_TRACE_TYPE] = type

    # Add description (from function docstring)
    if description:
        attrs[GALTEA_ATTR_TRACE_DESCRIPTION] = description

    # Add input/output (explicit None check to allow falsy values like 0, "", {})
    if input_data is not None:
        input_json = _serialize_to_json(input_data)
        attrs[GALTEA_ATTR_TRACE_INPUT] = input_json

    if output_data is not None:
        output_json = _serialize_to_json(output_data)
        attrs[GALTEA_ATTR_TRACE_OUTPUT] = output_json

    if error is not None:
        attrs[GALTEA_ATTR_TRACE_ERROR] = error

    if metadata is not None:
        metadata_json = _serialize_to_json(metadata)
        attrs[GALTEA_ATTR_TRACE_METADATA] = metadata_json

    return attrs
