"""Custom log processor to inject trace context into OpenTelemetry log records."""

import contextvars
import logging
from typing import Optional, Dict, Any
from opentelemetry import trace
try:
    # New OTel SDK API (1.39.1+)
    from opentelemetry.sdk._logs import LogRecordProcessor
    from opentelemetry._logs._internal import LogRecord as OTelLogRecord
    from opentelemetry.sdk._logs import ReadableLogRecord
except ImportError:
    try:
        # Older OTel SDK API
        from opentelemetry.sdk._logs import LogRecordProcessor
        from opentelemetry.sdk._logs import LogRecord as OTelLogRecord
        ReadableLogRecord = OTelLogRecord
    except ImportError:
        # Fallback for very old versions or if SDK is missing
        class LogRecordProcessor: pass
        OTelLogRecord = Any
        ReadableLogRecord = Any

# Import context from context module
try:
    from varicon_observability.context import tenant_context, username_context, request_id_context, client_type_context, request_body_context,app_name_context
except ImportError:
    try:
        from varicon_observability.celery_helpers import tenant_context
        import contextvars
        username_context = contextvars.ContextVar("username", default="")
        request_id_context = contextvars.ContextVar("request_id", default="")
        client_type_context = contextvars.ContextVar("client_type", default="")
        request_body_context = contextvars.ContextVar("request_body", default="")
        app_name_context = contextvars.ContextVar('app_name',default="")
    except ImportError:
        import contextvars
        tenant_context = contextvars.ContextVar("tenant_id", default="unknown-tenant")
        username_context = contextvars.ContextVar("username", default="")
        request_id_context = contextvars.ContextVar("request_id", default="")
        client_type_context = contextvars.ContextVar("client_type", default="")
        request_body_context = contextvars.ContextVar("request_body", default="")
        app_name_context = contextvars.ContextVar('app_name',default="")

