"""HTTP request -> Convex client spans. No payload/header/URL capture."""
import asyncio
import os

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.sdk.trace.sampling import ParentBased, TraceIdRatioBased
from opentelemetry.trace import SpanKind, Status, StatusCode
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
from quart import g, request

_PROPAGATOR = TraceContextTextMapPropagator()

class RequestTracingMiddleware:
  def __init__(self, app, tracer):
    self.app = app
    self.tracer = tracer

  async def __call__(self, scope, receive, send):
    if scope["type"] != "http":
      return await self.app(scope, receive, send)

    carrier = {
      key.decode("latin1"): value.decode("latin1")
      for key, value in scope.get("headers", [])
      if key.lower() in (b"traceparent", b"tracestate")
    }
    method = scope.get("method", "HTTP")
    with self.tracer.start_as_current_span(
      f"HTTP {method}",
      context=_PROPAGATOR.extract(carrier),
      kind=SpanKind.SERVER,
      attributes={"http.request.method": method},
      record_exception=False,
      set_status_on_exception=False,
    ) as span:
      async def traced_send(message):
        if message["type"] == "http.response.start":
          status = message["status"]
          span.set_attribute("http.response.status_code", status)
          if status >= 500:
            span.set_status(Status(StatusCode.ERROR))
          trace_id = f"{span.get_span_context().trace_id:032x}"
          headers = [(k, v) for k, v in message.get("headers", []) if k.lower() != b"x-trace-id"]
          message = {**message, "headers": [*headers, (b"x-trace-id", trace_id.encode("ascii"))]}
        await send(message)

      try:
        return await self.app(scope, receive, traced_send)
      except BaseException as error:
        span.set_attribute("error.type", type(error).__name__)
        span.set_status(Status(StatusCode.ERROR))
        raise

def register_request_tracing(app, *, provider=None):
  """Disabled without an OTLP endpoint. Inject a provider for isolated tests."""
  if provider is None:
    if not (os.getenv("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT") or os.getenv("OTEL_EXPORTER_OTLP_ENDPOINT")):
      return None
    from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter

    sample_rate = float(os.getenv("OTEL_TRACES_SAMPLER_ARG", "0.1"))
    if not 0 <= sample_rate <= 1:
      raise ValueError("OTEL_TRACES_SAMPLER_ARG must be between 0 and 1")
    provider = TracerProvider(
      resource=Resource.create({"service.name": os.getenv("OTEL_SERVICE_NAME") or os.getenv("SERVICE_NAME", "nanoshare")}),
      sampler=ParentBased(TraceIdRatioBased(sample_rate)),
    )
    provider.add_span_processor(BatchSpanProcessor(
      OTLPSpanExporter(timeout=2),
      max_queue_size=2048,
      max_export_batch_size=256,
      schedule_delay_millis=1000,
    ))
    trace.set_tracer_provider(provider)

    @app.after_serving
    async def shutdown_tracing():
      # SDK uses a bounded background queue, never an export on the request path.
      await asyncio.to_thread(provider.shutdown)

  app.asgi_app = RequestTracingMiddleware(app.asgi_app, provider.get_tracer(__name__))

  @app.before_request
  async def correlate_request():
    span = trace.get_current_span()
    context = span.get_span_context()
    event = getattr(g, "wide_event", None)
    if event is not None and context.is_valid:
      event["trace_id"] = f"{context.trace_id:032x}"
      request_id = event.get("request_id")
      if request_id:
        span.set_attribute("request.id", request_id)
    # Only the route template, not the raw path, query string or user values.
    if request.url_rule is not None:
      route = request.url_rule.rule
      span.update_name(f"{request.method} {route}")
      span.set_attribute("http.route", route)

  return provider