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

import logging
from typing import Optional, Any
from opentelemetry import trace

try:
    from opentelemetry.sdk._logs import LogRecordProcessor
    from opentelemetry._logs._internal import LogRecord as OTelLogRecord
    from opentelemetry.sdk._logs import ReadableLogRecord
except ImportError:
    try:
        from opentelemetry.sdk._logs import LogRecordProcessor
        from opentelemetry.sdk._logs import LogRecord as OTelLogRecord
        ReadableLogRecord = OTelLogRecord
    except ImportError:
        class LogRecordProcessor:
            pass
        OTelLogRecord = Any
        ReadableLogRecord = Any

from varicon_observability.context import (
    tenant_context,
    username_context,
    request_id_context,
    client_type_context,
    request_body_context,
    app_name_context,
    screen_name_context,
    trace_id_context,
)


class TraceContextLogProcessor(LogRecordProcessor):
    """
    Injects trace context (trace_id, span_id) and tenant/request metadata into
    OpenTelemetry log records so they are searchable as attributes in SigNoz.

    Runs after the LoggingHandler creates the record. Especially important for
    Celery tasks where the span is set up in task_prerun but the LoggingHandler
    may miss it.
    """

    def on_emit(self, log_data) -> None:
        """New OTel SDK API (1.39.1+)."""
        self._process_record(getattr(log_data, "log_record", log_data))

    def emit(self, log_data) -> None:
        """Older OTel SDK API fallback."""
        self._process_record(getattr(log_data, "log_record", log_data))

    def _process_record(self, log_record) -> None:
        """
        Inject trace context and request metadata into the log record attributes.

        Never replaces log_record.attributes with a plain dict — the OTel SDK
        uses BoundedAttributes which has a .dropped property required by the
        exporter. Write to it in-place via [] access.
        """
        try:
            attrs = getattr(log_record, "attributes", None)
            if attrs is None:
                return

            # Inject context variables
            self._set_if_present(attrs, "tenant_id", self._get_tenant_id(log_record))
            self._set_if_present(attrs, "username", self._get_ctx(username_context))
            self._set_if_present(attrs, "request_id", self._get_ctx(request_id_context))
            self._set_if_present(attrs, "client_type", self._get_ctx(client_type_context))
            self._set_if_present(attrs, "request_body", self._get_ctx(request_body_context))
            self._set_if_present(attrs, "app_name", self._get_ctx(app_name_context))
            self._set_if_present(attrs, "screen_name", self._get_ctx(screen_name_context))
            self._set_if_present(attrs, "trace_id", self._get_ctx(trace_id_context))

            # Inject active span's trace context
            span = trace.get_current_span()
            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 and span_context.trace_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)
                    log_record.trace_id = span_context.trace_id
                    log_record.span_id = span_context.span_id
                    log_record.trace_flags = span_context.trace_flags

            # Fallback: copy from record fields if attributes still missing
            if "trace_id" not in attrs:
                if hasattr(log_record, "trace_id") and log_record.trace_id:
                    attrs["trace_id"] = format(log_record.trace_id, "032x")
                    attrs["traceId"] = attrs["trace_id"]
                if hasattr(log_record, "span_id") and log_record.span_id:
                    attrs["span_id"] = format(log_record.span_id, "016x")
                    attrs["spanId"] = attrs["span_id"]

        except Exception:
            pass

    @staticmethod
    def _set_if_present(attrs, key, value):
        if value:
            attrs[key] = value

    @staticmethod
    def _get_ctx(context_var) -> Optional[str]:
        try:
            val = context_var.get()
            if val and str(val).strip():
                return str(val)
        except Exception:
            pass
        return None

    def _get_celery_span_fallback(self) -> Optional[trace.Span]:
        try:
            from varicon_observability.celery_helpers import get_celery_span_fallback
            return get_celery_span_fallback()
        except ImportError:
            return None

    def _get_tenant_id(self, log_record) -> Optional[str]:
        tenant_id = self._get_ctx(tenant_context)
        if tenant_id:
            return tenant_id
        try:
            if log_record.attributes and "tenant_id" in log_record.attributes:
                return str(log_record.attributes["tenant_id"])
        except Exception:
            pass
        try:
            if hasattr(log_record, "body") and log_record.body:
                body_str = str(log_record.body)
                if "Tenant:" in body_str:
                    part = body_str.split("Tenant:")[1].split("]")[0].strip()
                    if part and part != "unknown-tenant":
                        return part
        except Exception:
            pass
        return None

    def shutdown(self) -> None:
        pass

    def force_flush(self, timeout_millis: int = 30000) -> bool:
        return True
