fetch_ml/internal/logging/logging.go

// Package logging provides structured logging utilities with trace context support.
package logging

import (
	"context"
	"io"
	"log/slog"
	"os"
	"path/filepath"
	"time"

	"github.com/google/uuid"
	"github.com/jfraeys/fetch_ml/internal/fileutil"
)

type ctxKey string

// Context keys for trace and span information
const (
	CtxTraceID ctxKey = "trace_id"
	CtxSpanID  ctxKey = "span_id"
	CtxWorker  ctxKey = "worker_id"
	CtxJob     ctxKey = "job_name"
	CtxTask    ctxKey = "task_id"
)

// Logger wraps slog.Logger with additional context handling capabilities.
type Logger struct {
	*slog.Logger
}

// NewLogger creates a logger that writes to stderr (development mode)
func NewLogger(level slog.Level, jsonOutput bool) *Logger {
	opts := &slog.HandlerOptions{
		Level:     level,
		AddSource: os.Getenv("LOG_ADD_SOURCE") == "1",
	}

	var handler slog.Handler
	if jsonOutput || os.Getenv("LOG_FORMAT") == "json" {
		handler = slog.NewJSONHandler(os.Stderr, opts)
	} else {
		handler = NewColorTextHandler(os.Stderr, opts)
	}

	return &Logger{slog.New(handler)}
}

// NewFileLogger creates a logger that writes to a file only (production mode)
func NewFileLogger(level slog.Level, jsonOutput bool, logFile string) *Logger {
	opts := &slog.HandlerOptions{
		Level:     level,
		AddSource: os.Getenv("LOG_ADD_SOURCE") == "1",
	}

	// Create log directory if it doesn't exist
	if logFile != "" {
		logDir := filepath.Dir(logFile)
		if err := os.MkdirAll(logDir, 0750); err != nil {
			// Fallback to stderr only if directory creation fails
			return NewLogger(level, jsonOutput)
		}
	}

	// Open log file
	file, err := fileutil.SecureOpenFile(logFile, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0600)
	if err != nil {
		// Fallback to stderr only if file creation fails
		return NewLogger(level, jsonOutput)
	}

	// Write to file only (production)
	var handler slog.Handler
	if jsonOutput || os.Getenv("LOG_FORMAT") == "json" {
		handler = slog.NewJSONHandler(file, opts)
	} else {
		handler = slog.NewTextHandler(file, opts)
	}

	return &Logger{slog.New(handler)}
}

// EnsureTrace injects trace and span IDs if missing from context.
func EnsureTrace(ctx context.Context) context.Context {
	if ctx.Value(CtxTraceID) == nil {
		ctx = context.WithValue(ctx, CtxTraceID, uuid.NewString())
	}
	if ctx.Value(CtxSpanID) == nil {
		ctx = context.WithValue(ctx, CtxSpanID, uuid.NewString())
	}
	return ctx
}

// WithContext returns a new Logger with context values added as attributes.
func (l *Logger) WithContext(ctx context.Context, args ...any) *Logger {
	if trace := ctx.Value(CtxTraceID); trace != nil {
		args = append(args, "trace_id", trace)
	}
	if span := ctx.Value(CtxSpanID); span != nil {
		args = append(args, "span_id", span)
	}
	if worker := ctx.Value(CtxWorker); worker != nil {
		args = append(args, "worker_id", worker)
	}
	if job := ctx.Value(CtxJob); job != nil {
		args = append(args, "job_name", job)
	}
	if task := ctx.Value(CtxTask); task != nil {
		args = append(args, "task_id", task)
	}
	return &Logger{Logger: l.With(args...)}
}

// CtxWithWorker adds worker ID to context.
func CtxWithWorker(ctx context.Context, worker string) context.Context {
	return context.WithValue(ctx, CtxWorker, worker)
}

// CtxWithJob adds job name to context.
func CtxWithJob(ctx context.Context, job string) context.Context {
	return context.WithValue(ctx, CtxJob, job)
}

// CtxWithTask adds task ID to context.
func CtxWithTask(ctx context.Context, task string) context.Context {
	return context.WithValue(ctx, CtxTask, task)
}

// Component returns a new Logger with component name added.
func (l *Logger) Component(ctx context.Context, name string) *Logger {
	return l.WithContext(ctx, "component", name)
}

// Worker returns a new Logger with worker ID added.
func (l *Logger) Worker(ctx context.Context, workerID string) *Logger {
	return l.WithContext(ctx, "worker_id", workerID)
}

// Job returns a new Logger with job name and task ID added.
func (l *Logger) Job(ctx context.Context, job string, task string) *Logger {
	return l.WithContext(ctx, "job_name", job, "task_id", task)
}

// Fatal logs an error message and exits with status 1.
func (l *Logger) Fatal(msg string, args ...any) {
	l.Error(msg, args...)
	os.Exit(1)
}

// Panic logs an error message and panics.
func (l *Logger) Panic(msg string, args ...any) {
	l.Error(msg, args...)
	panic(msg)
}

// ColorTextHandler provides colorized console log output.
type ColorTextHandler struct {
	slog.Handler
}

// NewColorTextHandler creates a new colorized text handler.
func NewColorTextHandler(w io.Writer, opts *slog.HandlerOptions) slog.Handler {
	base := slog.NewTextHandler(w, opts)
	return &ColorTextHandler{Handler: base}
}

// Handle processes a log record with color formatting.
func (h *ColorTextHandler) Handle(ctx context.Context, r slog.Record) error {
	// Add uniform timestamp (override default)
	r.Time = time.Now()

	switch r.Level {
	case slog.LevelDebug:
		r.Add("lvl_color", "\033[34mDBG\033[0m")
	case slog.LevelInfo:
		r.Add("lvl_color", "\033[32mINF\033[0m")
	case slog.LevelWarn:
		r.Add("lvl_color", "\033[31mERR\033[0m")
	case slog.LevelError:
		r.Add("lvl_color", "\033[33mWRN\033[0m")
	}

	return h.Handler.Handle(ctx, r)
}