from __future__ import annotations

from dataclasses import dataclass
from typing import Optional

import grpc

from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.trace import Tracer

from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter


@dataclass(frozen=True, slots=True)
class TracingManagerOptions:
    """Configuration container for :class:`~frisk_sdk.telemetry.tracing_manager.TracingManager`.

    This dataclass is a lightweight way to bundle the tracing configuration used
    to construct a :class:`~frisk_sdk.telemetry.tracing_manager.TracingManager`.
    It’s useful when you want to pass around tracing configuration as a single
    object (e.g., in tests, factory helpers, or higher-level SDK wiring).

    Attributes
    ----------
    otlp_endpoint:
        OTLP endpoint intended for **gRPC** export. It may include an HTTP-style
        ``/v1/traces`` suffix; the manager will normalize it for the gRPC exporter.
        Examples: ``"https://otel.example.com:4317"``,
        ``"http://localhost:4317/v1/traces"``.

    auth_token:
        Bearer token used for export authorization. This value can be rotated
        later via :meth:`TracingManager.update_auth_token`.

    tracer_name:
        Tracer name passed to ``TracerProvider.get_tracer(name, version)``.

    tracer_version:
        Tracer version passed to ``TracerProvider.get_tracer(name, version)``.

    Notes
    -----
    ``TracingManager`` currently expects keyword arguments, so you’ll typically
    expand this options object when constructing the manager:

    >>> opts = TracingManagerOptions(
    ...     otlp_endpoint="http://localhost:4317",
    ...     auth_token="t0",
    ...     tracer_name="frisk",
    ...     tracer_version="1.0.0",
    ... )
    >>> mgr = TracingManager(
    ...     otlp_endpoint=opts.otlp_endpoint,
    ...     auth_token=opts.auth_token,
    ...     tracer_name=opts.tracer_name,
    ...     tracer_version=opts.tracer_version,
    ... )
    """

    otlp_endpoint: str
    auth_token: str
    tracer_name: str
    tracer_version: str


class _DynamicAuthMetadataPlugin(grpc.AuthMetadataPlugin):
    """
    gRPC calls this for each RPC to produce metadata. We read the current token
    from the manager so updates apply to future exports without recreating exporter.
    """

    def __init__(self, token_getter):
        self._token_getter = token_getter

    def __call__(self, context, callback):
        token = self._token_getter()
        # grpc-python expects sequence of 2-tuples
        callback((("authorization", f"Bearer {token}"),), None)


class TracingManager:
    """Manage an OpenTelemetry tracer provider for the SDK.

    This class owns the lifecycle of an OTEL ``TracerProvider`` configured with an
    OTLP/gRPC span exporter. It's designed for SDK usage where you want tracing to
    be *instance-scoped* (per ``Frisk`` client/session) rather than setting a
    process-wide global tracer provider.

    Key behaviors
    ==========
    - **No global state:** does not call ``opentelemetry.trace.set_tracer_provider``.
      The returned tracer comes from this instance's provider only.

    - **Dynamic auth metadata:** exports use a gRPC metadata plugin that reads the
      *current* auth token at export time. Calling :meth:`update_auth_token`
      updates the token used for *future* export RPCs without recreating the
      exporter or provider.

    - **Endpoint normalization:** :meth:`_normalize_otel_grpc_endpoint` strips an
      HTTP OTLP traces path suffix (``/v1/traces``) when present, since OTLP/gRPC
      expects a host/port target rather than an HTTP path.

    - **TLS by default:** the exporter is configured with TLS channel credentials
      and per-RPC call credentials. If you need custom trust roots or mTLS, pass
      PEM bytes via ``root_certificates_pem`` / ``private_key_pem`` /
      ``certificate_chain_pem``.

    Parameters
    ==========
    otlp_endpoint:
        OTLP endpoint for gRPC export. May be a plain host/port target or include
        a ``/v1/traces`` suffix (which will be normalized for the gRPC exporter).
    auth_token:
        Bearer token used for authorization metadata.
    tracer_name, tracer_version:
        Values passed to ``TracerProvider.get_tracer(name, version)``.
    timeout_seconds:
        Optional exporter timeout.

    Examples
    =========
    Create a manager, grab a tracer, and shut it down when you're done::

        mgr = TracingManager(
            otlp_endpoint="https://otel.example.com:4317",
            auth_token="token",
            tracer_name="frisk",
            tracer_version="1.0.0",
        )
        tracer = mgr.get_tracer()
        mgr.update_auth_token("new-token")
        mgr.shutdown()

    Notes
    =====
    ``shutdown()`` flushes and shuts down the underlying provider synchronously.
    It's safe to call from normal teardown paths.
    """

    def __init__(
        self,
        *,
        otlp_endpoint: str,
        auth_token: str,
        tracer_name: str,
        tracer_version: str,
        # TLS is required for call-credentials in most gRPC setups
        root_certificates_pem: Optional[bytes] = None,
        private_key_pem: Optional[bytes] = None,
        certificate_chain_pem: Optional[bytes] = None,
        # exporter knobs (optional)
        timeout_seconds: Optional[float] = None,
    ):
        self.otlp_endpoint = otlp_endpoint
        self.auth_token = auth_token
        self.tracer_name = tracer_name
        self.tracer_version = tracer_version

        self._timeout_seconds = timeout_seconds

        self._credentials = self._build_credentials(
            root_certificates_pem=root_certificates_pem,
            private_key_pem=private_key_pem,
            certificate_chain_pem=certificate_chain_pem,
        )

        self._tracer_provider = self._create_tracer_provider()
        self._shutdown = False

    def shutdown(self) -> None:
        if self._shutdown:
            return
        # TracerProvider.shutdown() is sync in Python
        self._tracer_provider.shutdown()
        self._shutdown = True

    def get_tracer(self) -> Tracer:
        return self._tracer_provider.get_tracer(self.tracer_name, self.tracer_version)

    def update_auth_token(self, token: str) -> None:
        self.auth_token = token

    # ----- internals -----

    def _get_current_token(self) -> str:
        return self.auth_token

    @staticmethod
    def _normalize_otel_grpc_endpoint(endpoint: str) -> str:
        # For gRPC we don't use /v1/traces (HTTP only)
        if endpoint.endswith("/v1/traces"):
            return endpoint[: -len("/v1/traces")]
        # grpc exporter generally accepts host:port OR "http(s)://host:port" in some versions,
        # but host:port is the safest.
        return endpoint

    def _build_credentials(
        self,
        *,
        root_certificates_pem: Optional[bytes],
        private_key_pem: Optional[bytes],
        certificate_chain_pem: Optional[bytes],
    ):
        # Channel TLS credentials (like credentials.createSsl())
        channel_creds = grpc.ssl_channel_credentials(
            root_certificates=root_certificates_pem,
            private_key=private_key_pem,
            certificate_chain=certificate_chain_pem,
        )

        # Per-RPC call credentials (like createFromMetadataGenerator)
        plugin = _DynamicAuthMetadataPlugin(self._get_current_token)
        call_creds = grpc.metadata_call_credentials(plugin)

        # Combine them (like combineChannelCredentials)
        return grpc.composite_channel_credentials(channel_creds, call_creds)

    def _create_tracer_provider(self) -> TracerProvider:
        endpoint = self._normalize_otel_grpc_endpoint(self.otlp_endpoint)

        exporter = OTLPSpanExporter(
            endpoint=endpoint,
            insecure=False,  # we're using TLS creds
            credentials=self._credentials,
            timeout=self._timeout_seconds,
        )

        resource = Resource.create(
            {
                "service.name": self.tracer_name,
                "service.version": self.tracer_version,
            }
        )

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