TaskPool/internal/logger/logger.go

package logger

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/engigu/taskpool/internal/systime"
	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
)

var Log *zap.Logger
var Sugar *zap.SugaredLogger
var atomicLevel zap.AtomicLevel

// ANSI 颜色代码
const (
	colorReset  = "\033[0m"
	colorRed    = "\033[31m" // error, fatal, panic
	colorYellow = "\033[33m" // warn
	colorBlue   = "\033[36m" // info
	colorGray   = "\033[37m" // debug
)

// customCore 实现 zapcore.Core 以提供与 logrus 一模一样的格式
type customCore struct {
	level  zapcore.LevelEnabler
	writer zapcore.WriteSyncer
}

func (c *customCore) Enabled(l zapcore.Level) bool {
	return c.level.Enabled(l)
}

func (c *customCore) With(fields []zapcore.Field) zapcore.Core {
	// 目前忽略字段，以保持与旧 logrus 格式一模一样（旧格式只输出 entry.Message）
	return c
}

func (c *customCore) Check(ent zapcore.Entry, ce *zapcore.CheckedEntry) *zapcore.CheckedEntry {
	if c.Enabled(ent.Level) {
		return ce.AddCore(ent, c)
	}
	return ce
}

func (c *customCore) Write(ent zapcore.Entry, fields []zapcore.Field) error {
	// 统一使用东八区时间
	timestamp := systime.InCST(ent.Time).Format("2006-01-02 15:04:05")
	level := strings.ToUpper(ent.Level.String())

	var levelColor string
	switch ent.Level {
	case zapcore.DebugLevel:
		levelColor = colorGray
	case zapcore.InfoLevel:
		levelColor = colorBlue
	case zapcore.WarnLevel:
		levelColor = colorYellow
	case zapcore.ErrorLevel, zapcore.DPanicLevel, zapcore.PanicLevel, zapcore.FatalLevel:
		levelColor = colorRed
	default:
		levelColor = colorBlue
	}

	msg := fmt.Sprintf("[%s]%s[%s]%s %s\n", timestamp, levelColor, level, colorReset, ent.Message)
	_, err := c.writer.Write([]byte(msg))
	return err
}

func (c *customCore) Sync() error {
	return c.writer.Sync()
}

func newLogger(output zapcore.WriteSyncer) *zap.Logger {
	core := &customCore{
		level:  atomicLevel,
		writer: output,
	}
	return zap.New(core)
}

func init() {
	// 强制设置全局时区为东八区
	time.Local = systime.CST
	atomicLevel = zap.NewAtomicLevelAt(zap.InfoLevel)
	Log = newLogger(zapcore.AddSync(os.Stdout))
	Sugar = Log.Sugar()
}

// SetupFileOutput 设置文件输出
func SetupFileOutput(logDir string) error {
	if err := os.MkdirAll(logDir, 0755); err != nil {
		return err
	}

	logFile := filepath.Join(logDir, systime.FormatDate(time.Now())+".log")
	file, err := os.OpenFile(logFile, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666)
	if err != nil {
		return err
	}

	Log = newLogger(zapcore.AddSync(file))
	Sugar = Log.Sugar()
	return nil
}

// SetOutput 直接设置 Log 实例
func SetOutput(l *zap.Logger) {
	Log = l
	Sugar = l.Sugar()
}

// SetSugar 直接设置 Sugar 实例
func SetSugar(s *zap.SugaredLogger) {
	Sugar = s
}

// SetLevel 设置日志级别
func SetLevel(level string) {
	switch level {
	case "debug":
		atomicLevel.SetLevel(zap.DebugLevel)
	case "info":
		atomicLevel.SetLevel(zap.InfoLevel)
	case "warn":
		atomicLevel.SetLevel(zap.WarnLevel)
	case "error":
		atomicLevel.SetLevel(zap.ErrorLevel)
	default:
		atomicLevel.SetLevel(zap.InfoLevel)
	}
}

// 便捷方法
func Debug(args ...interface{}) { Sugar.Debug(args...) }
func Info(args ...interface{})  { Sugar.Info(args...) }
func Warn(args ...interface{})  { Sugar.Warn(args...) }
func Error(args ...interface{}) { Sugar.Error(args...) }
func Fatal(args ...interface{}) { Sugar.Fatal(args...) }

func Debugf(format string, args ...interface{}) { Sugar.Debugf(format, args...) }
func Infof(format string, args ...interface{})  { Sugar.Infof(format, args...) }
func Warnf(format string, args ...interface{})  { Sugar.Warnf(format, args...) }
func Errorf(format string, args ...interface{}) { Sugar.Errorf(format, args...) }
func Fatalf(format string, args ...interface{}) { Sugar.Fatalf(format, args...) }

// WithField 带字段的日志
func WithField(key string, value interface{}) *zap.SugaredLogger {
	return Sugar.With(key, value)
}

// WithFields 带多个字段的日志
func WithFields(fields map[string]interface{}) *zap.SugaredLogger {
	f := make([]interface{}, 0, len(fields)*2)
	for k, v := range fields {
		f = append(f, k, v)
	}
	return Sugar.With(f...)
}

// SchedulerLogger 兼容 internal/executor 的日志接口
type SchedulerLogger struct{}

func (s *SchedulerLogger) Infof(format string, args ...interface{}) {
	Sugar.Infof(format, args...)
}
func (s *SchedulerLogger) Warnf(format string, args ...interface{}) {
	Sugar.Warnf(format, args...)
}
func (s *SchedulerLogger) Errorf(format string, args ...interface{}) {
	Sugar.Errorf(format, args...)
}

// NewSchedulerLogger 创建一个兼容 executor.SchedulerLogger 的实例
func NewSchedulerLogger() *SchedulerLogger {
	return &SchedulerLogger{}
}