package middleware

import (
	"crypto/rand"
	"encoding/hex"
	"log/slog"

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

const traceIDKey = "trace_id"

// Trace 透传或生成 trace_id，并注入可携带 trace_id 的 slog 日志句柄，
// 供后续日志统一携带链路标识。
func Trace() gin.HandlerFunc {
	return func(c *gin.Context) {
		tid := c.GetHeader("X-Trace-ID")
		if tid == "" {
			tid = newTraceID()
		}
		c.Set(traceIDKey, tid)
		c.Set("slog_ctx", slog.With(slog.String("trace_id", tid)))
		c.Next()
	}
}

func newTraceID() string {
	b := make([]byte, 16)
	if _, err := rand.Read(b); err != nil {
		return "unknown"
	}
	return hex.EncodeToString(b)
}

// TraceID 从 gin.Context 取出 trace_id。
func TraceID(c *gin.Context) string {
	if v, ok := c.Get(traceIDKey); ok {
		if s, ok := v.(string); ok {
			return s
		}
	}
	return ""
}

// SlogCtx 返回携带 trace_id 的 slog 日志句柄（无则回退默认）。
func SlogCtx(c *gin.Context) *slog.Logger {
	if v, ok := c.Get("slog_ctx"); ok {
		if l, ok := v.(*slog.Logger); ok {
			return l
		}
	}
	return slog.Default()
}