// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Author: wsfuyibing <682805@qq.com>
// Date: 2024-11-12

package common

import (
	"context"
	"fmt"
	"gitee.com/go-libs/log/conf"
	"net/http"
)

const (
	ContextKey = "__GO_LIB_LOG_TRACING__"
)

// AttachRequest
// 向 <http.Request> 添加兼容 OpenTracing 的 <common.Tracing> 参数.
func AttachRequest(ctx context.Context, req *http.Request) {
	if v := GetContextTracing(ctx); v != nil {
		cfg := conf.New()
		sid, sp := v.GetSpan()
		tid := v.GetTraceId()

		req.Header.Set(cfg.OpenTracingSampled, "1")
		req.Header.Set(cfg.OpenTracingSpanId, sid.String())
		req.Header.Set(cfg.OpenTracingSpanPosition, fmt.Sprintf(`%v`, sp))
		req.Header.Set(cfg.OpenTracingTraceId, tid.String())
	}
}

// GetContextHeader
// 从 <context.Context> 提取 <http.Header> 参数.
func GetContextHeader(ctx context.Context) http.Header {
	if v := GetContextTracing(ctx); v != nil {
		cfg := conf.New()
		sid, sp := v.GetSpan()
		tid := v.GetTraceId()

		if sp > 0 {
			sp -= 1
		}

		header := make(http.Header)
		header.Set(cfg.OpenTracingSampled, "1")
		header.Set(cfg.OpenTracingSpanId, sid.String())
		header.Set(cfg.OpenTracingSpanPosition, fmt.Sprintf(`%v`, sp))
		header.Set(cfg.OpenTracingTraceId, tid.String())
		return header
	}
	return nil
}

// GetContextTracing
// 从 <context.Context> 提取 <common.Tracing>.
func GetContextTracing(ctx context.Context) Tracing {
	if v, ok := ctx.Value(ContextKey).(Tracing); ok {
		return v
	}
	return nil
}

// SetContextTracing
// 设置 <common.Tracing> 到 <context.Context>.
func SetContextTracing(ctx context.Context, tracing Tracing) context.Context {
	return context.WithValue(ctx, ContextKey, tracing)
}

// NewContext
// 创建上下文.
func NewContext(cts ...context.Context) context.Context {
	if len(cts) > 0 && cts[0] != nil {
		if v := GetContextTracing(cts[0]); v != nil {
			return SetContextTracing(cts[0], v.Child())
		}
	}
	return SetContextTracing(context.Background(), NewTracing())
}

func NewContextWithHeader(header http.Header) context.Context {
	return SetContextTracing(context.Background(), NewTracingWithHeader(header))
}