# © 2023 SolarWinds Worldwide, LLC. All rights reserved.
#
# 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.

"""Module to initialize OpenTelemetry SDK components for SolarWinds backend"""

from __future__ import annotations

import logging
import math
import os
import uuid

from opentelemetry import trace
from opentelemetry._logs import set_logger_provider
from opentelemetry._logs._internal import NoOpLoggerProvider
from opentelemetry.environment_variables import (
    OTEL_METRICS_EXPORTER,
    OTEL_PROPAGATORS,
)
from opentelemetry.instrumentation.propagators import (
    set_global_response_propagator,
)
from opentelemetry.metrics import set_meter_provider
from opentelemetry.metrics._internal import NoOpMeterProvider
from opentelemetry.propagate import set_global_textmap
from opentelemetry.propagators.composite import CompositePropagator
from opentelemetry.sdk._configuration import (
    _get_exporter_names,
    _import_exporters,
    _init_logging,
    _OTelSDKConfigurator,
)
from opentelemetry.sdk.metrics import (
    Counter,
    Histogram,
    MeterProvider,
    ObservableCounter,
    ObservableGauge,
    ObservableUpDownCounter,
    UpDownCounter,
)
from opentelemetry.sdk.metrics.export import (
    AggregationTemporality,
    MetricExporter,
    MetricReader,
    PeriodicExportingMetricReader,
)
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace.export import (
    BatchSpanProcessor,
    SimpleSpanProcessor,
    SpanExporter,
)
from opentelemetry.sdk.trace.id_generator import IdGenerator
from opentelemetry.sdk.trace.sampling import Sampler
from opentelemetry.trace import NoOpTracerProvider, set_tracer_provider
from opentelemetry.util._importlib_metadata import entry_points

from solarwinds_apm import apm_logging
from solarwinds_apm.apm_config import SolarWindsApmConfig
from solarwinds_apm.apm_constants import INTL_SWO_DEFAULT_PROPAGATORS
from solarwinds_apm.response_propagator import (
    SolarWindsTraceResponsePropagator,
)
from solarwinds_apm.sampler import ParentBasedSwSampler
from solarwinds_apm.trace import (
    ResponseTimeProcessor,
    ServiceEntrySpanProcessor,
)
from solarwinds_apm.tracer_provider import SolarwindsTracerProvider
from solarwinds_apm.version import __version__

solarwinds_apm_logger = apm_logging.logger
logger = logging.getLogger(__name__)


