# tracer_handle.py

from __future__ import annotations

import threading
from typing import Optional, List, Tuple

from opentelemetry import trace
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter


GrpcHeaders = List[Tuple[str, str]]


class TracerHandle:
    """OTLP/gRPC tracer handle with runtime auth token rotation.

    This class is a **process-wide singleton**: regardless of how many `Frisk` instances
    are created, there will be only one global tracer provider/exporter.

    Use `TracerHandle.get_instance(...)` to get the singleton.

    Shutdown is reference counted:
    - `acquire()` increments a refcount
    - `release()` decrements the refcount and only shuts down when it reaches 0

    Key point: grpc metadata must be a sequence of 2-tuples, e.g.
      [("authorization", "Bearer ...")]
    NOT a dict.
    """

    _singleton_lock = threading.RLock()
    _singleton: Optional["TracerHandle"] = None

    def __init__(
        self,
        token: str,
        *,
        endpoint: str,
        insecure: bool,
        service_name: str,
        service_namespace: str,
    ):
        # NOTE: do not call this directly; prefer get_instance().
        self._lock = threading.RLock()
        self._endpoint = endpoint
        self._insecure = insecure
        self._service_name = service_name
        self._service_namespace = service_namespace

        self._exporter: Optional[OTLPSpanExporter] = None
        self._provider: Optional[TracerProvider] = None

        self._refcount = 0
        self._closed = False

        self.set_global_tracer_provider(token)

    @classmethod
    def get_instance(
        cls,
        token: str,
        *,
        endpoint: str,
        insecure: bool,
        service_name: str,
        service_namespace: str,
    ) -> "TracerHandle":
        """Get (and possibly create) the process-wide TracerHandle singleton.

        If it already exists, we *update* the token to the latest provided value.

        If you try to get the singleton with different static config (endpoint, insecure,
        service_name, service_namespace), we raise to avoid silently misconfiguring global
        tracing.
        """
        with cls._singleton_lock:
            if cls._singleton is None:
                cls._singleton = cls(
                    token,
                    endpoint=endpoint,
                    insecure=insecure,
                    service_name=service_name,
                    service_namespace=service_namespace,
                )
            else:
                inst = cls._singleton
                # Fail fast on config mismatch: global tracer provider is shared.
                if (
                    inst._endpoint != endpoint
                    or inst._insecure != insecure
                    or inst._service_name != service_name
                    or inst._service_namespace != service_namespace
                ):
                    raise RuntimeError(
                        "TracerHandle singleton already initialized with different configuration. "
                        "Create Frisk instances with consistent telemetry configuration."
                    )
                inst.update_token(token)

            return cls._singleton

    def acquire(self) -> "TracerHandle":
        """Increment the singleton refcount (used by each Frisk instance)."""
        with self._lock:
            if self._closed:
                raise RuntimeError("TracerHandle is closed")
            self._refcount += 1
            return self

    def release(self) -> None:
        """Decrement refcount and shut down tracing only when it reaches 0."""
        with self._lock:
            if self._refcount > 0:
                self._refcount -= 1

            if self._refcount == 0:
                self._shutdown_internal()

    def update_token(self, token: str) -> None:
        """Called from your token-changed callback.

        Mutates exporter headers in a thread-safe way.
        """
        with self._lock:
            if self._closed:
                return

            if self._exporter is None:
                self.set_global_tracer_provider(token)
                return

            headers: GrpcHeaders = [("authorization", f"Bearer {token}")]

            exp = self._exporter

            # Different versions store headers differently; handle the common cases.
            if hasattr(exp, "_headers"):
                setattr(exp, "_headers", headers)
                return

            if hasattr(exp, "headers"):
                setattr(exp, "headers", headers)
                return

            raise RuntimeError(
                "Exporter headers field not found; cannot update token dynamically."
            )

    def set_global_tracer_provider(self, token: str) -> None:
        with self._lock:
            if self._closed:
                return

            resource = Resource.create(
                {
                    "service.name": self._service_name,
                    "service.namespace": self._service_namespace,
                }
            )

            exporter = OTLPSpanExporter(
                endpoint=self._endpoint,
                insecure=self._insecure,
                headers=[("authorization", f"Bearer {token}")],
            )

            provider = TracerProvider(resource=resource)
            provider.add_span_processor(BatchSpanProcessor(exporter))

            trace.set_tracer_provider(provider)

            self._exporter = exporter
            self._provider = provider

    def force_flush(self) -> bool:
        with self._lock:
            if self._closed or self._provider is None:
                return True
            try:
                return bool(self._provider.force_flush())
            except Exception:
                return False

    def _shutdown_internal(self) -> None:
        if self._closed:
            return

        provider = self._provider
        self._provider = None
        self._exporter = None
        self._closed = True

        if provider is None:
            return

        provider.shutdown()

        # Allow re-init after full shutdown.
        with TracerHandle._singleton_lock:
            if TracerHandle._singleton is self:
                TracerHandle._singleton = None

    def shutdown(self) -> None:
        """Backward-compatible alias.

        Prefer `release()` for singleton/refcounted usage.
        """
        self.release()
