package middleware

import (
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"time"

	"github.com/gin-gonic/gin"
)

const (
	// TraceIDHeader 追踪 ID 的 HTTP 响应头名称
	TraceIDHeader = "X-Trace-ID"
	// TraceIDKey gin.Context 中的 key
	TraceIDKey = "trace_id"
)

// TraceID 中间件：为每个请求生成唯一追踪 ID
// 追踪 ID 写入 gin.Context 和响应头，供日志和下游服务关联
func TraceID() gin.HandlerFunc {
	return func(c *gin.Context) {
		// 优先复用上游传入的 Trace ID（如 API 网关、前端）
		traceID := c.GetHeader(TraceIDHeader)
		if traceID == "" {
			traceID = generateTraceID()
		}

		c.Set(TraceIDKey, traceID)
		c.Header(TraceIDHeader, traceID)

		c.Next()
	}
}

// generateTraceID 生成 16 字节随机 hex 字符串，格式：时间前缀+随机后缀
// 例：20260405-a1b2c3d4e5f60718
func generateTraceID() string {
	b := make([]byte, 8)
	_, err := rand.Read(b)
	if err != nil {
		// 降级：使用时间戳
		return fmt.Sprintf("%d", time.Now().UnixNano())
	}
	return fmt.Sprintf("%s-%s", time.Now().Format("20060102"), hex.EncodeToString(b))
}

// GetTraceID 从 gin.Context 获取 trace ID（供 handler 使用）
func GetTraceID(c *gin.Context) string {
	if v, exists := c.Get(TraceIDKey); exists {
		if id, ok := v.(string); ok {
			return id
		}
	}
	return ""
}