"""Main client for the Lumenova Beacon SDK."""

from __future__ import annotations

import copy
import logging
import os

from contextvars import ContextVar
from datetime import datetime
from typing import TYPE_CHECKING, Any, Callable

if TYPE_CHECKING:
    from opentelemetry.sdk.trace import TracerProvider as OtelTracerProvider

from lumenova_beacon.core.config import BeaconConfig
from lumenova_beacon.core.transport import FileTransport, HTTPTransport
from lumenova_beacon.exceptions import ConfigurationError
from lumenova_beacon.masking.engine import apply_masking
from lumenova_beacon.tracing.span import Span
from lumenova_beacon.tracing.trace import TraceContext, get_current_span, get_current_trace_id
from lumenova_beacon.types import SpanKind, SpanType


logger = logging.getLogger(__name__)


class _UserIdentitySpanProcessor:
    """SpanProcessor that injects user identity attributes into every span.

    This ensures instrumentor-created spans (LangChain, CrewAI, etc.) carry
    user identity even though they don't go through BeaconClient.create_span().
    """

    def __init__(self, identity_attrs: dict[str, str]):
        self._identity_attrs = identity_attrs

    def on_start(self, span, parent_context=None):
        for key, value in self._identity_attrs.items():
            span.set_attribute(key, value)

    def on_end(self, span):
        pass

    def shutdown(self):
        pass

    def force_flush(self, timeout_millis=None):
        return True


# Thread-safe client instance using contextvars (for span context propagation)
_current_client: ContextVar['BeaconClient | None'] = ContextVar('current_client', default=None)

# Module-level singleton client (persists across async contexts)
_singleton_client: 'BeaconClient | None' = None

# Global flag to track if Beacon OTLP exporter has been added (prevents duplicates)
_beacon_global_exporter_added: bool = False

# Track export failures to warn users (only warn once to avoid log spam)
_export_failure_warned: bool = False


