zen/core/zlog/zlog.go

package zlog

import (
	"context"
	"errors"
	"io"
	"log/slog"
	"os"
	"testing"
	"time"

	"github.com/fatih/color"
	"github.com/mattn/go-isatty"
	"github.com/tidwall/pretty"
	"gitlab.bareksa.com/backend/zen/internal/bufferpool"
)

type ZLog struct {
	bufPool   *bufferpool.Pool
	jsonPool  *jsonHandlerPool
	withAttrs []slog.Attr
	withGroup []string
	opts      *ZLogOptions
	writer    WriteLocker
}

type ZLogOptions struct {
	// HandlerOptions handles the options for the underlying slog.Handler.
	//
	// If nil, default options are used.
	HandlerOptions *slog.HandlerOptions
	// Pretty enables pretty-printing of log messages.
	Pretty bool
	// Color enables colorized output.
	Color bool
	// NotificationHandler is a handler that is called when a log message is
	// intended to be sent to a notification service.
	//
	// If nil, no notification is sent.
	NotificationHandler NotificationHandler
}

// New creates a new Logger that writes to w.
//
// if opts is nil, default options are used.
func New(w io.Writer, opts *ZLogOptions) *ZLog {
	if w == nil {
		panic(errors.New("zlog: New: w is nil"))
	}
	pretty := isatty.IsTerminal(os.Stderr.Fd())
	if opts == nil {
		opts = defaultOpts
	} else {
		if opts.HandlerOptions == nil {
			opts.HandlerOptions = defaultHandlerOptions
		} else {
			opts.HandlerOptions.ReplaceAttr = wrapPrettyReplaceAttr(pretty, opts.HandlerOptions.ReplaceAttr)
		}
	}
	wl := WrapLocker(w)

	return &ZLog{
		bufPool:   bufferpool.New(bufferpool.SharedCap, bufferpool.BufCap),
		jsonPool:  newJsonHandlerPool(opts.HandlerOptions),
		withAttrs: []slog.Attr{},
		withGroup: []string{},
		opts:      opts,
		writer:    wl,
	}
}

// Clone clones the Logger, replacing the output writer with w.
//
// If w is nil, the clone shares the original Logger's output writer.
func (log *ZLog) Clone(w io.Writer) *ZLog {
	var output WriteLocker
	if w == nil {
		output = log.writer
	} else {
		output = WrapLocker(w)
	}
	withAttrs := make([]slog.Attr, len(log.withAttrs))
	copy(withAttrs, log.withAttrs)
	withGroup := make([]string, len(log.withGroup))
	copy(withGroup, log.withGroup)
	return &ZLog{
		withAttrs: withAttrs,
		withGroup: withGroup,
		opts:      log.opts,
		bufPool:   bufferpool.New(bufferpool.SharedCap, bufferpool.BufCap),
		jsonPool:  newJsonHandlerPool(log.opts.HandlerOptions),
		writer:    output,
	}
}

// Enabled implements the slog.Handler interface.
func (lo *ZLog) Enabled(ctx context.Context, lvl slog.Level) bool {
	return lvl >= lo.opts.HandlerOptions.Level.Level()
}

// Handle implements the slog.Handler interface.
func (lo *ZLog) Handle(ctx context.Context, record slog.Record) error {
	if !lo.opts.Pretty {
		j := lo.jsonPool.Get()
		defer j.Put()
		return j.Handle(ctx, record)
	}

	var levelColor *color.Color
	switch {
	case record.Level >= slog.LevelError:
		levelColor = color.New(color.FgRed)
	case record.Level >= slog.LevelWarn:
		levelColor = color.New(color.FgYellow)
	case record.Level >= slog.LevelInfo:
		levelColor = color.New(color.FgGreen)
	default:
		levelColor = color.New(color.FgWhite)
	}

	buf := lo.bufPool.Get()
	defer buf.Close()

	if record.PC != 0 && lo.opts.HandlerOptions.AddSource {
		f := frame{Caller(record.PC).Frame()}
		levelColor.Fprint(buf, f.ShortFile())
		levelColor.Fprint(buf, ":")
		levelColor.Fprint(buf, f.Line)
		levelColor.Fprint(buf, " -- ")
		levelColor.Fprint(buf, f.ShortFunction())
		buf.WriteByte('\n')
	}

	if !record.Time.IsZero() {
		const format = `[` + time.DateTime + `] `
		b := record.Time.AppendFormat(nil, format)
		buf.Write(b)
	}

	buf.WriteByte('[')
	levelColor.Add(color.Bold).Fprint(buf, record.Level.String())
	buf.WriteString("] ")

	if record.Message != "" {
		buf.WriteString(record.Message)
	}

	buf.WriteByte('\n')

	jHandler := lo.jsonPool.Get()
	defer jHandler.Put()

	_ = jHandler.Handle(ctx, record)
	if jHandler.buf.Len() > 3 { // Skip empty objects like "{}\n"
		jsonData := jHandler.buf.Bytes()
		jsonData = pretty.Pretty(jsonData)
		if lo.opts.Color {
			jsonData = pretty.Color(jsonData, nil)
		}
		buf.Write(jsonData)
	}

	buf.WriteByte('\n')

	lo.writer.Lock()
	defer lo.writer.Unlock()
	_, err := buf.WriteTo(lo.writer)
	return err
}

// WithAttrs implements the slog.Handler interface.
func (lo *ZLog) WithAttrs(attrs []slog.Attr) slog.Handler {
	l := lo.Clone(nil)
	l.withAttrs = append(l.withAttrs, attrs...)
	return l
}

// WithGroup implements the slog.Handler interface.
func (lo *ZLog) WithGroup(name string) slog.Handler {
	l := lo.Clone(nil)
	l.withGroup = append(l.withGroup, name)
	return l
}

type replaceAttrFunc = func([]string, slog.Attr) slog.Attr

func wrapPrettyReplaceAttr(pretty bool, f replaceAttrFunc) replaceAttrFunc {
	if f == nil {
		f = func(s []string, a slog.Attr) slog.Attr { return a }
	}
	return func(s []string, a slog.Attr) slog.Attr {
		if len(s) > 0 {
			return f(s, a)
		}
		if pretty {
			return f(s, a)
		}
		switch a.Key {
		case slog.TimeKey, slog.LevelKey, slog.SourceKey, slog.MessageKey:
			return slog.Attr{}
		}
		return f(s, a)
	}
}

func defaultLevel() slog.Leveler {
	if testing.Testing() {
		return slog.Level(999) // Disable logging in tests
	}
	return slog.LevelInfo
}

var defaultOpts = &ZLogOptions{
	Pretty:         isatty.IsTerminal(os.Stderr.Fd()),
	Color:          isatty.IsTerminal(os.Stderr.Fd()),
	HandlerOptions: defaultHandlerOptions,
}

var defaultHandlerOptions = &slog.HandlerOptions{
	AddSource: true,
	Level:     defaultLevel(),
	ReplaceAttr: wrapPrettyReplaceAttr(
		isatty.IsTerminal(os.Stderr.Fd()),
		nil,
	),
}