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
from opentelemetry.instrumentation.threading import ThreadingInstrumentor

ThreadingInstrumentor().instrument()


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)