class BeaconClient:
    """Main client for creating and sending spans."""

    def __init__(
        self,
        endpoint: str | None = None,
        api_key: str | None = None,
        file_directory: str | None = None,
        session_id: str | None = None,
        user_email: str | None = None,
        user_name: str | None = None,
        user_id: str | None = None,
        service_name: str | None = None,
        environment: str | None = None,
        auto_instrument_opentelemetry: bool = True,
        isolated: bool = False,
        auto_instrument_litellm: bool = False,
        masking_function: Callable[[Any], Any] | None = None,
        verify: bool | None = None,
        eager_export: bool | None = None,
        record_stacktrace: bool = False,
        database_callback: Callable[[Span], None] | None = None,
        **kwargs,
    ):
        """Initialize the Beacon client.

        Configuration priority (constructor params override env vars):
        - Constructor parameters take precedence
        - Environment variables as fallback

        Spans are sent via OpenTelemetry's OTLPSpanExporter to the /v1/traces endpoint.
        Authentication is handled via 'headers' parameter or OTEL_EXPORTER_OTLP_HEADERS env var.

        Args:
            endpoint: Base server URL (e.g., "https://api.example.com"). Falls back to BEACON_ENDPOINT.
            api_key: API key for authentication (falls back to BEACON_API_KEY)
            file_directory: Directory to save span files (File mode)
            session_id: Default session ID for all spans (falls back to BEACON_SESSION_ID)
            user_email: Default user email for attribution (falls back to BEACON_USER_EMAIL)
            user_name: Default user display name for attribution (falls back to BEACON_USER_NAME)
            user_id: Default user ID for attribution (falls back to BEACON_USER_ID)
            service_name: Logical service name for OpenTelemetry resource (falls back to BEACON_SERVICE_NAME)
            environment: Deployment environment name (e.g., "production", "staging"). Falls back to BEACON_ENVIRONMENT.
            auto_instrument_opentelemetry: Automatically set up OpenTelemetry integration (default: True)
            auto_instrument_litellm: Automatically set up LiteLLM integration (default: False)
            masking_function: Custom function to mask sensitive data before transmission (optional)
            verify: Enable SSL certificate verification (default: True, falls back to BEACON_VERIFY)
            record_stacktrace: Capture full stack traces when recording exceptions (default: False)
            database_callback: Callback function for direct database export (internal use only)
            **kwargs: Additional configuration options (timeout, enabled, headers, etc.)
        """

        # Store database callback for internal use
        self._database_callback = database_callback

        # Constructor params override env vars (simple priority)
        # Handle verify separately to allow env var fallback
        if verify is None:
            verify_env = os.getenv('BEACON_VERIFY', 'true').lower()
            verify = verify_env not in ('false', '0', 'no')

        endpoint = endpoint if endpoint is not None else os.getenv('BEACON_ENDPOINT', '')
        api_key = api_key if api_key is not None else os.getenv('BEACON_API_KEY')

        # Third-tier fallback: AWS Secrets Manager (only if no api_key yet)
        if api_key is None:
            from lumenova_beacon.core.secrets import resolve_api_key_from_aws
            api_key = resolve_api_key_from_aws()

        service_name = service_name if service_name is not None else os.getenv('BEACON_SERVICE_NAME', os.getenv('OTEL_SERVICE_NAME'))
        session_id = session_id if session_id is not None else os.getenv('BEACON_SESSION_ID')
        user_email = user_email if user_email is not None else os.getenv('BEACON_USER_EMAIL')
        user_name = user_name if user_name is not None else os.getenv('BEACON_USER_NAME')
        user_id = user_id if user_id is not None else os.getenv('BEACON_USER_ID')
        environment = environment if environment is not None else os.getenv('BEACON_ENVIRONMENT')

        # Resolve eager_export from env var
        if eager_export is None:
            eager_export_env = os.getenv('BEACON_EAGER_EXPORT', 'true').lower()
            eager_export = eager_export_env not in ('false', '0', 'no')

        # Resolve isolated from env var (constructor param overrides)
        if not isolated:
            isolated_env = os.getenv('BEACON_ISOLATED', 'false').lower()
            isolated = isolated_env not in ('false', '0', 'no')

        # Build headers - start with custom headers if provided, then add auth
        headers = dict(kwargs.pop('headers', {}))
        if api_key:
            headers["Authorization"] = f"Bearer {api_key}"

        config = BeaconConfig(
            endpoint=endpoint,
            api_key=api_key,
            file_directory=file_directory or '',
            session_id=session_id,
            user_email=user_email,
            user_name=user_name,
            user_id=user_id,
            service_name=service_name,
            environment=environment,
            verify=verify,
            masking_function=masking_function,
            headers=headers,
            eager_export=eager_export,
            record_stacktrace=record_stacktrace,
            **kwargs,
        )

        # Skip validation if database_callback is provided (internal mode)
        if not database_callback:
            config.validate()
        self.config = config

        # Configure logger to always show warnings, and show debug when debug=True
        # This ensures export failures are visible to users even without debug mode
        if not logger.handlers:
            handler = logging.StreamHandler()
            formatter = logging.Formatter(
                '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
            )
            handler.setFormatter(formatter)
            logger.addHandler(handler)

        if config.debug:
            logger.setLevel(logging.DEBUG)
            for h in logger.handlers:
                h.setLevel(logging.DEBUG)
        else:
            # Ensure WARNING level so critical errors are visible
            if logger.level == logging.NOTSET or logger.level > logging.WARNING:
                logger.setLevel(logging.WARNING)
            for h in logger.handlers:
                if h.level == logging.NOTSET or h.level > logging.WARNING:
                    h.setLevel(logging.WARNING)

        if config.file_directory:
            # File mode: file_directory is provided
            self.transport = FileTransport(
                directory=config.file_directory,
                filename_pattern=config.file_filename_pattern,
                pretty_print=config.file_pretty_print,
            )
        else:
            # HTTP mode: endpoint is provided
            self.transport = HTTPTransport(
                endpoint=config.endpoint,
                api_key=config.api_key,
                timeout=config.timeout,
                headers=config.headers,
                verify=config.verify,
            )

        global _singleton_client
        _current_client.set(self)
        _singleton_client = self

        # Hidden timestamp override for seeding historical data
        self._timestamp_override: datetime | None = None

        # Store OTEL resource for span export (built from service_name + session_id)
        self._otel_resource = None

        # Store reference to the TracerProvider Beacon is using (may differ from global)
        self._otel_provider: 'OtelTracerProvider | None' = None

        # Automatically set up OpenTelemetry integration if requested
        if auto_instrument_opentelemetry:
            self._setup_opentelemetry(isolated=isolated)
            self._setup_additional_otlp_endpoints()

        # Automatically set up LiteLLM integration if requested
        if auto_instrument_litellm:
            self._setup_litellm()


    def _check_endpoint_connectivity(self) -> None:
        """Check if the endpoint is reachable and log WARNING if not.

        This is called during setup to provide early warning about connectivity
        issues that would prevent spans from being exported.
        """
        import socket
        from urllib.parse import urlparse

        parsed = urlparse(self.config.endpoint)
        host = parsed.hostname
        if not host:
            logger.warning(
                f"Invalid endpoint URL: {self.config.endpoint}. "
                "Spans will not be exported."
            )
            return

        port = parsed.port or (443 if parsed.scheme == 'https' else 80)

        try:
            socket.create_connection((host, port), timeout=5)
        except socket.gaierror as e:
            logger.warning(
                f"DNS resolution failed for {host}: {e}. "
                "Spans will not be exported until this is resolved."
            )
        except socket.timeout:
            logger.warning(
                f"Connection to {host}:{port} timed out. "
                "Spans may not be exported if the endpoint is unreachable."
            )
        except (ConnectionRefusedError, OSError) as e:
            logger.warning(
                f"Cannot connect to {self.config.endpoint}: {e}. "
                "Spans will not be exported until this is resolved."
            )

    def _setup_opentelemetry(self, isolated: bool = False) -> None:
        """Automatically configure OpenTelemetry to export spans.

        Sets up either OTLPSpanExporter (HTTP mode) or FileSpanExporter (file-only mode).
        This is called automatically when auto_instrument_opentelemetry=True (default).

        Args:
            isolated: If True, create a new TracerProvider without touching the global.
                      If False (default), attach to existing global or create new global.

        Note:
            If OpenTelemetry SDK is not installed, this method logs a warning
            and returns without raising an error (optional dependency).
        """
        global _beacon_global_exporter_added

        try:
            from opentelemetry import trace
            from opentelemetry.sdk.trace import TracerProvider, SpanLimits
            from opentelemetry.sdk.trace.export import BatchSpanProcessor
            from opentelemetry.sdk.resources import Resource
        except ImportError:
            logger.warning(
                'OpenTelemetry SDK not installed. Skipping automatic instrumentation. '
                'Install with: pip install lumenova-beacon[opentelemetry]'
            )
            return

        # Enable OpenTelemetry logging to surface export errors
        # Always enable WARNING level; DEBUG level only when debug=True
        import logging as stdlib_logging
        otel_logger = stdlib_logging.getLogger('opentelemetry')

        # Ensure OTel logger has a handler so messages are visible
        if not otel_logger.handlers:
            handler = stdlib_logging.StreamHandler()
            formatter = stdlib_logging.Formatter(
                '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
            )
            handler.setFormatter(formatter)
            otel_logger.addHandler(handler)

        if self.config.debug:
            otel_logger.setLevel(stdlib_logging.DEBUG)
            for h in otel_logger.handlers:
                h.setLevel(stdlib_logging.DEBUG)
        else:
            # Ensure at least WARNING level so export errors are visible
            if otel_logger.level == stdlib_logging.NOTSET or otel_logger.level > stdlib_logging.WARNING:
                otel_logger.setLevel(stdlib_logging.WARNING)
            for h in otel_logger.handlers:
                if h.level == stdlib_logging.NOTSET or h.level > stdlib_logging.WARNING:
                    h.setLevel(stdlib_logging.WARNING)

        # Build the resource for this client (used in export_span regardless of provider)
        try:
            from importlib.metadata import version as pkg_version
            sdk_version = pkg_version("lumenova-beacon")
        except Exception:
            from lumenova_beacon import __version__
            sdk_version = __version__
        resource_attrs = {
            "service.name": self.config.service_name or "unknown_service",
            "telemetry.sdk.name": "lumenova-beacon",
            "telemetry.sdk.version": sdk_version,
            "telemetry.sdk.language": "python",
        }
        if self.config.session_id:
            resource_attrs["beacon.session.id"] = self.config.session_id
        if self.config.user_email:
            resource_attrs["user.email"] = self.config.user_email
        if self.config.user_name:
            resource_attrs["user.name"] = self.config.user_name
        if self.config.user_id:
            resource_attrs["user.id"] = self.config.user_id
        if self.config.environment:
            resource_attrs["deployment.environment.name"] = self.config.environment
        self._otel_resource = Resource.create(resource_attrs)

        # Build user identity attributes to inject into all spans
        user_identity_attrs = {}
        if self.config.user_email:
            user_identity_attrs["user.email"] = self.config.user_email
        if self.config.user_name:
            user_identity_attrs["user.name"] = self.config.user_name
        if self.config.user_id:
            user_identity_attrs["user.id"] = self.config.user_id

        # Non-isolated path: reuse existing global provider if present
        if not isolated:
            if _beacon_global_exporter_added:
                logger.debug('Beacon already registered with global provider, skipping')
                return
            current_provider = trace.get_tracer_provider()
            if isinstance(current_provider, TracerProvider):
                self._otel_provider = current_provider
                # Add user identity processor so instrumentor spans carry user info
                if user_identity_attrs:
                    current_provider.add_span_processor(
                        _UserIdentitySpanProcessor(user_identity_attrs)
                    )
                _beacon_global_exporter_added = True
                logger.info('Using existing TracerProvider for Beacon span export')
                return

        # Create exporter — only needed when creating a new provider
        exporter = self._create_exporter()
        if exporter is None:
            return

        # Configure span limits to allow unlimited attribute value length
        # Default OTEL limit is 4096 chars which truncates large JSON outputs
        span_limits = SpanLimits(max_attribute_length=None)

        # Create a new TracerProvider (isolated, or no existing global)
        try:
            provider = TracerProvider(resource=self._otel_resource, span_limits=span_limits)
            provider.add_span_processor(BatchSpanProcessor(exporter))
            self._otel_provider = provider

            if not isolated:
                trace.set_tracer_provider(provider)
                _beacon_global_exporter_added = True
                # Add user identity processor to global provider
                if user_identity_attrs:
                    provider.add_span_processor(
                        _UserIdentitySpanProcessor(user_identity_attrs)
                    )
                logger.info('OpenTelemetry integration configured automatically')
            else:
                logger.info('OpenTelemetry integration configured with isolated TracerProvider')
                # When isolated, also inject user identity into the global provider
                # so that instrumentor spans (which use the global provider) carry user info
                if user_identity_attrs:
                    global_provider = trace.get_tracer_provider()
                    if isinstance(global_provider, TracerProvider):
                        global_provider.add_span_processor(
                            _UserIdentitySpanProcessor(user_identity_attrs)
                        )

            if self.config.endpoint:
                self._check_endpoint_connectivity()
        except Exception as e:
            logger.warning(f'Failed to set up OpenTelemetry provider: {e}')

    def _create_exporter(self):
        """Create the appropriate span exporter based on configuration.

        Returns:
            FileSpanExporter for file-only mode, OTLPSpanExporter for HTTP mode,
            or None if creation fails.
        """
        if self.config.endpoint:
            return self._create_otlp_exporter()

        if self.config.file_directory:
            from lumenova_beacon.tracing.exporters.opentelemetry import FileSpanExporter
            logger.info(f"File-only mode: spans will be exported to {self.config.file_directory}")
            return FileSpanExporter(
                directory=self.config.file_directory,
                filename_pattern=self.config.file_filename_pattern,
                pretty_print=self.config.file_pretty_print,
            )

        return None

    def _setup_litellm(self) -> None:
        """Automatically configure LiteLLM to send spans to Beacon.

        This method creates a BeaconLiteLLMLogger and registers it with
        litellm.callbacks. This is called automatically when
        auto_instrument_litellm=True is set (default: False).

        Note:
            If litellm is not installed, this method logs a warning
            and returns without raising an error (optional dependency).

            If a BeaconLiteLLMLogger is already registered, this method
            skips registration to avoid duplicates.
        """
        try:
            import litellm
        except ImportError:
            logger.warning(
                'litellm not installed. Skipping automatic instrumentation. '
                'Install with: pip install litellm'
            )
            return

        # Check if BeaconLiteLLMLogger is already registered
        if litellm.callbacks:
            for callback in litellm.callbacks:
                if callback.__class__.__name__ == 'BeaconLiteLLMLogger':
                    logger.debug('LiteLLM already has BeaconLiteLLMLogger registered, skipping auto-setup')
                    return

        try:
            from lumenova_beacon.tracing.integrations.litellm import BeaconLiteLLMLogger

            # Create logger instance (it will use this client via get_client())
            beacon_logger = BeaconLiteLLMLogger()

            # Register with LiteLLM
            if litellm.callbacks is None:
                litellm.callbacks = [beacon_logger]
            else:
                litellm.callbacks.append(beacon_logger)

            logger.info('LiteLLM integration configured automatically')
        except Exception as e:
            logger.warning(f'Failed to set up LiteLLM integration: {e}')

    def _create_otlp_exporter(self):
        """Create and configure an OTLP exporter.

        Returns:
            Configured OTLPSpanExporter instance

        Raises:
            ImportError: If OTLP exporter package is not installed
        """
        try:
            from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
        except ImportError:
            logger.error(
                "OTLP exporter not installed. "
                "Install with: pip install opentelemetry-exporter-otlp-proto-http"
            )
            raise

        # Build OTLP endpoint URL
        endpoint = self.config.endpoint.rstrip('/')

        # Add /v1/traces if not present (required for OTLPSpanExporter endpoint parameter)
        if not endpoint.endswith('/v1/traces'):
            endpoint = f"{endpoint}/v1/traces"

        # Build headers for OTLP
        headers = dict(self.config.headers)

        # Add session_id as HTTP headers for backend routing
        if self.config.session_id:
            headers["x-beacon-session-id"] = self.config.session_id

        # Handle SSL verification
        # Note: OTLPSpanExporter has a bug where certificate_file=False doesn't work
        # because of `certificate_file or environ.get(..., True)` logic.
        # We work around this by using a custom session that forces verify=False.
        session = None
        if not self.config.verify:
            import requests
            import urllib3
            urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

            # Create a custom session that forces verify=False on all requests
            class InsecureSession(requests.Session):
                def request(self, method, url, **kwargs):
                    kwargs['verify'] = False
                    return super().request(method, url, **kwargs)

            session = InsecureSession()

        # Temporarily clear OTEL env vars so OTLPSpanExporter doesn't read them
        # internally and overwrite our explicit endpoint/headers. These env vars are
        # intended for the global (infra) provider — not for Beacon's isolated exporter.
        _saved_otlp_headers = os.environ.pop('OTEL_EXPORTER_OTLP_HEADERS', None)
        _saved_otlp_endpoint = os.environ.pop('OTEL_EXPORTER_OTLP_ENDPOINT', None)
        try:
            exporter = OTLPSpanExporter(
                endpoint=endpoint,
                headers=headers,
                timeout=int(self.config.timeout),
                session=session,
            )
        finally:
            if _saved_otlp_headers is not None:
                os.environ['OTEL_EXPORTER_OTLP_HEADERS'] = _saved_otlp_headers
            if _saved_otlp_endpoint is not None:
                os.environ['OTEL_EXPORTER_OTLP_ENDPOINT'] = _saved_otlp_endpoint
        return exporter

    def add_span_exporter(self, exporter) -> None:
        """Add an additional span exporter to the TracerProvider.

        Allows sending spans to multiple OTLP-compatible destinations
        in parallel. The exporter is wrapped in a BatchSpanProcessor
        and added to the existing TracerProvider.

        Args:
            exporter: An OpenTelemetry SpanExporter instance

        Raises:
            RuntimeError: If no TracerProvider is configured
        """
        from opentelemetry import trace
        from opentelemetry.sdk.trace import TracerProvider
        from opentelemetry.sdk.trace.export import BatchSpanProcessor

        provider = self._otel_provider or trace.get_tracer_provider()
        if not isinstance(provider, TracerProvider):
            raise RuntimeError(
                "No TracerProvider configured. "
                "Ensure BeaconClient is initialized with auto_instrument_opentelemetry=True."
            )

        processor = BatchSpanProcessor(exporter)
        provider.add_span_processor(processor)
        logger.info("Added additional span exporter to TracerProvider")

    def _setup_additional_otlp_endpoints(self) -> None:
        """Configure additional OTLP exporter endpoints from BEACON_EXTRA_OTLP_ENDPOINTS.

        Parses a comma-separated list of OTLP endpoint URLs and creates a gRPC
        OTLPSpanExporter for each. Headers are shared except for the Authorization
        header, which is excluded from extra endpoints.

        Example:
            BEACON_EXTRA_OTLP_ENDPOINTS=http://collector:4317,http://other-collector:4317
        """
        endpoints_raw = os.getenv('BEACON_EXTRA_OTLP_ENDPOINTS')
        if not endpoints_raw:
            return

        endpoints = [ep.strip() for ep in endpoints_raw.split(',') if ep.strip()]
        if not endpoints:
            return

        try:
            from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
        except ImportError:
            logger.warning(
                'BEACON_EXTRA_OTLP_ENDPOINTS configured but gRPC OTLP exporter not installed. '
                'Install with: pip install opentelemetry-exporter-otlp-proto-grpc'
            )
            return

        for endpoint in endpoints:
            try:
                url = endpoint.rstrip('/')

                # Share headers but exclude authorization for extra endpoints
                headers = {
                    k: v for k, v in self.config.headers.items()
                    if k.lower() != "authorization"
                }
                if self.config.session_id:
                    headers["x-beacon-session-id"] = self.config.session_id

                exporter = OTLPSpanExporter(
                    endpoint=url,
                    headers=headers,
                    timeout=int(self.config.timeout),
                )
                self.add_span_exporter(exporter)
                logger.info(f"Added additional OTLP endpoint: {endpoint}")
            except Exception as e:
                logger.warning(f"Failed to add OTLP endpoint {endpoint}: {e}")

    @property
    def otel_provider(self) -> 'OtelTracerProvider':
        """Return the TracerProvider Beacon is exporting spans to.

        Returns the stored isolated provider when initialized with isolated=True,
        otherwise falls back to the global TracerProvider.
        """
        from opentelemetry.sdk.trace import TracerProvider as _SDKTracerProvider
        from opentelemetry import trace as _trace
        if self._otel_provider is not None:
            return self._otel_provider
        provider = _trace.get_tracer_provider()
        if not isinstance(provider, _SDKTracerProvider):
            raise RuntimeError(
                "No SDK TracerProvider configured. "
                "Ensure BeaconClient is initialized with auto_instrument_opentelemetry=True, "
                "or that a global SDK TracerProvider is set before accessing otel_provider."
            )
        return provider

    def should_sample(self) -> bool:
        """Determine if the current operation should be sampled.

        Returns:
            True if should sample, False otherwise
        """
        if not self.config.enabled:
            return False
        return True

    def create_span(
        self,
        name: str,
        *,
        kind: SpanKind = SpanKind.INTERNAL,
        span_type: SpanType = SpanType.SPAN,
        auto_parent: bool = True,
        session_id: str | None = None,
        user_email: str | None = None,
        user_name: str | None = None,
        user_id: str | None = None,
    ) -> Span:
        """Create a new span.

        Args:
            name: Name of the span
            kind: Type of span (default: INTERNAL)
            span_type: Type of span (default: SPAN)
            auto_parent: Automatically set parent from context (default: True)
            session_id: Session ID (optional, inherits from parent or client default)
            user_email: User email for attribution (optional, inherits from parent or client default)
            user_name: User display name for attribution (optional, inherits from parent or client default)
            user_id: User ID for attribution (optional, inherits from parent or client default)

        Returns:
            New Span instance
        """
        trace_id = get_current_trace_id()
        parent_id = None

        if auto_parent:
            # Check SDK's ContextVar for current span
            current_span = get_current_span()
            if current_span:
                trace_id = current_span.trace_id
                parent_id = current_span.span_id

                # Inherit session_id from parent if not explicitly provided
                if session_id is None:
                    session_id = current_span.session_id

                # Inherit user identity from parent if not explicitly provided
                if user_email is None:
                    user_email = current_span.user_email
                if user_name is None:
                    user_name = current_span.user_name
                if user_id is None:
                    user_id = current_span.user_id

        # Fall back to client defaults if no parent and no explicit value provided
        if session_id is None:
            session_id = self.config.session_id
        if user_email is None:
            user_email = self.config.user_email
        if user_name is None:
            user_name = self.config.user_name
        if user_id is None:
            user_id = self.config.user_id

        span = Span(
            name=name,
            trace_id=trace_id,
            parent_id=parent_id,
            kind=kind,
            span_type=span_type,
            session_id=session_id,
            user_email=user_email,
            user_name=user_name,
            user_id=user_id,
            record_stacktrace=self.config.record_stacktrace,
        )

        if self.config.environment:
            span.set_attribute("deployment.environment.name", self.config.environment)

        return span

    def trace(
        self,
        name: str,
        *,
        kind: SpanKind = SpanKind.INTERNAL,
        span_type: SpanType = SpanType.SPAN,
        session_id: str | None = None,
        user_email: str | None = None,
        user_name: str | None = None,
        user_id: str | None = None,
    ) -> TraceContext:
        """Create a span and return a context manager for it.

        Usage:
            with client.trace("operation_name"):
                # Your code here
                pass

        Args:
            name: Name of the span
            kind: Type of span (default: INTERNAL)
            span_type: Type of span (default: SPAN)
            session_id: Session ID (optional, stored in attributes)
            user_email: User email for attribution (optional)
            user_name: User display name for attribution (optional)
            user_id: User ID for attribution (optional)

        Returns:
            TraceContext that can be used as a context manager
        """
        span = self.create_span(
            name=name,
            kind=kind,
            span_type=span_type,
            session_id=session_id,
            user_email=user_email,
            user_name=user_name,
            user_id=user_id,
        )
        span.start()
        return TraceContext(span, client=self)

    def _apply_masking_to_span(self, span: Span) -> Span:
        """Apply the configured masking function to a span's attributes.

        Creates a shallow copy of the span to preserve the original, then
        applies masking to a copy of the span's attributes.

        Args:
            span: The span to mask

        Returns:
            A new span with masked attributes
        """
        masked_span = copy.copy(span)
        masked_span.attributes = apply_masking(
            dict(span.attributes), self.config.masking_function
        )
        return masked_span

    def export_span(self, span: Span) -> bool:
        """Export a span via database callback or OpenTelemetry exporter.

        This method either calls the database callback (if provided) or
        converts the Beacon span to an OpenTelemetry ReadableSpan and
        sends it through the configured TracerProvider's span processors.

        Args:
            span: The Beacon span to export

        Returns:
            True if export succeeded, False otherwise
        """
        # Apply masking if configured
        if self.config.masking_function is not None:
            span = self._apply_masking_to_span(span)

        # If database callback is provided, use it for direct export
        if self._database_callback:
            try:
                self._database_callback(span)
                logger.debug(
                    f"Exported span via database callback: name={span.name}, "
                    f"trace_id={span.trace_id}, span_id={span.span_id}"
                )
                return True
            except Exception as e:
                logger.error(f"Database callback failed for span {span.span_id}: {e}", exc_info=True)
                return False

        # Otherwise, use OpenTelemetry export
        try:
            from opentelemetry import trace
            from opentelemetry.sdk.trace import TracerProvider

            provider = self._otel_provider or trace.get_tracer_provider()
            if not isinstance(provider, TracerProvider):
                logger.warning("No OpenTelemetry TracerProvider configured, cannot export span")
                return False

            # Convert Beacon span to OTel ReadableSpan using client's resource
            otel_span = span.to_otel_span(resource=self._otel_resource)

            # Send through all span processors.
            # _active_span_processor is a private OTEL SDK attribute — there is no
            # public API for directly submitting a ReadableSpan to the batch processor
            # without going through a tracer. Monitor OTEL SDK releases for a cleaner
            # alternative.
            if hasattr(provider, '_active_span_processor'):
                provider._active_span_processor.on_end(otel_span)
                logger.debug(
                    f"Queued span for export: name={span.name}, "
                    f"trace_id={span.trace_id}, span_id={span.span_id}"
                )
                return True

            logger.warning("TracerProvider has no _active_span_processor")
            return False
        except ImportError:
            logger.warning("OpenTelemetry SDK not installed, cannot export span")
            return False
        except Exception as e:
            logger.error(f"Failed to export span: {e}")
            return False

    def export_eager_span(self, span: Span) -> bool:
        """Export an eager (incomplete) span with end_time=0.

        Used by integrations (e.g. LangChain) to send span metadata
        before the operation completes. The backend's upsert logic
        will merge this with the complete span later.

        Args:
            span: The Beacon span to eagerly export

        Returns:
            True if export succeeded, False otherwise
        """
        if not self.config.eager_export:
            return False

        # Apply masking if configured
        if self.config.masking_function is not None:
            span = self._apply_masking_to_span(span)

        try:
            from opentelemetry import trace
            from opentelemetry.sdk.trace import TracerProvider

            provider = self._otel_provider or trace.get_tracer_provider()
            if not isinstance(provider, TracerProvider):
                return False

            otel_span = span.to_otel_span_eager(resource=self._otel_resource)

            # See comment in export_span() about _active_span_processor.
            if hasattr(provider, '_active_span_processor'):
                provider._active_span_processor.on_end(otel_span)
                logger.debug(
                    f"Eager-exported span: name={span.name}, "
                    f"trace_id={span.trace_id}, span_id={span.span_id}"
                )
                return True

            return False
        except ImportError:
            return False
        except Exception as e:
            logger.error(f"Failed to eager-export span: {e}")
            return False

    def export_spans(self, spans: list[Span]) -> int:
        """Export multiple spans via the configured OpenTelemetry exporter.

        Args:
            spans: List of Beacon spans to export

        Returns:
            Number of spans successfully exported
        """
        count = 0
        for span in spans:
            if self.export_span(span):
                count += 1
        return count

    def flush(self, timeout_millis: int = 30000) -> bool:
        """Flush all pending spans to the backend.

        This forces the OpenTelemetry BatchSpanProcessor to export any
        buffered spans immediately.

        Args:
            timeout_millis: Maximum time to wait for flush in milliseconds

        Returns:
            True if flush succeeded, False otherwise
        """
        logger.debug(f"Flushing spans (timeout={timeout_millis}ms)...")
        try:
            from opentelemetry import trace
            from opentelemetry.sdk.trace import TracerProvider

            global _export_failure_warned

            provider = self._otel_provider or trace.get_tracer_provider()
            if isinstance(provider, TracerProvider):
                result = provider.force_flush(timeout_millis)
                if result:
                    logger.debug("Flush completed successfully")
                else:
                    # Only warn once to avoid log spam on repeated failures
                    if not _export_failure_warned:
                        logger.warning(
                            f"OTLP export failed - spans are not being sent to {self.config.endpoint}. "
                            "This usually means the endpoint is unreachable (DNS failure, network issue, or wrong URL). "
                            "Check your BEACON_ENDPOINT configuration and network connectivity."
                        )
                        _export_failure_warned = True
                    else:
                        logger.debug("Flush returned False (export failure already reported)")
                return result
            logger.debug("No TracerProvider to flush")
            return True  # No provider to flush
        except ImportError:
            logger.debug("OpenTelemetry not installed, nothing to flush")
            return True  # OpenTelemetry not installed, nothing to flush
        except Exception as e:
            logger.error(f"Failed to flush spans: {e}", exc_info=True)
            return False

    def __enter__(self) -> 'BeaconClient':
        """Enter sync context manager.

        Returns:
            Self for use in with statement
        """
        return self

    def __exit__(self, _exc_type, _exc_val, _exc_tb) -> None:
        """Exit sync context manager, flushing pending spans."""
        self.flush()

    async def __aenter__(self) -> 'BeaconClient':
        """Enter async context manager.

        Returns:
            Self for use in async with statement
        """
        return self

    async def __aexit__(self, _exc_type, _exc_val, _exc_tb) -> None:
        """Exit async context manager, flushing pending spans."""
        self.flush()

    def get_base_url(self) -> str:
        """Get the base URL for API requests.

        Returns the base endpoint URL (e.g., "https://api.example.com").

        Returns:
            Base URL string

        Raises:
            ConfigurationError: If transport is not HTTPTransport
        """
        if not isinstance(self.transport, HTTPTransport):
            raise ConfigurationError(
                'Dataset operations require HTTPTransport. '
                'File transport is not supported for datasets.'
            )

        return self.transport.endpoint


def get_client() -> 'BeaconClient':
    """Get the current Beacon client or init a new one.

    Priority:
    1. Current async context (ContextVar) - for span context propagation
    2. Module-level singleton - persists across async contexts
    3. Create new client - only if neither exists

    Returns:
        BeaconClient instance
    """
    # First check async context
    client = _current_client.get()
    if client is not None:
        return client

    # Fall back to module-level singleton (persists across async contexts)
    if _singleton_client is not None:
        return _singleton_client

    # Create new client (will set both _current_client and _singleton_client)
    return BeaconClient()
