# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd.
# SPDX-License-Identifier: AGPL-3.0

"""
Global metrics bootstrap and event subscription.

This module provides:
- A process-global MetricRegistry and PrometheusExporter instance for the server.
- A metrics subscriber that consumes shared observability events without making
  DataSources depend on MetricRegistry internals.

Why an event router?
Business code (models, encryption, HTTP middleware, etc.) calls DataSource APIs such as
`VLMEventDataSource.record_call(...)`. DataSources publish observability events. Metrics
collectors subscribe to those events and remain the only layer allowed to write into
MetricRegistry, keeping the architecture consistent while letting other subscribers consume
the same signal.
"""

from __future__ import annotations

import asyncio
import inspect
import threading
from typing import Optional

from openviking.observability.events import (
    make_payload_subscriber,
    register_event_subscriber,
    try_publish_event,
    unregister_event_subscriber,
)
from openviking.server.config import ServerConfig
from openviking_cli.utils import get_logger

from .account_dimension import (
    configure_metric_account_dimension as _configure_metric_account_dimension_runtime,
)
from .account_dimension import (
    reset_metric_account_dimension,
)
from .bootstrap import create_default_collector_manager
from .collectors.cache import CacheCollector
from .collectors.embedding import EmbeddingCollector
from .collectors.encryption import EncryptionCollector
from .collectors.http import HTTPCollector
from .collectors.rerank import RerankCollector
from .collectors.resource import ResourceIngestionCollector
from .collectors.retrieval import RetrievalCollector
from .collectors.session import SessionCollector
from .collectors.telemetry_bridge import TelemetryBridgeCollector
from .collectors.vlm import VLMCollector
from .core.registry import MetricRegistry
from .core.runtime import EventCollectorRouter
from .exporters.otel import OTelMetricExporter
from .exporters.prometheus import PrometheusExporter

logger = get_logger(__name__)

_lock = threading.Lock()
_registry: Optional[MetricRegistry] = None
_exporters: list = []
_event_router: EventCollectorRouter | None = None
_METRICS_EVENT_SUBSCRIBER = "metrics"


def _shutdown_exporters_best_effort(exporters: list) -> None:
    """Shutdown exporters without raising.

    This is used both for explicit shutdown and for re-initialization, to avoid leaking
    background tasks (e.g. OTel periodic export loops) across multiple init cycles.

    Note: This is a synchronous function. For proper async shutdown in async contexts,
    use `shutdown_metrics_async()` instead, which properly awaits async shutdown methods.
    """
    for exporter in exporters:
        if hasattr(exporter, "shutdown"):
            try:
                result = exporter.shutdown()
                if inspect.isawaitable(result):
                    try:
                        loop = asyncio.get_running_loop()
                    except RuntimeError:
                        # No running event loop - use asyncio.run() to execute the coroutine
                        asyncio.run(result)
                    else:
                        # There is a running event loop. We cannot block-wait from within
                        # the event loop thread (would cause deadlock). Instead, schedule
                        # the shutdown as a task and rely on the async shutdown path
                        # (shutdown_metrics_async) for proper cleanup in async contexts.
                        #
                        # This is best-effort for synchronous re-initialization paths.
                        # The proper async shutdown path is used in app.py's lifespan.
                        loop.create_task(result)
            except Exception as e:
                logger.warning(
                    "[_shutdown_exporters_best_effort] failed to shutdown exporter: %s", e
                )


def configure_metric_account_dimension(
    *,
    enabled: bool | None = None,
    metric_allowlist: set[str] | list[str] | tuple[str, ...] | None = None,
    max_active_accounts: int | None = None,
    policy=None,
    resolver=None,
) -> None:
    """
    Configure the process-global account-dimension runtime used by collector write helpers.

    This function is a thin façade over `account_dimension.configure_metric_account_dimension`
    so callers do not need to import policy internals directly from the metrics package.
    """
    _configure_metric_account_dimension_runtime(
        enabled=enabled,
        metric_allowlist=metric_allowlist,
        max_active_accounts=max_active_accounts,
        policy=policy,
        resolver=resolver,
    )


