package litellm

import (
	"context"
	"fmt"
	"os"
	"time"

	"go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp"
	"go.opentelemetry.io/otel/sdk/resource"
	sdktrace "go.opentelemetry.io/otel/sdk/trace"
	semconv "go.opentelemetry.io/otel/semconv/v1.36.0"
	"google.golang.org/adk/telemetry"
)

// SetupTelemetry wires ADK OTLP tracing from environment variables.
//
// Reads:
//   - OTLP_ENDPOINT       full URL base, e.g. http://jaeger.d-ma.be:4318 (skip if empty)
//   - ADK_SERVICE_NAME    service name in Jaeger (default: "agent")
//   - ADK_SERVICE_VERSION semver label (default: "0.1.0")
//
// Returns a shutdown func to call on exit with a short-timeout context.
// No-op (nil error, noop shutdown) when OTLP_ENDPOINT is unset.
func SetupTelemetry(ctx context.Context) (shutdown func(context.Context) error, err error) {
	endpoint := os.Getenv("OTLP_ENDPOINT")
	if endpoint == "" {
		return func(context.Context) error { return nil }, nil
	}

	svcName := os.Getenv("ADK_SERVICE_NAME")
	if svcName == "" {
		svcName = "agent"
	}
	svcVersion := os.Getenv("ADK_SERVICE_VERSION")
	if svcVersion == "" {
		svcVersion = "0.1.0"
	}

	exporter, err := otlptracehttp.New(ctx,
		otlptracehttp.WithEndpointURL(endpoint+"/v1/traces"),
	)
	if err != nil {
		return nil, fmt.Errorf("otlp exporter: %w", err)
	}

	res, err := resource.New(ctx,
		resource.WithAttributes(
			semconv.ServiceName(svcName),
			semconv.ServiceVersion(svcVersion),
		),
	)
	if err != nil {
		return nil, fmt.Errorf("resource: %w", err)
	}

	providers, err := telemetry.New(ctx,
		telemetry.WithSpanProcessors(sdktrace.NewBatchSpanProcessor(exporter)),
		telemetry.WithResource(res),
	)
	if err != nil {
		return nil, fmt.Errorf("telemetry.New: %w", err)
	}
	providers.SetGlobalOtelProviders()

	return func(ctx context.Context) error {
		shutCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
		defer cancel()
		return providers.Shutdown(shutCtx)
	}, nil
}