fetch_ml/internal/logging/logging.go

package logging

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

	"github.com/google/uuid"
)

type ctxKey string

const (
	CtxTraceID ctxKey = "trace_id"
	CtxSpanID  ctxKey = "span_id"
	CtxWorker  ctxKey = "worker_id"
	CtxJob     ctxKey = "job_name"
	CtxTask    ctxKey = "task_id"
)

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, 0755); err != nil {
			// Fallback to stderr only if directory creation fails
			return NewLogger(level, jsonOutput)
		}
	}

	// Open log file
	file, err := os.OpenFile(logFile, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666)
	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)}
}

// Inject trace + span if missing
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
}

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...)}
}

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

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

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

func (l *Logger) Component(ctx context.Context, name string) *Logger {
	return l.WithContext(ctx, "component", name)
}

func (l *Logger) Worker(ctx context.Context, workerID string) *Logger {
	return l.WithContext(ctx, "worker_id", workerID)
}

func (l *Logger) Job(ctx context.Context, job string, task string) *Logger {
	return l.WithContext(ctx, "job_name", job, "task_id", task)
}

func (l *Logger) Fatal(msg string, args ...any) {
	l.Error(msg, args...)
	os.Exit(1)
}

func (l *Logger) Panic(msg string, args ...any) {
	l.Error(msg, args...)
	panic(msg)
}

// -----------------------------------------------------
// Colorized human-friendly console logs
// -----------------------------------------------------

type ColorTextHandler struct {
	slog.Handler
}

func NewColorTextHandler(w io.Writer, opts *slog.HandlerOptions) slog.Handler {
	base := slog.NewTextHandler(w, opts)
	return &ColorTextHandler{Handler: base}
}

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[33mWRN\033[0m")
	case slog.LevelError:
		r.Add("lvl_color", "\033[31mERR\033[0m")
	}

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