custom_handler.go

package logger

import (
	"context"
	"fmt"
	"io"
	"log/slog"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

const (
	// Placeholders
	PlaceholderTime    = "{time}"
	PlaceholderLevel   = "{level}"
	PlaceholderMessage = "{message}"
	PlaceholderFile    = "{file}"
	PlaceholderAttrs   = "{attrs}"

	// ANSI escape codes
	ansiReset          = "\033[0m"
	ansiFaint          = "\033[2m"
	ansiBrightCyan     = "\033[96m"
	ansiBrightRed      = "\033[91m"
	ansiBrightRedFaint = "\033[91;2m"
	ansiBrightGreen    = "\033[92m"
	ansiBrightYellow   = "\033[93m"
	ansiBrightMagenta  = "\033[95m"
)

// TokenType represents the type of a template token
type TokenType int

const (
	TokenTypeText TokenType = iota
	TokenTypeTime
	TokenTypeLevel
	TokenTypeMessage
	TokenTypeFile
	TokenTypeAttrs
)

// Token represents a parsed template component
type Token struct {
	Type TokenType
	Text string // For static text tokens
}

// ParsedTemplate holds the pre-parsed template tokens
type ParsedTemplate struct {
	tokens []Token
}

type groupOrAttrs struct {
	group string
	attrs []slog.Attr
	next  *groupOrAttrs
}

// handlerConfig stores immutable configuration data for atomic access
type handlerConfig struct {
	globalCfg      *Config
	outputCfg      outputConfig
	attrsIndex     int
	goa            *groupOrAttrs
	groupPrefix    string
	opts           slog.HandlerOptions
	parsedTemplate *ParsedTemplate // Pre-parsed template for efficient formatting
}

type customHandler struct {
	// Lightweight mutex to protect write operations
	writeMu *sync.Mutex
	out     io.Writer

	// Configuration data, accessed using atomic operations
	config atomic.Value // *handlerConfig

	// Byte buffer pool, thread-safe
	pool *sync.Pool

	hotBuf atomic.Pointer[[]byte]
}

// outputConfig interface for unified access to Console and File configurations
type outputConfig interface {
	GetFormat() OutputFormat
	GetColor() bool
	GetFormatter() string
}

// ConsoleConfig implements outputConfig interface
func (c *ConsoleConfig) GetFormat() OutputFormat {
	return c.Format
}

func (c *ConsoleConfig) GetColor() bool {
	return c.Color
}

func (c *ConsoleConfig) GetFormatter() string {
	return c.Formatter
}

// FileConfig implements outputConfig interface
func (c *FileConfig) GetFormat() OutputFormat {
	return c.Format
}

func (c *FileConfig) GetColor() bool {
	// File output does not support color
	return false
}

func (c *FileConfig) GetFormatter() string {
	return c.Formatter
}

// parseTemplate parses a format template into tokens for efficient rendering
func parseTemplate(template string) *ParsedTemplate {
	if template == "" {
		template = DefaultFormatter
	}

	var tokens []Token
	remaining := template

	for len(remaining) > 0 {
		// Find the next placeholder
		nextPlaceholder := -1
		var placeholderType TokenType
		var placeholderLen int

		// Check for each placeholder type
		placeholders := []struct {
			text      string
			tokenType TokenType
		}{
			{PlaceholderTime, TokenTypeTime},
			{PlaceholderLevel, TokenTypeLevel},
			{PlaceholderMessage, TokenTypeMessage},
			{PlaceholderFile, TokenTypeFile},
			{PlaceholderAttrs, TokenTypeAttrs},
		}

		for _, p := range placeholders {
			if idx := strings.Index(remaining, p.text); idx != -1 {
				if nextPlaceholder == -1 || idx < nextPlaceholder {
					nextPlaceholder = idx
					placeholderType = p.tokenType
					placeholderLen = len(p.text)
				}
			}
		}

		if nextPlaceholder == -1 {
			// No more placeholders, add remaining text as static token
			if len(remaining) > 0 {
				tokens = append(tokens, Token{Type: TokenTypeText, Text: remaining})
			}
			break
		}

		// Add static text before placeholder (if any)
		if nextPlaceholder > 0 {
			tokens = append(tokens, Token{Type: TokenTypeText, Text: remaining[:nextPlaceholder]})
		}

		// Add placeholder token
		tokens = append(tokens, Token{Type: placeholderType})

		// Move to after the placeholder
		remaining = remaining[nextPlaceholder+placeholderLen:]
	}

	return &ParsedTemplate{tokens: tokens}
}

func newCustomHandler(w io.Writer, globalCfg *Config, outputCfg outputConfig, opts *slog.HandlerOptions) (slog.Handler, error) {
	formatter := outputCfg.GetFormatter()
	if formatter == "" {
		formatter = DefaultFormatter
	}

	// Parse template at startup time for efficient formatting
	parsedTemplate := parseTemplate(formatter)

	// Create configuration object
	cfg := &handlerConfig{
		globalCfg:      globalCfg,
		outputCfg:      outputCfg,
		attrsIndex:     strings.Index(formatter, PlaceholderAttrs),
		goa:            nil,
		groupPrefix:    "",
		parsedTemplate: parsedTemplate,
	}

	if opts != nil {
		cfg.opts = *opts
	} else {
		cfg.opts = slog.HandlerOptions{
			Level:       globalCfg.Level,
			AddSource:   globalCfg.AddSource,
			ReplaceAttr: globalCfg.ReplaceAttr,
		}
	}

	h := &customHandler{
		writeMu: &sync.Mutex{},
		out:     w,
		pool: &sync.Pool{
			New: func() any {
				buf := make([]byte, 0, 1024)
				return &buf
			},
		},
	}

	// Atomically set the configuration
	h.config.Store(cfg)

	return h, nil
}

// getConfig atomically retrieves the current configuration
func (h *customHandler) getConfig() *handlerConfig {
	cfg, ok := h.config.Load().(*handlerConfig)
	if !ok || cfg == nil {
		// Should never happen — defensive fallback instead of crash
		return &handlerConfig{
			globalCfg:      DefaultConfig(),
			parsedTemplate: parseTemplate(""),
			opts:           slog.HandlerOptions{Level: slog.LevelInfo},
		}
	}
	return cfg
}

func collectGroupOrAttrs(goa *groupOrAttrs) []groupOrAttrs {
	var nodes []groupOrAttrs
	for curr := goa; curr != nil; curr = curr.next {
		nodes = append(nodes, *curr)
	}

	for i, j := 0, len(nodes)-1; i < j; i, j = i+1, j-1 {
		nodes[i], nodes[j] = nodes[j], nodes[i]
	}

	return nodes
}

func (h *customHandler) Enabled(_ context.Context, level slog.Level) bool {
	cfg := h.getConfig()
	minLevel := slog.LevelInfo
	if cfg.opts.Level != nil {
		minLevel = cfg.opts.Level.Level()
	}
	return level >= minLevel
}

func (h *customHandler) Handle(ctx context.Context, r slog.Record) error {
	// Lock-free access to config and formatting
	cfg := h.getConfig()

	// Lock-free log formatting (CPU-intensive operation)
	buf := h.getBuffer()
	defer h.putBuffer(buf)

	h.formatLogLine(buf, r, cfg)

	if err := h.lockedWrite(*buf); err != nil {
		return fmt.Errorf("custom handler write: %w", err)
	}
	return nil
}

// lockedWrite serialises writes to the underlying writer.
// Using defer ensures the mutex is released even if Write panics.
func (h *customHandler) lockedWrite(buf []byte) error {
	h.writeMu.Lock()
	defer h.writeMu.Unlock()
	_, err := h.out.Write(buf)
	return err
}

func (h *customHandler) getBuffer() *[]byte {
	if buf := h.hotBuf.Swap(nil); buf != nil {
		return buf
	}
	return h.pool.Get().(*[]byte)
}

const maxPooledBufCap = 16 << 10 // 16 KiB — discard oversized buffers to limit pool memory

func (h *customHandler) putBuffer(buf *[]byte) {
	if cap(*buf) > maxPooledBufCap {
		return
	}

	*buf = (*buf)[:0]
	if h.hotBuf.CompareAndSwap(nil, buf) {
		return
	}
	h.pool.Put(buf)
}

func (h *customHandler) WithAttrs(attrs []slog.Attr) slog.Handler {
	if len(attrs) == 0 {
		return h
	}

	attrsCopy := append([]slog.Attr(nil), attrs...)

	// Lock-free operation: copy config and add new attributes
	oldCfg := h.getConfig()
	newCfg := &handlerConfig{
		globalCfg:      oldCfg.globalCfg,
		outputCfg:      oldCfg.outputCfg,
		attrsIndex:     oldCfg.attrsIndex,
		goa:            &groupOrAttrs{attrs: attrsCopy, next: oldCfg.goa},
		groupPrefix:    oldCfg.groupPrefix,
		opts:           oldCfg.opts,
		parsedTemplate: oldCfg.parsedTemplate, // Share the parsed template
	}

	newHandler := &customHandler{
		writeMu: h.writeMu,
		out:     h.out,
		pool:    h.pool,
	}
	newHandler.config.Store(newCfg)

	return newHandler
}

func (h *customHandler) WithGroup(name string) slog.Handler {
	if name == "" {
		return h
	}

	// Lock-free operation: copy config and add new group
	oldCfg := h.getConfig()
	newCfg := &handlerConfig{
		globalCfg:      oldCfg.globalCfg,
		outputCfg:      oldCfg.outputCfg,
		attrsIndex:     oldCfg.attrsIndex,
		goa:            &groupOrAttrs{group: name, next: oldCfg.goa},
		groupPrefix:    oldCfg.groupPrefix + name + ".",
		opts:           oldCfg.opts,
		parsedTemplate: oldCfg.parsedTemplate, // Share the parsed template
	}

	newHandler := &customHandler{
		writeMu: h.writeMu,
		out:     h.out,
		pool:    h.pool,
	}
	newHandler.config.Store(newCfg)

	return newHandler
}

func (h *customHandler) formatLogLine(buf *[]byte, r slog.Record, cfg *handlerConfig) {
	// Process built-in attributes through ReplaceAttr like standard slog handlers
	rep := cfg.opts.ReplaceAttr

	// Pre-compute all the parts that might be needed
	var timePart, levelPart, msgPart, filePart, attrsPart []byte

	// Handle time (built-in attribute)
	if !r.Time.IsZero() {
		timeAttr := slog.Time(slog.TimeKey, r.Time.In(cfg.globalCfg.TimeZone))
		if rep != nil {
			timeAttr = rep(nil, timeAttr) // Built-ins are not in any group
		}
		if !timeAttr.Equal(slog.Attr{}) { // Check if not removed by ReplaceAttr
			timeValue := timeAttr.Value.Any()
			if t, ok := timeValue.(time.Time); ok {
				timePart = h.appendColorized(timePart, t.Format(cfg.globalCfg.TimeFormat), ansiFaint, cfg)
			} else {
				// ReplaceAttr changed the type, use the new value
				timePart = h.appendColorized(timePart, fmt.Sprintf("%v", timeValue), ansiFaint, cfg)
			}
		}
	}

	// Handle level (built-in attribute)
	levelAttr := slog.Any(slog.LevelKey, r.Level)
	if rep != nil {
		levelAttr = rep(nil, levelAttr) // Built-ins are not in any group
	}
	if !levelAttr.Equal(slog.Attr{}) { // Check if not removed by ReplaceAttr
		levelValue := levelAttr.Value.Any()
		if level, ok := levelValue.(slog.Level); ok {
			levelPart = h.appendColorizedLevel(levelPart, level, cfg)
		} else {
			// ReplaceAttr changed the type, use the new value
			levelPart = h.appendColorized(levelPart, fmt.Sprintf("%v", levelValue), ansiBrightGreen, cfg)
		}
	}

	// Handle message (built-in attribute)
	msgAttr := slog.String(slog.MessageKey, r.Message)
	if rep != nil {
		msgAttr = rep(nil, msgAttr) // Built-ins are not in any group
	}
	if !msgAttr.Equal(slog.Attr{}) { // Check if not removed by ReplaceAttr
		msgPart = h.appendColorizedMessage(msgPart, msgAttr.Value.String(), r.Level, cfg)
	}

	// Handle source/file (built-in attribute)
	if cfg.opts.AddSource {
		// Create source attribute like standard slog handlers
		source := r.Source()
		sourceAttr := slog.Any(slog.SourceKey, source)
		if rep != nil {
			sourceAttr = rep(nil, sourceAttr) // Built-ins are not in any group
		}
		if !sourceAttr.Equal(slog.Attr{}) { // Check if not removed by ReplaceAttr
			sourceValue := sourceAttr.Value.Any()
			if src, ok := sourceValue.(*slog.Source); ok && src != nil {
				if src.File != "" {
					// Standard format: filename:function:line
					filePart = h.appendColorized(filePart, fmt.Sprintf("%s:%s:%d", filepath.Base(src.File), filepath.Base(src.Function), src.Line), ansiFaint, cfg)
				}
			} else {
				// ReplaceAttr changed the type, use the new value
				filePart = h.appendColorized(filePart, fmt.Sprintf("%v", sourceValue), ansiFaint, cfg)
			}
		}
	}

	// Handle user attributes
	if cfg.attrsIndex >= 0 {
		attrBuilder := h.getBuffer()
		defer h.putBuffer(attrBuilder)

		isFirst := true
		preStoredPrefix := ""
		groups := make([]string, 0)

		for _, node := range collectGroupOrAttrs(cfg.goa) {
			if node.group != "" {
				groups = append(groups, node.group)
				preStoredPrefix += node.group + "."
				continue
			}

			for _, a := range node.attrs {
				if h.appendColorizedAttr(attrBuilder, a, preStoredPrefix, r.Level, isFirst, cfg, rep, groups) {
					isFirst = false
				}
			}
		}

		r.Attrs(func(a slog.Attr) bool {
			if h.appendColorizedAttr(attrBuilder, a, cfg.groupPrefix, r.Level, isFirst, cfg, rep, groups) {
				isFirst = false
			}
			return true
		})
		attrsPart = *attrBuilder
	}

	// Use parsed template for efficient formatting
	h.renderTemplate(buf, cfg.parsedTemplate, timePart, levelPart, msgPart, filePart, attrsPart)
	*buf = append(*buf, '\n')
}

// renderTemplate efficiently renders the parsed template by iterating through tokens
func (h *customHandler) renderTemplate(buf *[]byte, template *ParsedTemplate, timePart, levelPart, msgPart, filePart, attrsPart []byte) {
	tokens := template.tokens
	for i, token := range tokens {
		switch token.Type {
		case TokenTypeText:
			// Handle text tokens, but be smart about spaces around empty placeholders
			text := token.Text

			// If this is a space before an empty placeholder, skip it when appropriate
			if text == " " && i+1 < len(tokens) {
				nextToken := tokens[i+1]

				// Check if next token is an empty placeholder
				isEmpty := false
				switch nextToken.Type {
				case TokenTypeTime:
					isEmpty = len(timePart) == 0
				case TokenTypeLevel:
					isEmpty = len(levelPart) == 0
				case TokenTypeMessage:
					isEmpty = len(msgPart) == 0
				case TokenTypeFile:
					isEmpty = len(filePart) == 0
				case TokenTypeAttrs:
					isEmpty = len(attrsPart) == 0
				}

				if isEmpty {
					// Skip trailing space before empty placeholder at end of template
					if i+2 >= len(tokens) {
						continue
					}
					// Skip space before empty placeholder followed by another space
					if tokens[i+2].Type == TokenTypeText && strings.HasPrefix(tokens[i+2].Text, " ") {
						continue
					}
				}
			}

			*buf = append(*buf, text...)
		case TokenTypeTime:
			if len(timePart) > 0 {
				*buf = append(*buf, timePart...)
			}
		case TokenTypeLevel:
			if len(levelPart) > 0 {
				*buf = append(*buf, levelPart...)
			}
		case TokenTypeMessage:
			if len(msgPart) > 0 {
				*buf = append(*buf, msgPart...)
			}
		case TokenTypeFile:
			if len(filePart) > 0 {
				*buf = append(*buf, filePart...)
			}
		case TokenTypeAttrs:
			if len(attrsPart) > 0 {
				*buf = append(*buf, attrsPart...)
			}
		}
	}
}

func (h *customHandler) appendColorized(dst []byte, s, color string, cfg *handlerConfig) []byte {
	if !cfg.outputCfg.GetColor() {
		return append(dst, s...)
	}
	dst = append(dst, color...)
	dst = append(dst, s...)
	dst = append(dst, ansiReset...)
	return dst
}

func (h *customHandler) appendColorizedLevel(dst []byte, level slog.Level, cfg *handlerConfig) []byte {
	var color string
	switch {
	case level <= slog.LevelDebug:
		color = ansiBrightCyan
	case level <= slog.LevelInfo:
		color = ansiBrightGreen
	case level <= slog.LevelWarn:
		color = ansiBrightYellow
	case level <= slog.LevelError:
		color = ansiBrightRed
	default:
		color = ansiBrightMagenta
	}

	if cfg.outputCfg.GetColor() {
		dst = append(dst, color...)
	}

	if appended, err := level.AppendText(dst); err == nil {
		dst = appended
	} else {
		dst = append(dst, level.String()...)
	}

	if cfg.outputCfg.GetColor() {
		dst = append(dst, ansiReset...)
	}

	return dst
}

func (h *customHandler) appendColorizedMessage(dst []byte, msg string, level slog.Level, cfg *handlerConfig) []byte {
	if level >= slog.LevelError {
		return h.appendColorized(dst, msg, ansiBrightRed, cfg)
	}
	return append(dst, msg...)
}

func (h *customHandler) appendColorizedToBuffer(buf *[]byte, s, color string, cfg *handlerConfig) {
	*buf = h.appendColorized(*buf, s, color, cfg)
}

func (h *customHandler) appendColorizedAttr(buf *[]byte, a slog.Attr, prefix string, level slog.Level, isFirst bool, cfg *handlerConfig, rep func([]string, slog.Attr) slog.Attr, groups []string) bool {
	appendValue := func(dst *[]byte, v slog.Value) {
		switch v.Kind() {
		case slog.KindString:
			*dst = append(*dst, v.String()...)
		case slog.KindInt64:
			*dst = strconv.AppendInt(*dst, v.Int64(), 10)
		case slog.KindUint64:
			*dst = strconv.AppendUint(*dst, v.Uint64(), 10)
		case slog.KindFloat64:
			*dst = strconv.AppendFloat(*dst, v.Float64(), 'g', -1, 64)
		case slog.KindBool:
			*dst = strconv.AppendBool(*dst, v.Bool())
		case slog.KindTime:
			*dst = v.Time().AppendFormat(*dst, time.RFC3339Nano)
		case slog.KindDuration:
			*dst = append(*dst, v.Duration().String()...)
		case slog.KindAny:
			*dst = fmt.Appendf(*dst, "%v", v.Any())
		default:
			*dst = fmt.Appendf(*dst, "%v", v.Any())
		}
	}

	a.Value = a.Value.Resolve()

	if a.Equal(slog.Attr{}) {
		return false
	}

	if a.Value.Kind() != slog.KindGroup && rep != nil {
		a = rep(groups, a)
		a.Value = a.Value.Resolve()

		if a.Equal(slog.Attr{}) {
			return false
		}
	}

	if a.Value.Kind() == slog.KindGroup {
		attrs := a.Value.Group()
		if len(attrs) == 0 {
			return false
		}

		childPrefix := prefix
		childGroups := groups
		if a.Key != "" {
			childPrefix += a.Key + "."
			childGroups = append(append([]string(nil), groups...), a.Key)
		}

		wrote := false
		for _, child := range attrs {
			if h.appendColorizedAttr(buf, child, childPrefix, level, isFirst && !wrote, cfg, rep, childGroups) {
				wrote = true
			}
		}

		return wrote
	}

	if !isFirst {
		*buf = append(*buf, ' ')
	}

	// Build the key with group prefixes (slog standard behavior)
	key := a.Key
	if prefix != "" {
		key = prefix + a.Key
	}

	if level >= slog.LevelError && a.Key == "error" {
		h.appendColorizedToBuffer(buf, key, ansiBrightRedFaint, cfg)
		h.appendColorizedToBuffer(buf, "=", ansiBrightRedFaint, cfg)
		start := len(*buf)
		appendValue(buf, a.Value)
		if cfg.outputCfg.GetColor() && len(*buf) > start {
			value := append([]byte(nil), (*buf)[start:]...)
			*buf = (*buf)[:start]
			*buf = append(*buf, ansiBrightRed...)
			*buf = append(*buf, value...)
			*buf = append(*buf, ansiReset...)
		}
	} else {
		h.appendColorizedToBuffer(buf, key, ansiFaint, cfg)
		h.appendColorizedToBuffer(buf, "=", ansiFaint, cfg)
		appendValue(buf, a.Value)
	}

	return true
}