// ConsoleHandler formats slog.Logger output in console format, a bit similar with Uber's zap ConsoleEncoder
// The log format is designed to be human-readable. 
// 
// Performance can definitely be improved, however it's not in my priority as 
// this should only be used in development environment.
//
// e.g. log output:
// 2022-11-24T11:40:20+08:00 DEBUG ./main.go:162 Debug message {"hello":"world","!BADKEY":"bad kv"}
// 2022-11-24T11:40:20+08:00 INFO ./main.go:167 Info message {"with_key_1":"with_value_1","group_1":{"with_key_2":"with_value_2","hello":"world"}}
// 2022-11-24T11:40:20+08:00 WARN ./main.go:168 Warn message {"with_key_1":"with_value_1","group_1":{"with_key_2":"with_value_2","hello":"world"}}
// 2022-11-24T11:40:20+08:00 ERROR ./main.go:169 Error message {"with_key_1":"with_value_1","group_1":{"with_key_2":"with_value_2","hello":"world","err":"an error"}}

package main

import (
	"bytes"
	"errors"
	"fmt"
	"golang.org/x/exp/slog"
	"io"
	"os"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"
)

type ConsoleHandler struct {
	opts            ConsoleHandlerOptions
	internalHandler slog.Handler

	mu sync.Mutex
	w  io.Writer
}

type ConsoleHandlerOptions struct {
	SlogOpts slog.HandlerOptions
	UseColor bool
}

func NewConsoleHandler(w io.Writer) *ConsoleHandler {
	return ConsoleHandlerOptions{}.NewConsoleHandler(w)
}

func (opts ConsoleHandlerOptions) NewConsoleHandler(w io.Writer) *ConsoleHandler {
	internalOpts := opts.SlogOpts
	internalOpts.AddSource = false
	internalOpts.ReplaceAttr = func(a slog.Attr) slog.Attr {
		if a.Key == "time" || a.Key == "level" || a.Key == "msg" {
			return slog.String("", "")
		}
		rep := opts.SlogOpts.ReplaceAttr
		if rep != nil {
			return rep(a)
		}
		return a
	}
	return &ConsoleHandler{opts: opts, w: w, internalHandler: internalOpts.NewJSONHandler(w)}
}

func (h *ConsoleHandler) Enabled(level slog.Level) bool {
	return h.internalHandler.Enabled(level)
}

func (h *ConsoleHandler) Handle(r slog.Record) error {
	var buf bytes.Buffer

	buf.WriteString(r.Time.Format(time.RFC3339))
	buf.WriteString(" ")

	level := r.Level.String()
	if h.opts.UseColor {
		level = addColorToLevel(level)
	}
	buf.WriteString(level)
	buf.WriteString(" ")

	if h.opts.SlogOpts.AddSource {
		file, line := r.SourceLine()
		if file != "" {
			buf.WriteString(trimRootPath(file))
			buf.WriteString(":")
			buf.Write([]byte(strconv.Itoa(line)))
			buf.WriteString(" ")
		}
	}

	buf.WriteString(r.Message)
	buf.WriteString(" ")

	h.mu.Lock()
	defer h.mu.Unlock()
	_, err := h.w.Write(buf.Bytes())
	if err != nil {
		return err
	}

	return h.internalHandler.Handle(r)
}

func (h *ConsoleHandler) WithAttrs(attrs []slog.Attr) slog.Handler {
	return &ConsoleHandler{
		opts:            h.opts,
		w:               h.w,
		internalHandler: h.internalHandler.WithAttrs(attrs),
	}
}

func (h *ConsoleHandler) WithGroup(name string) slog.Handler {
	return &ConsoleHandler{
		opts:            h.opts,
		w:               h.w,
		internalHandler: h.internalHandler.WithGroup(name),
	}
}

var (
	_, callerFile, _, _ = runtime.Caller(0)
	rootPath            = filepath.Dir(callerFile)
)

func trimRootPath(p string) string {
	return strings.Replace(p, rootPath, ".", 1)
}

type Color uint8

const (
	Black Color = iota + 30
	Red
	Green
	Yellow
	Blue
	Magenta
	Cyan
	White
)

// Add adds the coloring to the given string.
func (c Color) Add(s string) string {
	return fmt.Sprintf("\x1b[%dm%s\x1b[0m", uint8(c), s)
}

var (
	levelToColor = map[string]Color{
		slog.DebugLevel.String(): Magenta,
		slog.InfoLevel.String():  Blue,
		slog.WarnLevel.String():  Yellow,
		slog.ErrorLevel.String(): Red,
	}
	unknownLevelColor = Red
)

func addColorToLevel(level string) string {
	color, ok := levelToColor[level]
	if !ok {
		color = unknownLevelColor
	}
	return color.Add(level)
}

func main() {
	logHandler := ConsoleHandlerOptions{
		SlogOpts: slog.HandlerOptions{
			AddSource: true,
			Level:     slog.DebugLevel,
		},
		UseColor: true,
	}.NewConsoleHandler(os.Stderr)

	logger := slog.New(logHandler)
	logger.Debug("Debug message", "hello", "world", "bad kv")
	logger = logger.
		With("with_key_1", "with_value_1").
		WithGroup("group_1").
		With("with_key_2", "with_value_2")
	logger.Info("Info message", "hello", "world")
	logger.Warn("Warn message", "hello", "world")
	logger.Error("Error message", errors.New("an error"), "hello", "world")
}