# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.
from __future__ import annotations

import atexit
import datetime
import logging
import random
import warnings
from collections.abc import Callable
from typing import TYPE_CHECKING

from opentelemetry import metrics
from opentelemetry.sdk.metrics import MeterProvider
from opentelemetry.sdk.metrics._internal.export import (
    ConsoleMetricExporter,
    PeriodicExportingMetricReader,
)
from opentelemetry.sdk.metrics.view import ExponentialBucketHistogramAggregation, View
from opentelemetry.sdk.resources import SERVICE_NAME, Resource

from ..common import get_otel_data_exporter
from ..otel_env_config import load_metrics_env_config
from .protocols import Timer
from .validators import (
    OTEL_NAME_MAX_LENGTH,
    ListValidator,
    PatternAllowListValidator,
    get_validator,
    stat_name_otel_handler,
)

if TYPE_CHECKING:
    from opentelemetry.metrics import Instrument
    from opentelemetry.util.types import Attributes

    from .protocols import DeltaType

log = logging.getLogger(__name__)

GaugeValues = int | float

DEFAULT_GAUGE_VALUE = 0.0

# "airflow.dag_processing.processes" is currently the only UDC used in Airflow.  If more are added,
# we should add a better system for this.
#
# Generally in OTel a Counter is monotonic (can only go up) and there is an UpDownCounter which,
# as you can guess, is non-monotonic; it can go up or down. The choice here is to either drop
# this one metric and implement the rest as monotonic Counters, implement all counters as
# UpDownCounters, or add a bit of logic to do it intelligently. The catch is that the Collector
# which transmits these metrics to the upstream dashboard tools (Prometheus, Grafana, etc.) assigns
# the type of Gauge to any UDC instead of Counter. Adding this logic feels like the best compromise
# where normal Counters still get typed correctly, and we don't lose an existing metric.
# See:
# https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/metrics/api.md#counter-creation
# https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/metrics/api.md#updowncounter
UP_DOWN_COUNTERS = {"airflow.dag_processing.processes"}

DEFAULT_METRIC_NAME_PREFIX = "airflow"
# Delimiter is placed between the universal metric prefix and the unique metric name.
DEFAULT_METRIC_NAME_DELIMITER = "."


def full_name(name: str, *, prefix: str = DEFAULT_METRIC_NAME_PREFIX) -> str:
    """Assembles the prefix, delimiter, and name and returns it as a string."""
    return f"{prefix}{DEFAULT_METRIC_NAME_DELIMITER}{name}"


def _is_up_down_counter(name):
    return name in UP_DOWN_COUNTERS


def _generate_key_name(name: str, attributes: Attributes = None):
    if attributes:
        key = name
        for item in attributes.items():
            key += f"_{item[0]}_{item[1]}"
    else:
        key = name

    return key


def name_is_otel_safe(prefix: str, name: str) -> bool:
    """
    Return True if the provided name and prefix would result in a name that meets the OpenTelemetry standard.

    Legal names are defined here:
    https://opentelemetry.io/docs/reference/specification/metrics/api/#instrument-name-syntax
    """
    return bool(stat_name_otel_handler(prefix, name, max_length=OTEL_NAME_MAX_LENGTH))


def _type_as_str(obj: Instrument) -> str:
    """
    Given an OpenTelemetry Instrument, returns the type of the instrument as a string.

    :param obj: An OTel Instrument or subclass
    :returns: The type() of the Instrument without all the nested class info
    """
    # type().__name__ will return something like: '_Counter',
    # this drops the leading underscore for cleaner logging.

    return type(obj).__name__[1:]


def _get_otel_safe_name(name: str) -> str:
    """
    Verify that the provided name does not exceed OpenTelemetry's maximum length for metric names.

    :param name: The original metric name
    :returns: The name, truncated to an OTel-acceptable length if required.
    """
    otel_safe_name = name[:OTEL_NAME_MAX_LENGTH]
    if name != otel_safe_name:
        warnings.warn(
            f"Metric name `{name}` exceeds OpenTelemetry's name length limit of "
            f"{OTEL_NAME_MAX_LENGTH} characters and will be truncated to `{otel_safe_name}`.",
            category=UserWarning,
            stacklevel=2,
        )
    return otel_safe_name


