llm-api-benchmark-tool/logger/logger.go

package logger

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

// 日志级别
const (
	DEBUG = iota
	INFO
	WARN
	ERROR
	FATAL
)

// 日志级别名称
var levelNames = []string{
	"调试",
	"信息",
	"警告",
	"错误",
	"致命",
}

// 日志级别颜色
var levelColors = []string{
	"\033[36m", // 青色 - 调试
	"\033[32m", // 绿色 - 信息
	"\033[33m", // 黄色 - 警告
	"\033[31m", // 红色 - 错误
	"\033[35m", // 紫色 - 致命
}

// 重置颜色
const resetColor = "\033[0m"

// Logger 日志记录器
type Logger struct {
	level      int
	logger     *log.Logger
	fileLogger *log.Logger
	useColors  bool
	timeFormat string
}

// 全局日志记录器
var defaultLogger *Logger

// 初始化默认日志记录器
func init() {
	defaultLogger = NewLogger(INFO, true)
}

// NewLogger 创建新的日志记录器
func NewLogger(level int, useColors bool) *Logger {
	return &Logger{
		level:      level,
		logger:     log.New(os.Stdout, "", 0),
		fileLogger: nil, // 文件日志记录器初始为nil
		useColors:  useColors,
		timeFormat: "2006-01-02 15:04:05.000",
	}
}

// SetLogFile 设置日志文件
func SetLogFile(filePath string) error {
	// 确保目录存在
	dir := filepath.Dir(filePath)
	if err := os.MkdirAll(dir, 0755); err != nil {
		return fmt.Errorf("创建日志目录失败: %w", err)
	}
	
	// 打开日志文件
	file, err := os.OpenFile(filePath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644)
	if err != nil {
		return fmt.Errorf("打开日志文件失败: %w", err)
	}
	
	// 设置文件日志记录器
	defaultLogger.fileLogger = log.New(file, "", 0)
	Info("日志文件已设置: %s", filePath)
	return nil
}

// SetLevel 设置日志级别
func SetLevel(level int) {
	defaultLogger.level = level
}

// SetUseColors 设置是否使用颜色
func SetUseColors(useColors bool) {
	defaultLogger.useColors = useColors
}

// formatLog 格式化日志消息
func (l *Logger) formatLog(level int, format string, args ...interface{}) string {
	// 格式化消息
	var msg string
	if len(args) > 0 {
		msg = fmt.Sprintf(format, args...)
	} else {
		msg = format
	}

	// 获取当前时间
	now := time.Now().Format(l.timeFormat)

	// 获取级别名称
	levelName := levelNames[level]

	// 构建日志前缀
	prefix := fmt.Sprintf("[%s] [%s] ", now, levelName)

	// 添加颜色
	if l.useColors {
		prefix = levelColors[level] + prefix
		msg = msg + resetColor
	}

	// 处理多行消息，确保每行都有正确的前缀
	lines := strings.Split(msg, "\n")
	for i := 1; i < len(lines); i++ {
		if len(lines[i]) > 0 {
			padding := strings.Repeat(" ", len(now)+3) // +3 for "[] "
			lines[i] = fmt.Sprintf("%s%s", padding, lines[i])
		}
	}

	return prefix + strings.Join(lines, "\n")
}

// log 记录日志
func (l *Logger) log(level int, format string, args ...interface{}) {
	if level >= l.level {
		logMsg := l.formatLog(level, format, args...)
		l.logger.Println(logMsg)
		
		// 同时写入文件日志（如果已设置）
		if l.fileLogger != nil {
			// 去除颜色代码
			for _, color := range levelColors {
				logMsg = strings.ReplaceAll(logMsg, color, "")
			}
			logMsg = strings.ReplaceAll(logMsg, resetColor, "")
			
			l.fileLogger.Println(logMsg)
		}
	}
}

// LogToFile 直接写入文件日志，不输出到控制台
func (l *Logger) LogToFile(message string) {
	if l.fileLogger != nil {
		l.fileLogger.Println(message)
	}
}

// LogToFile 直接写入文件日志，不输出到控制台
func LogToFile(message string) {
	defaultLogger.LogToFile(message)
}

// Debug 调试日志
func (l *Logger) Debug(format string, args ...interface{}) {
	l.log(DEBUG, format, args...)
}

// Info 信息日志
func (l *Logger) Info(format string, args ...interface{}) {
	l.log(INFO, format, args...)
}

// Warn 警告日志
func (l *Logger) Warn(format string, args ...interface{}) {
	l.log(WARN, format, args...)
}

// Error 错误日志
func (l *Logger) Error(format string, args ...interface{}) {
	l.log(ERROR, format, args...)
}

// Fatal 致命日志
func (l *Logger) Fatal(format string, args ...interface{}) {
	l.log(FATAL, format, args...)
	os.Exit(1)
}

// Debug 调试日志
func Debug(format string, args ...interface{}) {
	defaultLogger.Debug(format, args...)
}

// Info 信息日志
func Info(format string, args ...interface{}) {
	defaultLogger.Info(format, args...)
}

// Warn 警告日志
func Warn(format string, args ...interface{}) {
	defaultLogger.Warn(format, args...)
}

// Error 错误日志
func Error(format string, args ...interface{}) {
	defaultLogger.Error(format, args...)
}

// Fatal 致命日志
func Fatal(format string, args ...interface{}) {
	defaultLogger.Fatal(format, args...)
}

// 进度条相关
const (
	progressBarWidth = 50
)

// PrintProgressBar 打印进度条
func PrintProgressBar(current, total int, prefix string) {
	percentage := float64(current) / float64(total)
	filled := int(progressBarWidth * percentage)
	empty := progressBarWidth - filled

	// 构建进度条
	bar := fmt.Sprintf("\r%s [%s%s] %.1f%% (%d/%d)",
		prefix,
		strings.Repeat("█", filled),
		strings.Repeat("░", empty),
		percentage*100,
		current,
		total,
	)

	// 打印进度条（不换行）
	fmt.Print(bar)

	// 如果完成，打印换行
	if current == total {
		fmt.Println()
	}
}