def _create_metric_registry() -> MetricRegistry:
    """
    Create a new `MetricRegistry` instance.

    This helper centralizes the registry instantiation to avoid security issues
    introduced by bridging classes (e.g., an external registry implementation
    bypassing internal safety checks).
    """
    return MetricRegistry()


def init_metrics_from_server_config(
    config: ServerConfig, *, app=None, service=None, registry: MetricRegistry | None = None
) -> None:
    """
    Initialize the process-global metrics registry, router, and exporters from server config.

    When metrics are enabled:
    - Creates a MetricRegistry (or uses provided one)
    - Registers default collectors (State/Probe/DomainStats)
    - Builds the in-process event router for Event-style metrics
    - Creates configured exporters and stores them in app.state.metrics_exporters
    When metrics are disabled:
    - Clears all global metrics state
    - Resets account-dimension runtime to its disabled defaults
    - Detaches any exporter references previously stored on `app.state`
    """
    enabled = config.observability.metrics.enabled
    global _registry, _exporters, _event_router
    with _lock:
        if not enabled:
            _shutdown_exporters_best_effort(_exporters)
            _registry = None
            _exporters = []
            _event_router = None
            unregister_event_subscriber(_METRICS_EVENT_SUBSCRIBER)
            reset_metric_account_dimension()
            if app is not None:
                app.state.metrics_exporters = []
            return

        # If re-initializing while already enabled, clean up existing exporters first.
        if _exporters:
            _shutdown_exporters_best_effort(_exporters)
        account_dimension = config.observability.metrics.account_dimension
        configure_metric_account_dimension(
            enabled=account_dimension.enabled,
            metric_allowlist=list(account_dimension.metric_allowlist or []),
            max_active_accounts=int(account_dimension.max_active_accounts),
        )
        _registry = registry or _create_metric_registry()
        collector_manager = create_default_collector_manager(app=app, service=service, config=config)
        _event_router = _build_event_router(_registry)
        register_event_subscriber(
            _METRICS_EVENT_SUBSCRIBER,
            make_payload_subscriber(_event_router.dispatch),
        )

        _exporters = []
        exporters_config = config.observability.metrics.exporters

        # Initialize Prometheus exporter if enabled
        if exporters_config.prometheus.enabled:
            prometheus_exporter = PrometheusExporter(
                registry=_registry, collector_manager=collector_manager
            )
            _exporters.append(prometheus_exporter)

        # Initialize OTel exporter if enabled
        if exporters_config.otel.enabled:
            otel_exporter = OTelMetricExporter(
                registry=_registry,
                collector_manager=collector_manager,
                protocol=exporters_config.otel.protocol,
                insecure=exporters_config.otel.tls.insecure,
                endpoint=exporters_config.otel.endpoint,
                service_name=exporters_config.otel.service_name,
                export_interval_ms=exporters_config.otel.export_interval_ms,
                headers=exporters_config.otel.headers,
                enabled=True,
            )
            otel_exporter.start()
            _exporters.append(otel_exporter)

        if app is not None:
            app.state.metrics_exporters = _exporters


def shutdown_metrics(*, app=None) -> None:
    """
    Clear all process-global metrics objects and detach the exporters from the application.

    Shutdown is intentionally idempotent so tests and service teardown paths can call it
    repeatedly without needing to coordinate current metrics initialization state.
    """
    global _registry, _exporters, _event_router
    with _lock:
        _shutdown_exporters_best_effort(_exporters)

        _registry = None
        _exporters = []
        _event_router = None
        unregister_event_subscriber(_METRICS_EVENT_SUBSCRIBER)
        reset_metric_account_dimension()
        if app is not None:
            app.state.metrics_exporters = []