def _skip_due_to_rate(rate: float) -> bool:
    if rate < 0:
        raise ValueError("rate must be a positive value.")
    return rate < 1 and random.random() > rate


class _OtelTimer(Timer):
    """
    An implementation of stats.Timer() which records the result in the OTel Metrics Map.

    OpenTelemetry does not have a native timer; values are stored as a Histogram so that
    all observations (count, sum, bucket distribution) are preserved across multiple recordings.

    :param name: The name of the timer.
    :param tags: Tags to append to the timer.
    """

    def __init__(self, otel_logger: SafeOtelLogger, name: str | None, tags: Attributes):
        super().__init__()
        self.otel_logger = otel_logger
        self.name = name
        self.tags = tags

    def stop(self, send: bool = True) -> None:
        super().stop(send)
        if self.name and send and self.duration is not None:
            self.otel_logger.metrics_map.record_histogram_value(
                full_name(prefix=self.otel_logger.prefix, name=self.name), self.duration, self.tags
            )


class SafeOtelLogger:
    """Otel Logger."""

    def __init__(
        self,
        otel_provider,
        prefix: str = DEFAULT_METRIC_NAME_PREFIX,
        metrics_validator: ListValidator = PatternAllowListValidator(),
        stat_name_handler: Callable[[str], str] | None = None,
        statsd_influxdb_enabled: bool = False,
    ):
        self.otel: Callable = otel_provider
        self.prefix: str = prefix
        self.metrics_validator = metrics_validator
        self.meter = otel_provider.get_meter(__name__)
        self.metrics_map = MetricsMap(self.meter)
        self.stat_name_handler = stat_name_handler
        self.statsd_influxdb_enabled = statsd_influxdb_enabled

    def incr(
        self,
        stat: str,
        count: int = 1,
        rate: float = 1,
        tags: Attributes = None,
    ):
        """
        Increment stat by count.

        :param stat: The name of the stat to increment.
        :param count: A positive integer to add to the current value of stat.
        :param rate: value between 0 and 1 that represents the sample rate at
            which the metric is going to be emitted.
        :param tags: Tags to append to the stat.
        """
        if _skip_due_to_rate(rate):
            return
        if count < 0:
            raise ValueError("count must be a positive value.")

        if self.metrics_validator.test(stat) and name_is_otel_safe(self.prefix, stat):
            counter = self.metrics_map.get_counter(full_name(prefix=self.prefix, name=stat), attributes=tags)
            counter.add(count, attributes=tags)
            return counter

    def decr(
        self,
        stat: str,
        count: int = 1,
        rate: float = 1,
        tags: Attributes = None,
    ):
        """
        Decrement stat by count.

        :param stat: The name of the stat to decrement.
        :param count: A positive integer to subtract from current value of stat.
        :param rate: value between 0 and 1 that represents the sample rate at
            which the metric is going to be emitted.
        :param tags: Tags to append to the stat.
        """
        if _skip_due_to_rate(rate):
            return
        if count < 0:
            raise ValueError("count must be a positive value.")

        if self.metrics_validator.test(stat) and name_is_otel_safe(self.prefix, stat):
            counter = self.metrics_map.get_counter(full_name(prefix=self.prefix, name=stat))
            counter.add(-count, attributes=tags)
            return counter

    def gauge(
        self,
        stat: str,
        value: int | float,
        rate: float = 1,
        delta: bool = False,
        *,
        tags: Attributes = None,
        back_compat_name: str = "",
    ) -> None:
        """
        Record a new value for a Gauge.

        :param stat: The name of the stat to update.
        :param value: The new value of stat, either a float or an int.
        :param rate: value between 0 and 1 that represents the sample rate at
            which the metric is going to be emitted.
        :param delta: If true, the provided value will be added to the previous value.
            If False the new value will override the previous.
        :param tags: Tags to append to the stat.
        :param back_compat_name:  If an alternative name is provided, the
            stat will be emitted using both names if possible.
        """
        if _skip_due_to_rate(rate):
            return

        if back_compat_name and self.metrics_validator.test(back_compat_name):
            self.metrics_map.set_gauge_value(
                full_name(prefix=self.prefix, name=back_compat_name), value, delta, tags
            )

        if self.metrics_validator.test(stat):
            self.metrics_map.set_gauge_value(full_name(prefix=self.prefix, name=stat), value, delta, tags)

    def timing(
        self,
        stat: str,
        dt: DeltaType,
        *,
        tags: Attributes = None,
    ) -> None:
        """Record a timing observation as a Histogram to preserve distribution information."""
        if self.metrics_validator.test(stat) and name_is_otel_safe(self.prefix, stat):
            if isinstance(dt, datetime.timedelta):
                dt = dt.total_seconds() * 1000.0
            self.metrics_map.record_histogram_value(full_name(prefix=self.prefix, name=stat), float(dt), tags)

    def timer(
        self,
        stat: str | None = None,
        *args,
        tags: Attributes = None,
        **kwargs,
    ) -> Timer:
        """Timer context manager returns the duration and can be cancelled."""
        return _OtelTimer(self, stat, tags)


