Bluemage/go/pkg/log/pretty.go

package log

import (
	"bytes"
	"context"
	"io"
	"log/slog"
	"os"
	"strings"
	"sync"
	"sync/atomic"

	"github.com/fatih/color"
	"github.com/tidwall/pretty"
	"github.com/tigorlazuardi/bluemage/go/pkg/caller"
)

type PrettyHandler struct {
	opts         *slog.HandlerOptions
	output       WriteLocker
	replaceAttr  func(groups []string, attr slog.Attr) slog.Attr
	withAttrs    []slog.Attr
	withGroup    []string
	isFlushing   atomic.Bool
	flushingChan chan struct{}
	wg           sync.WaitGroup
}

// NewPrettyHandler creates a human friendly readable logs.
func NewPrettyHandler(writer WriteLocker, opts *slog.HandlerOptions) *PrettyHandler {
	if opts == nil {
		opts = &slog.HandlerOptions{Level: slog.LevelDebug}
	}
	if opts.ReplaceAttr == nil {
		opts.ReplaceAttr = func(groups []string, attr slog.Attr) slog.Attr { return attr }
	}
	return &PrettyHandler{
		opts:   opts,
		output: writer,
		replaceAttr: func(groups []string, attr slog.Attr) slog.Attr {
			if len(groups) > 0 {
				return opts.ReplaceAttr(groups, attr)
			}
			switch attr.Key {
			case slog.TimeKey, slog.LevelKey, slog.SourceKey, slog.MessageKey:
				return slog.Attr{}
			default:
				return opts.ReplaceAttr(groups, attr)
			}
		},
		flushingChan: make(chan struct{}),
	}
}

// Flush waits until all logs are written to the underlying writer
// and then flush the buffer.
//
// Flush blocks other log writes until it's done.
func (pr *PrettyHandler) Flush() {
	pr.isFlushing.Store(true)
	defer func() { pr.isFlushing.Store(false) }()
	pr.wg.Wait()
	for {
		select {
		// signal to handlers that flushing is done.
		case pr.flushingChan <- struct{}{}:
		default:
			return
		}
	}
}

// Sync attemps to flush the buffer of the underlying writer.
func (pr *PrettyHandler) Close() error {
	pr.Flush()
	return nil
}

// Enabled implements slog.Handler interface.
func (pr *PrettyHandler) Enabled(ctx context.Context, lvl slog.Level) bool {
	return pr.opts.Level.Level() <= lvl
}

var bufferPool = sync.Pool{
	New: func() interface{} {
		buf := &bytes.Buffer{}
		buf.Grow(1024)
		return buf
	},
}

func putBuffer(buf *bytes.Buffer) {
	const limit = 1024 * 512 // 512KB
	if buf.Cap() < limit {
		buf.Reset()
		bufferPool.Put(buf)
	}
}

var jsonColorStyle = &pretty.Style{
	Key:      [2]string{"\x1B[95m", "\x1B[0m"},
	String:   [2]string{"\x1B[32m", "\x1B[0m"},
	Number:   [2]string{"\x1B[33m", "\x1B[0m"},
	True:     [2]string{"\x1B[36m", "\x1B[0m"},
	False:    [2]string{"\x1B[36m", "\x1B[0m"},
	Null:     [2]string{"\x1B[2m", "\x1B[0m"},
	Escape:   [2]string{"\x1B[35m", "\x1B[0m"},
	Brackets: [2]string{"\x1B[0m", "\x1B[0m"},
	Append:   pretty.TerminalStyle.Append,
}

var jsonPrettyOpts = &pretty.Options{
	Width:    80,
	Prefix:   "",
	Indent:   "  ",
	SortKeys: false,
}

// Handle implements slog.Handler interface.
func (pr *PrettyHandler) Handle(ctx context.Context, record slog.Record) error {
	if pr.isFlushing.Load() {
		// pr is currently flushing. Wait until the flushing process is done.
		<-pr.flushingChan
	}
	pr.wg.Add(1)
	defer pr.wg.Done()

	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 := bufferPool.Get().(*bytes.Buffer)
	jsonBuf := bufferPool.Get().(*bytes.Buffer)
	defer putBuffer(buf)
	defer putBuffer(jsonBuf)

	if record.PC != 0 && pr.opts.AddSource {
		frame := caller.From(record.PC).Frame
		cwd, _ := os.Getwd()
		f := strings.TrimPrefix(frame.File, cwd)
		f = strings.TrimPrefix(f, string(os.PathSeparator))
		levelColor.Fprint(buf, f)
		levelColor.Fprint(buf, ":")
		levelColor.Fprint(buf, frame.Line)
		levelColor.Fprint(buf, " :: ")
		split := strings.Split(frame.Function, string(os.PathSeparator))
		fnName := split[len(split)-1]
		levelColor.Fprint(buf, fnName)
		buf.WriteByte('\n')
	}

	if !record.Time.IsZero() {
		buf.WriteString(record.Time.Format("[2006-01-02 15:04:05] "))
	}

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

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

	buf.WriteByte('\n')

	serializer := pr.createSerializer(jsonBuf)
	_ = serializer.Handle(ctx, record)
	if jsonBuf.Len() > 3 { // Ignore empty json like "{}\n"
		jsonData := jsonBuf.Bytes()
		jsonData = pretty.PrettyOptions(jsonData, jsonPrettyOpts)
		jsonData = pretty.Color(jsonData, jsonColorStyle)
		buf.Write(jsonData)
	}
	buf.WriteByte('\n')

	_, err := buf.WriteTo(pr.output)
	return err
}

func (pr *PrettyHandler) createSerializer(w io.Writer) slog.Handler {
	var jsonHandler slog.Handler = slog.NewJSONHandler(w, &slog.HandlerOptions{
		Level:       slog.LevelDebug,
		ReplaceAttr: pr.replaceAttr,
	})

	if len(pr.withAttrs) > 0 {
		jsonHandler = jsonHandler.WithAttrs(pr.withAttrs)
	}

	if len(pr.withGroup) > 0 {
		for _, group := range pr.withGroup {
			jsonHandler = jsonHandler.WithGroup(group)
		}
	}

	return jsonHandler
}

func (pr *PrettyHandler) clone() *PrettyHandler {
	return &PrettyHandler{
		opts:        pr.opts,
		output:      pr.output,
		replaceAttr: pr.replaceAttr,
		withAttrs:   pr.withAttrs,
		withGroup:   pr.withGroup,
	}
}

// WithAttrs implements slog.Handler interface.
func (pr *PrettyHandler) WithAttrs(attrs []slog.Attr) slog.Handler {
	p := pr.clone()
	p.withAttrs = append(p.withAttrs, attrs...)
	return p
}

// WithGroup implements slog.Handler interface.
func (pr *PrettyHandler) WithGroup(name string) slog.Handler {
	p := pr.clone()
	p.withGroup = append(p.withGroup, name)
	return p
}