logging/logging.go

package logging

import (
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"sync"
	"time"

	"golang.org/x/term"
)

const (
	COLOR_RESET  = "\033[0m"
	COLOR_RED    = "\033[91m"
	COLOR_YELLOW = "\033[93m"
	COLOR_GREEN  = "\033[92m"
	COLOR_BLUE   = "\033[94m"
	COLOR_CYAN   = "\033[96m"
	COLOR_WHITE  = "\033[97m"
	COLOR_GRAY   = "\033[90m"
	COLOR_GREY   = COLOR_GRAY
)

type LogLevel int

const (
	Debug LogLevel = iota
	Info
	Warn
	Error
	Critical
)

func (l LogLevel) String() string {
	switch l {
	case Debug:
		return "DEBUG"
	case Info:
		return "INFO"
	case Warn:
		return "WARN"
	case Error:
		return "ERROR"
	case Critical:
		return "CRITICAL"
	default:
		return "UNKNOWN"
	}
}

var LevelColorMap = map[LogLevel]string{
	Debug:    COLOR_GRAY,
	Info:     COLOR_WHITE,
	Warn:     COLOR_YELLOW,
	Error:    COLOR_RED,
	Critical: COLOR_RED,
}

type Formatter interface {
	Format(level LogLevel, timestamp time.Time, message string) string
}

type ConsoleFormatter struct {
	TimestampFormatter string
	UseColor           bool
	LevelColors        map[LogLevel]string
	mu                 sync.Mutex
}

func NewConsoleFormatter(timestampFormat string, useColor bool, levelColors map[LogLevel]string) *ConsoleFormatter {
	if timestampFormat == "" {
		timestampFormat = "[02/01/06 15:04:05]"
	}
	if levelColors == nil {
		levelColors = LevelColorMap
	}

	return &ConsoleFormatter{
		TimestampFormatter: timestampFormat,
		UseColor:           useColor,
		LevelColors:        levelColors,
	}
}

func (f *ConsoleFormatter) Format(level LogLevel, timestamp time.Time, message string) string {
	f.mu.Lock()
	defer f.mu.Unlock()

	timestampStr := timestamp.Format(f.TimestampFormatter)
	levelStr := level.String()
	color, ok := f.LevelColors[level]
	if !ok {
		color = COLOR_WHITE
	}

	if f.UseColor && isTTY() {
		return fmt.Sprintf("%s %s[%s]%s %s", timestampStr, color, levelStr, COLOR_RESET, message)
	}
	return fmt.Sprintf("%s [%s] %s", timestampStr, levelStr, message)
}

type FileFormatter struct {
	TimestampFormat string
	LoggerName      string
}

func NewFileFormatter(timestampFormat, loggerName string) *FileFormatter {
	if timestampFormat == "" {
		timestampFormat = "2006-01-02 15:04:05"
	}
	return &FileFormatter{
		TimestampFormat: timestampFormat,
		LoggerName:      loggerName,
	}
}

func (f *FileFormatter) Format(level LogLevel, timestamp time.Time, message string) string {
	return fmt.Sprintf("[ %s | %s | %s ] %s", timestamp.Format(f.TimestampFormat), f.LoggerName, level.String(), message)
}

type Logger struct {
	name               string
	level              LogLevel
	consoleColor       bool
	logDir             string
	logFilenamePattern string
	formatter          Formatter
	consoleWriter      io.Writer
	fileWriter         io.WriteCloser
	fileLogger         *log.Logger
	fileFormatter      *FileFormatter
	mu                 sync.Mutex
	initialized        bool
}

var (
	globalLogger *Logger
	once         sync.Once
)

func NewLogger(
	loggerName string,
	logLevel LogLevel,
	consoleColor bool,
	logDir string,
	logFilenamePattern string,
) *Logger {
	if loggerName == "" {
		loggerName = "main"
	}
	if logDir == "" {
		logDir = "logs"
	}
	if logFilenamePattern == "" {
		logFilenamePattern = "2006-01-02"
	}

	return &Logger{
		name:               loggerName,
		level:              logLevel,
		consoleColor:       consoleColor,
		logDir:             logDir,
		logFilenamePattern: logFilenamePattern,
		consoleWriter:      os.Stdout,
	}
}

func (l *Logger) init() {
	l.mu.Lock()
	defer l.mu.Unlock()

	if l.initialized {
		return
	}

	if err := os.MkdirAll(l.logDir, 0755); err != nil {
		fmt.Fprintf(os.Stderr, "Error creating log directory %s: %v\n", l.logDir, err)
	} else {
		logFilePath := l.genUniqueLogFilename()
		file, err := os.OpenFile(logFilePath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error opening log file %s: %v\n", logFilePath, err)
		} else {
			l.fileWriter = file
			if l.fileWriter != nil {
				l.fileFormatter = NewFileFormatter("", l.name)
			}
			l.fileLogger = log.New(l.fileWriter, "", 0)
			fmt.Printf("Log file will be saved to %s\n", logFilePath)
		}
	}

	l.formatter = NewConsoleFormatter("", l.consoleColor, nil)

	l.initialized = true
}

func (l *Logger) genUniqueLogFilename() string {
	now := time.Now()
	baseFilename := now.Format(l.logFilenamePattern)
	logFilePath := filepath.Join(l.logDir, fmt.Sprintf("%s.log", baseFilename))

	if _, err := os.Stat(logFilePath); os.IsNotExist(err) {
		return logFilePath
	}

	counter := 1
	for {
		suffixedFilePath := filepath.Join(l.logDir, fmt.Sprintf("%s_%d.log", baseFilename, counter))
		if _, err := os.Stat(suffixedFilePath); os.IsNotExist(err) {
			return suffixedFilePath
		}
		counter += 1
	}
}

func (l *Logger) Log(level LogLevel, message string) {
	if !l.initialized {
		l.init()
	}

	if level < l.level {
		return
	}

	now := time.Now()

	if l.formatter != nil {
		formattedMsg := l.formatter.Format(level, now, message)
		fmt.Fprintln(l.consoleWriter, formattedMsg)
	} else {
		fmt.Fprintf(l.consoleWriter, "%s [%s] %s\n", now.Format("[15:04:05]"), level.String(), message)
	}

	if l.fileLogger != nil && l.fileFormatter != nil {
		formattedFileMsg := l.fileFormatter.Format(level, now, message)
		l.fileLogger.Println(formattedFileMsg)
	}
}

func (l *Logger) Debug(message string, v ...any) {
	l.Log(Debug, fmt.Sprintf(message, v...))
}

func (l *Logger) Warn(message string, v ...any) {
	l.Log(Warn, fmt.Sprintf(message, v...))
}

func (l *Logger) Info(message string, v ...any) {
	l.Log(Info, fmt.Sprintf(message, v...))
}

func (l *Logger) Error(message string, v ...any) {
	l.Log(Error, fmt.Sprintf(message, v...))
}

func (l *Logger) Critical(message string, v ...any) {
	l.Log(Critical, fmt.Sprintf(message, v...))
	os.Exit(1)
}

func (l *Logger) Close() error {
	l.mu.Lock()
	defer l.mu.Unlock()
	if l.fileWriter != nil {
		err := l.fileWriter.Close()
		l.fileWriter = nil
		l.fileLogger = nil
		return err
	}
	return nil
}

func GetLogger() *Logger {
	once.Do(func() {
		globalLogger = NewLogger("shizuku", Debug, true, "logs", "2006-01-02")
	})
	return globalLogger
}

func isTTY() bool {
	return os.Getenv("TERM") != "" && term.IsTerminal(int(os.Stdout.Fd()))
}