class InternalGauge:
    """Stores sync gauge instrument and current value to support delta feature."""

    def __init__(self, meter, name: str, tags: Attributes):
        self.attributes = tags
        otel_safe_name = _get_otel_safe_name(name)
        self.gauge = meter.create_gauge(name=otel_safe_name)
        log.debug("Created %s as type: %s", otel_safe_name, _type_as_str(self.gauge))
        self.value = DEFAULT_GAUGE_VALUE
        self.gauge.set(self.value, attributes=self.attributes)

    def set_value(self, new_value: int | float, delta: bool):
        """Delta feature to increase old value with new value and metric export."""
        if delta:
            new_value += self.value
        self.value = new_value
        self.gauge.set(new_value, attributes=self.attributes)


class InternalHistogram:
    """Stores a histogram instrument for timer/timing metrics."""

    def __init__(self, meter, name: str):
        otel_safe_name = _get_otel_safe_name(name)
        self.histogram = meter.create_histogram(name=otel_safe_name, unit="ms")
        log.debug("Created %s as type: %s", otel_safe_name, _type_as_str(self.histogram))

    def record(self, value: float, tags: Attributes) -> None:
        self.histogram.record(value, attributes=tags)


class MetricsMap:
    """Stores Otel Instruments."""

    def __init__(self, meter):
        self.meter = meter
        self.map = {}
        self.histograms: dict[str, InternalHistogram] = {}

    def clear(self) -> None:
        self.map.clear()
        self.histograms.clear()

    def _create_counter(self, name):
        """Create a new counter or up_down_counter for the provided name."""
        otel_safe_name = _get_otel_safe_name(name)

        if _is_up_down_counter(name):
            counter = self.meter.create_up_down_counter(name=otel_safe_name)
        else:
            counter = self.meter.create_counter(name=otel_safe_name)

        log.debug("Created %s as type: %s", otel_safe_name, _type_as_str(counter))
        return counter

    def get_counter(self, name: str, attributes: Attributes = None):
        """
        Return the counter; creates a new one if it did not exist.

        :param name: The name of the counter to fetch or create.
        :param attributes:  Counter attributes, used to generate a unique key to store the counter.
        """
        key = _generate_key_name(name, attributes)
        if key not in self.map:
            self.map[key] = self._create_counter(name)
        return self.map[key]

    def del_counter(self, name: str, attributes: Attributes = None) -> None:
        """
        Delete a counter.

        :param name: The name of the counter to delete.
        :param attributes: Counter attributes which were used to generate a unique key to store the counter.
        """
        key = _generate_key_name(name, attributes)
        if key in self.map.keys():
            del self.map[key]

    def set_gauge_value(self, name: str, value: int | float, delta: bool, tags: Attributes):
        """
        Override the last reading for a Gauge with a new value.

        :param name: The name of the gauge to record.
        :param value: The new reading to record.
        :param delta: If True, value is added to the previous reading, else it overrides.
        :param tags: Gauge attributes which were used to generate a unique key to store the counter.
        :returns: None
        """
        key: str = _generate_key_name(name, tags)

        if key not in self.map:
            self.map[key] = InternalGauge(meter=self.meter, name=name, tags=tags)

        self.map[key].set_value(value, delta)

    def record_histogram_value(self, name: str, value: float, tags: Attributes) -> None:
        """
        Record a timing observation in a Histogram instrument.

        Unlike a Gauge, a Histogram accumulates all observations so that count, sum,
        and bucket distribution are preserved across multiple recordings.

        :param name: The name of the histogram to record.
        :param value: The timing observation in milliseconds.
        :param tags: Attributes to attach to the observation.
        """
        if name not in self.histograms:
            self.histograms[name] = InternalHistogram(meter=self.meter, name=name)
        self.histograms[name].record(value, tags)


