handler_json.go

package logi

import (
	"context"
	"io"
	"log/slog"
)

// JSONHandler wraps slog.JSONHandler and adds color support
type JSONHandler struct {
	handler slog.Handler
	cw      *ColorWriter
}

type ColorWriter struct {
	w      io.Writer
	enable bool
	color  string
}

func (c *ColorWriter) SetColor(color string) {
	c.color = color
}

func (c *ColorWriter) Write(p []byte) (n int, err error) {
	if !c.enable {
		c.w.Write(p)
		return
	}

	if len(p) > 0 && p[len(p)-1] == '\n' {
		p = p[:len(p)-1]
		coloredOutput := make([]byte, 0, len(c.color)+len(p)+len(colorReset)+1)
		coloredOutput = append(coloredOutput, c.color...)
		coloredOutput = append(coloredOutput, p...)
		coloredOutput = append(coloredOutput, colorReset...)
		coloredOutput = append(coloredOutput, '\n')
		p = coloredOutput
	}
	return c.w.Write(p)
}

// NewJSONHandler creates a JSONHandler that writes to w
// using the given options with optional color support.
func NewJSONHandler(w io.Writer, opts *slog.HandlerOptions, color bool) *JSONHandler {
	if opts == nil {
		opts = &slog.HandlerOptions{}
	}

	cw := &ColorWriter{
		w:      w,
		enable: color,
	}

	return &JSONHandler{
		handler: slog.NewJSONHandler(cw, opts), // will be created per Handle call
		cw:      cw,
	}
}

func (h *JSONHandler) Enabled(ctx context.Context, l slog.Level) bool {
	return h.handler.Enabled(ctx, l)
}

func (h *JSONHandler) WithAttrs(attrs []slog.Attr) slog.Handler {
	return &JSONHandler{
		handler: h.handler.WithAttrs(attrs),
		cw:      h.cw,
	}
}

func (h *JSONHandler) WithGroup(name string) slog.Handler {
	return &JSONHandler{
		handler: h.handler.WithGroup(name),
		cw:      h.cw,
	}
}

// Handle processes the record and formats the JSON output with optional colors
func (h *JSONHandler) Handle(ctx context.Context, r slog.Record) error {
	h.cw.SetColor(getLevelColor(r.Level))
	return h.handler.Handle(ctx, r)
}