class TraceContextLogProcessor(LogRecordProcessor):
    """
    Injects trace context (trace_id, span_id) and tenant_id into OpenTelemetry log records.
    
    This processor runs AFTER the LoggingHandler creates the log record, so it can:
    1. Override trace_id/span_id if they weren't captured correctly
    2. Add tenant_id from context variable
    3. Add these as searchable attributes for SigNoz
    
    This is especially important for Celery tasks where the span might be
    set up in task_prerun but the LoggingHandler might miss it.
    """
    
    def on_emit(self, log_data) -> None:
        """
        New OTel SDK API method. Receives LogData wrapper.
        """
        log_record = getattr(log_data, 'log_record', log_data)
        self._process_record(log_record)

    def emit(self, log_data) -> None:
        """
        Older OTel SDK API method. Also receives LogData wrapper.
        """
        log_record = getattr(log_data, 'log_record', log_data)
        self._process_record(log_record)

    def _process_record(self, log_record) -> None:
        """
        Internal method to process and inject trace context and tenant_id.

        IMPORTANT: Never replace log_record.attributes with a plain dict.
        The OTel SDK uses BoundedAttributes which has a .dropped property
        required by the exporter. Just write to it in-place via [] access.
        """
        try:
            attrs = getattr(log_record, 'attributes', None)
            if attrs is None:
                # No attributes object at all — nothing we can write to safely
                return

            # 1. Extract and add tenant_id from context variable
            tenant_id = self._get_tenant_id(log_record)
            if tenant_id:
                attrs["tenant_id"] = tenant_id

            # 2. Extract and add username and request_id from context variables
            username = self._get_context_var(username_context)
            if username:
                attrs["username"] = username

            request_id = self._get_context_var(request_id_context)
            if request_id:
                attrs["request_id"] = request_id

            # 2b. Extract and add client_type from context variable
            client_type = self._get_context_var(client_type_context)
            if client_type:
                attrs["client_type"] = client_type

            # 2c. Extract and add request_body from context variable
            request_body = self._get_context_var(request_body_context)
            if request_body:
                attrs["request_body"] = request_body

            app_name = self._get_context_var(app_name_context)
            if app_name:
                attrs["app_name"] = app_name

            # 3. Get current span and inject/override trace context
            span = trace.get_current_span()

            # Fallback: check if we are in a Celery task but lost context
            if not (span and span.get_span_context().is_valid):
                span = self._get_celery_span_fallback()

            if span is not None:
                span_context = span.get_span_context()
                if span_context is not None and span_context.is_valid:
                    if span_context.trace_id != 0 and span_context.span_id != 0:
                        trace_id_hex = format(span_context.trace_id, "032x")
                        span_id_hex = format(span_context.span_id, "016x")

                        attrs["trace_id"] = trace_id_hex
                        attrs["span_id"] = span_id_hex
                        attrs["traceId"] = trace_id_hex
                        attrs["spanId"] = span_id_hex
                        attrs["trace_flags"] = int(span_context.trace_flags)

                        # Override the log record's direct fields for OTLP export.
                        # The processor receives the inner API LogRecord (via
                        # ReadWriteLogRecord.log_record) where trace_id/span_id
                        # are plain writable instance attributes.
                        # Use direct assignment (matches working v1 approach).
                        log_record.trace_id = span_context.trace_id
                        log_record.span_id = span_context.span_id
                        log_record.trace_flags = span_context.trace_flags

            # 4. Fallback: copy trace context from record fields into attributes
            if "trace_id" not in attrs:
                if hasattr(log_record, 'trace_id') and log_record.trace_id and log_record.trace_id != 0:
                    trace_id_hex = format(log_record.trace_id, "032x")
                    attrs["trace_id"] = trace_id_hex
                    attrs["traceId"] = trace_id_hex
                if hasattr(log_record, 'span_id') and log_record.span_id and log_record.span_id != 0:
                    span_id_hex = format(log_record.span_id, "016x")
                    attrs["span_id"] = span_id_hex
                    attrs["spanId"] = span_id_hex

        except Exception:
            # Silently ignore errors to avoid recursion and logging failures
            pass
    
    def _get_celery_span_fallback(self) -> Optional[trace.Span]:
        """Retrieve span from Celery task context if not in current OTel context."""
        try:
            from varicon_observability.celery_helpers import get_celery_span_fallback
            return get_celery_span_fallback()
        except ImportError:
            return None
            
    def _get_context_var(self, context_var) -> Optional[str]:
        """Safely retrieve a value from a ContextVar."""
        try:
            val = context_var.get()
            if val and str(val).strip():
                return str(val)
        except LookupError:
            pass
        except Exception:
            pass
        return None
    
    def _get_tenant_id(self, log_record: OTelLogRecord) -> Optional[str]:
        """
        Extract tenant_id from context variable or log message.
        
        Args:
            log_record: The log record to extract tenant_id from
            
        Returns:
            The tenant_id string or None if not found
        """
        # First try the context variable (most reliable for Celery tasks)
        tenant_id = self._get_context_var(tenant_context)
        if tenant_id:
            return tenant_id
        
        # Fallback: try to extract from existing attributes
        try:
            if log_record.attributes and "tenant_id" in log_record.attributes:
                return str(log_record.attributes["tenant_id"])
        except Exception:
            pass
        
        # Final fallback: try to extract from log message
        try:
            if hasattr(log_record, 'body') and log_record.body:
                body_str = str(log_record.body)
                if 'Tenant:' in body_str:
                    tenant_part = body_str.split('Tenant:')[1].split(']')[0].strip()
                    if tenant_part and tenant_part != 'unknown-tenant':
                        return tenant_part
        except Exception:
            pass
        
        return None
    
    def shutdown(self) -> None:
        """Shutdown the processor. No resources to clean up."""
        pass
    
    def force_flush(self, timeout_millis: int = 30000) -> bool:
        """Force flush any buffered log records. No-op since we don't buffer."""
        return True