async def shutdown_metrics_async(*, app=None) -> None:
    """Async variant of `shutdown_metrics` that awaits async exporter cleanup."""
    global _registry, _exporters, _event_router
    with _lock:
        exporters = list(_exporters)
        _registry = None
        _exporters = []
        _event_router = None
        unregister_event_subscriber(_METRICS_EVENT_SUBSCRIBER)
        reset_metric_account_dimension()
        if app is not None:
            app.state.metrics_exporters = []

    for exporter in exporters:
        if hasattr(exporter, "shutdown"):
            try:
                result = exporter.shutdown()
                if inspect.isawaitable(result):
                    await result
            except Exception as e:
                logger.warning("[shutdown_metrics_async] failed to shutdown exporter: %s", e)


def get_metrics_registry() -> MetricRegistry:
    """
    Return the process-global `MetricRegistry`.

    Raises:
        RuntimeError: If metrics have not been initialized yet or were explicitly disabled.
    """
    reg = _registry
    if reg is None:
        raise RuntimeError("metrics registry not initialized")
    return reg


def try_get_metrics_registry() -> MetricRegistry | None:
    """
    Return the global MetricRegistry if metrics are initialized, otherwise `None`.

    This helper is intended for best-effort call sites that want to avoid raising an exception
    when metrics are disabled.
    """
    return _registry


def try_dispatch_event(event_name: str, payload: dict) -> None:
    """
    Publish an in-process observability event.

    Kept for older call sites that imported the metrics API directly. The event is now
    published to the shared bus, so non-metrics subscribers can consume it even when
    Prometheus metrics are disabled.
    """
    try_publish_event(event_name, payload)


def _build_event_router(registry: MetricRegistry) -> EventCollectorRouter:
    """
    Build the default in-process event router for event-driven metrics.

    The router binds each supported event name to the corresponding collector instance and
    ensures every handler receives the shared registry object that backs the active exporter.
    """
    router = EventCollectorRouter()
    http_collector = HTTPCollector()
    cache_collector = CacheCollector()
    embedding_collector = EmbeddingCollector()
    rerank_collector = RerankCollector()
    vlm_collector = VLMCollector()
    session_collector = SessionCollector()
    resource_collector = ResourceIngestionCollector()
    retrieval_collector = RetrievalCollector()
    encryption_collector = EncryptionCollector()
    telemetry_bridge_collector = TelemetryBridgeCollector()

    def _receiver(collector, event_name: str):
        """
        Create a payload handler bound to a specific collector and event name.

        The returned handler closes over the global registry and forwards the payload to the
        collector's `receive(...)` method, preserving the architectural rule that collectors are
        the only layer allowed to write into `MetricRegistry`.
        """

        def _handle(payload):
            """Forward one event payload into the bound collector without altering the payload."""
            collector.receive(event_name, payload, registry)

        return _handle

    for event_name, collector in (
        ("http.request", http_collector),
        ("http.inflight", http_collector),
        ("cache.hit", cache_collector),
        ("cache.miss", cache_collector),
        ("embedding.call", embedding_collector),
        ("embedding.success", embedding_collector),
        ("embedding.error", embedding_collector),
        ("rerank.call", rerank_collector),
        ("vlm.call", vlm_collector),
        ("session.lifecycle", session_collector),
        ("session.contexts_used", session_collector),
        ("session.archive", session_collector),
        ("resource.stage", resource_collector),
        ("resource.wait", resource_collector),
        ("retrieval.completed", retrieval_collector),
        ("encryption.operation", encryption_collector),
        ("encryption.bytes", encryption_collector),
        ("encryption.payload_size", encryption_collector),
        ("encryption.auth_failed", encryption_collector),
        ("encryption.key_derivation", encryption_collector),
        ("encryption.key_load", encryption_collector),
        ("encryption.key_cache_hit", encryption_collector),
        ("encryption.key_cache_miss", encryption_collector),
        ("encryption.key_version_usage", encryption_collector),
        ("telemetry.summary", telemetry_bridge_collector),
    ):
        router.register(event_name, _receiver(collector, event_name))
    return router