def flush_otel_metrics():
    provider = metrics.get_meter_provider()
    provider.force_flush()


def atexit_register_metrics_flush():
    atexit.register(flush_otel_metrics)


def get_otel_logger(
    *,
    host: str | None = None,
    port: int | None = None,
    prefix: str | None = None,
    ssl_active: bool = False,
    conf_interval: float | None = None,
    debug: bool = False,
    service_name: str | None = None,
    metrics_allow_list: str | None = None,
    metrics_block_list: str | None = None,
    stat_name_handler: Callable[[str], str] | None = None,
    statsd_influxdb_enabled: bool = False,
) -> SafeOtelLogger:
    """
    Build and return a :class:`SafeOtelLogger` backed by a configured :class:`MeterProvider`.

    Histogram instruments (used for ``timing()`` / ``timer()`` metrics) are aggregated with
    :class:`~opentelemetry.sdk.metrics.view.ExponentialBucketHistogramAggregation`
    so that bucket boundaries adapt automatically to the observed data range.  This avoids
    the need to hand-tune explicit bucket boundaries for metrics that span very different
    scales (milliseconds to hours).
    """
    otel_env_config = load_metrics_env_config()

    effective_service_name: str = otel_env_config.service_name or service_name or "airflow"
    effective_prefix: str = prefix or DEFAULT_METRIC_NAME_PREFIX
    resource = Resource.create(attributes={SERVICE_NAME: effective_service_name})

    # https://opentelemetry.io/docs/specs/otel/configuration/sdk-environment-variables/#periodic-exporting-metricreader
    interval = otel_env_config.interval_ms or conf_interval

    metric_exporter = get_otel_data_exporter(
        otel_env_config=otel_env_config,
        host=host,
        port=port,
        ssl_active=ssl_active,
    )

    readers = [
        PeriodicExportingMetricReader(
            exporter=metric_exporter,  # type: ignore[arg-type]
            export_interval_millis=interval,  # type: ignore[arg-type]
        )
    ]

    if otel_env_config.exporter:
        debug = otel_env_config.exporter == "console"

    if debug:
        export_to_console = PeriodicExportingMetricReader(
            ConsoleMetricExporter(),
            export_interval_millis=interval,
        )
        readers.append(export_to_console)

    # Reset the OTel SDK's Once() guard so set_meter_provider() can succeed.
    # This is necessary when get_otel_logger() is called after a process fork:
    # the parent's _METER_PROVIDER_SET_ONCE._done = True is inherited by the child,
    # causing set_meter_provider() to silently fail with "Overriding of current
    # MeterProvider is not allowed". The child then uses the parent's stale provider
    # whose PeriodicExportingMetricReader thread is dead after fork.
    # On first call (no fork), _done is already False so this is a no-op.
    # See: https://github.com/apache/airflow/issues/64690
    try:
        import opentelemetry.metrics._internal as _metrics_internal

        _metrics_internal._METER_PROVIDER_SET_ONCE._done = False
        _metrics_internal._METER_PROVIDER = None
    except (ImportError, AttributeError):
        pass

    metrics.set_meter_provider(
        MeterProvider(
            resource=resource,
            metric_readers=readers,
            views=[
                View(
                    instrument_type=metrics.Histogram,
                    aggregation=ExponentialBucketHistogramAggregation(),
                )
            ],
            shutdown_on_exit=False,
        ),
    )

    # Register a hook that flushes any in-memory metrics at shutdown.
    atexit_register_metrics_flush()

    validator = get_validator(metrics_allow_list, metrics_block_list)

    return SafeOtelLogger(
        metrics.get_meter_provider(), effective_prefix, validator, stat_name_handler, statsd_influxdb_enabled
    )
