package logging

import (
	"context"
	"log/slog"
)

type contextKey int

const (
	loggerKey contextKey = iota
	requestIDKey
	userIDKey
	traceIDKey
)

// NewContext returns a new context with the logger attached.
func NewContext(ctx context.Context, logger *Logger) context.Context {
	return context.WithValue(ctx, loggerKey, logger)
}

// FromContext extracts the logger from the context.
// Returns a no-op logger if none is found.
func FromContext(ctx context.Context) *Logger {
	if logger, ok := ctx.Value(loggerKey).(*Logger); ok {
		return logger
	}
	return Nop()
}

// WithRequestID adds a request ID to the context.
func WithRequestID(ctx context.Context, requestID string) context.Context {
	return context.WithValue(ctx, requestIDKey, requestID)
}

// RequestIDFromContext extracts the request ID from the context.
func RequestIDFromContext(ctx context.Context) string {
	if id, ok := ctx.Value(requestIDKey).(string); ok {
		return id
	}
	return ""
}

// WithUserID adds a user ID to the context.
func WithUserID(ctx context.Context, userID string) context.Context {
	return context.WithValue(ctx, userIDKey, userID)
}

// UserIDFromContext extracts the user ID from the context.
func UserIDFromContext(ctx context.Context) string {
	if id, ok := ctx.Value(userIDKey).(string); ok {
		return id
	}
	return ""
}

// WithTraceID adds a trace ID to the context.
func WithTraceID(ctx context.Context, traceID string) context.Context {
	return context.WithValue(ctx, traceIDKey, traceID)
}

// TraceIDFromContext extracts the trace ID from the context.
func TraceIDFromContext(ctx context.Context) string {
	if id, ok := ctx.Value(traceIDKey).(string); ok {
		return id
	}
	return ""
}

// ContextAttrs returns slog attributes from context values.
func ContextAttrs(ctx context.Context) []slog.Attr {
	var attrs []slog.Attr

	if id := RequestIDFromContext(ctx); id != "" {
		attrs = append(attrs, slog.String("request_id", id))
	}
	if id := UserIDFromContext(ctx); id != "" {
		attrs = append(attrs, slog.String("user_id", id))
	}
	if id := TraceIDFromContext(ctx); id != "" {
		attrs = append(attrs, slog.String("trace_id", id))
	}

	return attrs
}

// LoggerWithContext returns a logger enriched with context attributes.
func LoggerWithContext(ctx context.Context, logger *Logger) *Logger {
	attrs := ContextAttrs(ctx)
	if len(attrs) == 0 {
		return logger
	}

	args := make([]any, 0, len(attrs)*2)
	for _, attr := range attrs {
		args = append(args, attr.Key, attr.Value.Any())
	}
	return logger.With(args...)
}