import os

from constant import SOFTWARE_VERSION
from opentelemetry import trace
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter as OTLPSpanHttpExporter
from opentelemetry.sdk.resources import DEPLOYMENT_ENVIRONMENT, HOST_NAME, Resource, SERVICE_NAME, SERVICE_VERSION
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor, SimpleSpanProcessor


def get_tracer(name):
    return trace.get_tracer(name)

# 初始化 OpenTelemetry
def init_opentelemetry(batch=True):
    # 设置服务名、主机名
    resource = Resource(attributes={
        SERVICE_NAME: "RENDER_WORKER",
        SERVICE_VERSION: SOFTWARE_VERSION,
        DEPLOYMENT_ENVIRONMENT: "Python",
        HOST_NAME: os.getenv("ACCESS_KEY"),
    })

    # 使用HTTP协议上报
    if batch:
        span_processor = BatchSpanProcessor(OTLPSpanHttpExporter(
            endpoint="https://oltp.jerryyan.top/v1/traces",
        ))
    else:
        span_processor = SimpleSpanProcessor(OTLPSpanHttpExporter(
            endpoint="https://oltp.jerryyan.top/v1/traces",
        ))

    trace_provider = TracerProvider(resource=resource, active_span_processor=span_processor)
    trace.set_tracer_provider(trace_provider)