class SolarWindsConfigurator(_OTelSDKConfigurator):
    """Configure OpenTelemetry SDK components for SolarWinds APM reporting."""

    def __init__(self) -> None:
        """
        Initialize the SolarWinds Configurator.

        Creates APM configuration instance for SDK initialization.
        """
        super().__init__()
        self.apm_config = SolarWindsApmConfig()

    def _create_apm_resource(self) -> Resource:
        """Create new OpenTelemetry Resource for telemetry providers.

        Returns:
        Resource: Resource with SolarWinds APM attributes and service information.
        """
        apm_resource = Resource.create(
            {
                "sw.apm.version": __version__,
                "sw.data.module": "apm",
                "service.name": self.apm_config.service_name,
            }
        )
        # Prioritize service.instance.id set by any Resource Detectors
        updated_apm_resource = apm_resource.merge(
            Resource(
                {"service.instance.id": str(uuid.uuid4())}
                if "service.instance.id" not in apm_resource.attributes
                else {}
            )
        )
        return updated_apm_resource

    def _configure(self, **kwargs: int) -> None:
        """Configure SolarWinds APM and OpenTelemetry components.

        Parameters:
        **kwargs (int): Additional configuration keyword arguments.
        """
        if not self.apm_config.agent_enabled:
            # ApmConfig will log a more informative Info message if disabled
            logger.debug(
                "SWO APM agent disabled. Not configuring OpenTelemetry."
            )
            set_tracer_provider(NoOpTracerProvider())
            set_meter_provider(NoOpMeterProvider())
            set_logger_provider(NoOpLoggerProvider())
            return

        # Duplicated from _OTelSDKConfigurator in order to support custom settings
        trace_exporter_names = kwargs.get("trace_exporter_names", [])
        metric_exporter_names = kwargs.get("metric_exporter_names", [])
        log_exporter_names = kwargs.get("log_exporter_names", [])

        span_exporters, metric_exporters, log_exporters = _import_exporters(
            trace_exporter_names + _get_exporter_names("traces"),
            metric_exporter_names + _get_exporter_names("metrics"),
            log_exporter_names + _get_exporter_names("logs"),
        )

        # Custom initialization of OTel components
        apm_sampler = ParentBasedSwSampler(
            self.apm_config,
        )
        apm_resource = self._create_apm_resource()

        self._custom_init_tracing(
            exporters=span_exporters,
            id_generator=kwargs.get("id_generator"),
            sampler=apm_sampler,
            resource=apm_resource,
        )
        self._custom_init_metrics(
            exporters_or_readers=metric_exporters,
            resource=apm_resource,
        )

        # Initialize LoggerProvider, log processors, and log exporters.
        # As of OTel 1.40.0, auto-instrumentation of Python logging (setup_logging_handler)
        # is handled by opentelemetry-instrumentation-logging via instrumentation, not by the
        # deprecated LoggingHandler in the SDK. So we pass _init_logging false.
        _init_logging(log_exporters, apm_resource, setup_logging_handler=False)

        # Set up additional custom SW components
        self._configure_service_entry_span_processor()
        self._configure_response_time_processor()
        self._configure_propagator()
        self._configure_response_propagator()

    def _custom_init_tracing(
        self,
        exporters: dict[str, type[SpanExporter]],
        id_generator: IdGenerator | None = None,
        sampler: Sampler | None = None,
        resource: Resource | None = None,
    ) -> None:
        """
        Initialize custom tracing with SolarWinds APM span processors.

        Based on OpenTelemetry SDK Configurator with custom processors.

        Parameters:
        exporters (dict[str, type[SpanExporter]]): Span exporter classes by name.
        id_generator (IdGenerator | None): Custom ID generator. Defaults to None.
        sampler (Sampler | None): Custom sampler. Defaults to None.
        resource (Resource | None): Resource attributes. Defaults to None.
        """
        provider = SolarwindsTracerProvider(
            id_generator=id_generator,
            sampler=sampler,
            resource=resource,
        )
        set_tracer_provider(provider)

        for _, exporter_class in exporters.items():
            exporter_args = {}
            if self.apm_config.is_lambda:
                provider.add_span_processor(
                    SimpleSpanProcessor(exporter_class(**exporter_args))
                )
            else:
                provider.add_span_processor(
                    BatchSpanProcessor(exporter_class(**exporter_args))
                )

    def _custom_init_metrics(
        self,
        exporters_or_readers: dict[
            str, type[MetricExporter] | type[MetricReader]
        ],
        resource: Resource | None = None,
    ) -> None:
        """
        Initialize custom metrics export with delta temporality.

        Based on OpenTelemetry SDK Configurator with delta aggregation.

        Parameters:
        exporters_or_readers (dict[str, type[MetricExporter] | type[MetricReader]]): Metric exporter or reader classes.
        resource (Resource | None): Resource attributes. Defaults to None.
        """
        metric_readers = []

        for _, exporter_or_reader_class in exporters_or_readers.items():
            exporter_args = {
                "preferred_temporality": {
                    Counter: AggregationTemporality.DELTA,
                    UpDownCounter: AggregationTemporality.DELTA,
                    Histogram: AggregationTemporality.DELTA,
                    ObservableCounter: AggregationTemporality.DELTA,
                    ObservableUpDownCounter: AggregationTemporality.DELTA,
                    ObservableGauge: AggregationTemporality.DELTA,
                }
            }

            if issubclass(exporter_or_reader_class, MetricReader):
                metric_readers.append(
                    exporter_or_reader_class(**exporter_args)
                )
            elif self.apm_config.is_lambda:
                # Inf interval to not invoke periodic collection
                metric_readers.append(
                    PeriodicExportingMetricReader(
                        exporter_or_reader_class(**exporter_args),
                        export_interval_millis=math.inf,
                    )
                )
            else:
                # Use default interval 60s else OTEL_METRIC_EXPORT_INTERVAL
                metric_readers.append(
                    PeriodicExportingMetricReader(
                        exporter_or_reader_class(**exporter_args),
                    )
                )

        provider = MeterProvider(
            resource=resource, metric_readers=metric_readers
        )
        set_meter_provider(provider)

    def _configure_service_entry_span_processor(
        self,
    ) -> None:
        """Configure ServiceEntrySpanProcessor for the tracer provider."""
        trace.get_tracer_provider().add_span_processor(
            ServiceEntrySpanProcessor()
        )

    def _configure_response_time_processor(
        self,
    ) -> None:
        """Configure ResponseTimeProcessor if metrics exporters are configured.

        Only configures if OTEL_METRICS_EXPORTER is set.
        """
        # SolarWindsDistro._configure does setdefault before this is called
        environ_exporter = os.environ.get(
            OTEL_METRICS_EXPORTER,
        )
        if not environ_exporter:
            logger.debug(
                "No OTEL_METRICS_EXPORTER set, skipping init of metrics processors"
            )
            return

        trace.get_tracer_provider().add_span_processor(
            ResponseTimeProcessor(
                self.apm_config,
            )
        )

    def _configure_propagator(self) -> None:
        """Configure CompositePropagator with SolarWinds and other propagators.

        Uses default or environment-configured propagators.
        """
        propagators = []

        # SolarWindsDistro._configure does setdefault so this shouldn't
        # be None, but safer and more explicit this way
        environ_propagators_names = os.environ.get(
            OTEL_PROPAGATORS,
            ",".join(INTL_SWO_DEFAULT_PROPAGATORS),
        ).split(",")

        for propagator_name in environ_propagators_names:
            try:
                propagators.append(
                    next(
                        iter(
                            entry_points(
                                group="opentelemetry_propagator",
                                name=propagator_name,
                            )
                        )
                    ).load()()
                )
            except Exception:
                logger.exception(
                    "Failed to load configured propagator %s",
                    propagator_name,
                )
                raise
        logger.debug(
            "Setting CompositePropagator with %s",
            environ_propagators_names,
        )
        set_global_textmap(CompositePropagator(propagators))

    def _configure_response_propagator(self) -> None:
        """Configure global HTTP response propagator."""
        # Set global HTTP response propagator
        set_global_response_propagator(SolarWindsTraceResponsePropagator())
