fscan/Common/logging/Logger.go

package logging

import (
	"fmt"
	"io"
	"log"
	"path/filepath"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/fatih/color"
)

// Logger 日志管理器
type Logger struct {
	mu            sync.RWMutex
	config        *LoggerConfig
	formatter     LogFormatter
	handlers      []LogHandler
	scanStatus    *ScanStatus
	progressBar   ProgressDisplay
	outputMutex   *sync.Mutex
	initialized   bool
}

// NewLogger 创建新的日志管理器
func NewLogger(config *LoggerConfig) *Logger {
	if config == nil {
		config = DefaultLoggerConfig()
	}

	logger := &Logger{
		config:      config,
		formatter:   NewStandardFormatter(),
		handlers:    make([]LogHandler, 0),
		scanStatus:  NewScanStatus(),
		outputMutex: &sync.Mutex{},
		initialized: true,
	}

	// 设置格式化器的开始时间
	logger.formatter.SetStartTime(config.StartTime)

	// 添加默认的控制台处理器
	consoleHandler := NewConsoleHandler(config)
	logger.AddHandler(consoleHandler)

	return logger
}

// SetFormatter 设置格式化器
func (l *Logger) SetFormatter(formatter LogFormatter) {
	l.mu.Lock()
	defer l.mu.Unlock()
	l.formatter = formatter
	l.formatter.SetStartTime(l.config.StartTime)
}

// AddHandler 添加日志处理器
func (l *Logger) AddHandler(handler LogHandler) {
	l.mu.Lock()
	defer l.mu.Unlock()
	l.handlers = append(l.handlers, handler)
}

// SetProgressBar 设置进度条显示
func (l *Logger) SetProgressBar(progressBar ProgressDisplay) {
	l.mu.Lock()
	defer l.mu.Unlock()
	l.progressBar = progressBar
}

// SetOutputMutex 设置输出互斥锁
func (l *Logger) SetOutputMutex(mutex *sync.Mutex) {
	l.mu.Lock()
	defer l.mu.Unlock()
	l.outputMutex = mutex
}

// Log 记录日志
func (l *Logger) Log(level LogLevel, content string, metadata ...map[string]interface{}) {
	if !l.shouldLog(level) {
		return
	}

	entry := &LogEntry{
		Level:   level,
		Time:    time.Now(),
		Content: content,
	}

	// 添加元数据
	if len(metadata) > 0 {
		entry.Metadata = metadata[0]
	}

	// 对于错误级别，自动添加调用者信息
	if level == LevelError {
		if _, file, line, ok := runtime.Caller(2); ok {
			entry.Source = fmt.Sprintf("%s:%d", filepath.Base(file), line)
			entry.Content = fmt.Sprintf("%s:%d - %s", filepath.Base(file), line, content)
		}
	}

	l.handleLogEntry(entry)

	// 更新扫描状态
	if level == LevelSuccess {
		l.scanStatus.UpdateSuccess()
	} else if level == LevelError {
		l.scanStatus.UpdateError()
	}
}

// shouldLog 检查是否应该记录该级别的日志
func (l *Logger) shouldLog(level LogLevel) bool {
	switch l.config.Level {
	case LevelAll:
		return true
	case LevelBaseInfoSuccess:
		return level == LevelBase || level == LevelInfo || level == LevelSuccess
	case LevelInfoSuccess:
		return level == LevelInfo || level == LevelSuccess
	case LevelError:
		return level == LevelError
	case LevelBase:
		return level == LevelBase
	case LevelInfo:
		return level == LevelInfo
	case LevelSuccess:
		return level == LevelSuccess
	case LevelDebug:
		return level == LevelDebug
	default:
		// 向后兼容：如果是字符串 "debug"，显示所有
		if l.config.Level == "debug" {
			return true
		}
		// 默认显示base、info和success
		return level == LevelBase || level == LevelInfo || level == LevelSuccess
	}
}

// handleLogEntry 处理日志条目
func (l *Logger) handleLogEntry(entry *LogEntry) {
	l.outputMutex.Lock()
	defer l.outputMutex.Unlock()

	// 清除进度条
	l.clearProgress()

	// 使用所有处理器处理日志
	l.mu.RLock()
	for _, handler := range l.handlers {
		if handler.IsEnabled() {
			handler.Handle(entry)
		}
	}
	l.mu.RUnlock()

	// 恢复进度条
	l.restoreProgress()
}

// clearProgress 清除进度条
func (l *Logger) clearProgress() {
	if l.progressBar != nil {
		l.progressBar.Clear() // 忽略错误
		time.Sleep(10 * time.Millisecond)
	}
}

// restoreProgress 恢复进度条
func (l *Logger) restoreProgress() {
	if l.progressBar != nil {
		l.progressBar.RenderBlank() // 忽略错误
	}
}

// 便利方法
func (l *Logger) Debug(content string, metadata ...map[string]interface{}) {
	l.Log(LevelDebug, content, metadata...)
}

func (l *Logger) Base(content string, metadata ...map[string]interface{}) {
	l.Log(LevelBase, content, metadata...)
}

func (l *Logger) Info(content string, metadata ...map[string]interface{}) {
	l.Log(LevelInfo, content, metadata...)
}

func (l *Logger) Success(content string, metadata ...map[string]interface{}) {
	l.Log(LevelSuccess, content, metadata...)
}

func (l *Logger) Error(content string, metadata ...map[string]interface{}) {
	l.Log(LevelError, content, metadata...)
}

// GetScanStatus 获取扫描状态管理器
func (l *Logger) GetScanStatus() *ScanStatus {
	return l.scanStatus
}

// Initialize 初始化日志系统（兼容原接口）
func (l *Logger) Initialize() {
	// 禁用标准日志输出
	log.SetOutput(io.Discard)
}

// ConsoleHandler 控制台日志处理器
type ConsoleHandler struct {
	config    *LoggerConfig
	formatter LogFormatter
	enabled   bool
	mu        sync.RWMutex
}

// NewConsoleHandler 创建控制台处理器
func NewConsoleHandler(config *LoggerConfig) *ConsoleHandler {
	formatter := NewStandardFormatter()
	formatter.SetStartTime(config.StartTime)
	
	return &ConsoleHandler{
		config:    config,
		formatter: formatter,
		enabled:   true,
	}
}

// Handle 处理日志条目
func (h *ConsoleHandler) Handle(entry *LogEntry) {
	h.mu.RLock()
	defer h.mu.RUnlock()

	if !h.enabled {
		return
	}

	// 使用自己的格式化器格式化消息
	logMsg := h.formatter.Format(entry)

	// 根据颜色设置输出
	if h.config.EnableColor {
		if colorAttr, ok := h.config.LevelColors[entry.Level]; ok {
			color.New(colorAttr).Println(logMsg)
		} else {
			fmt.Println(logMsg)
		}
	} else {
		fmt.Println(logMsg)
	}

	// 根据慢速输出设置决定是否添加延迟
	if h.config.SlowOutput {
		time.Sleep(50 * time.Millisecond)
	}
}

// SetEnabled 设置处理器启用状态
func (h *ConsoleHandler) SetEnabled(enabled bool) {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.enabled = enabled
}

// IsEnabled 检查处理器是否启用
func (h *ConsoleHandler) IsEnabled() bool {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.enabled
}

// 全局日志管理器实例
var (
	globalLogger *Logger
	initOnce     sync.Once
)

// GetGlobalLogger 获取全局日志管理器
func GetGlobalLogger() *Logger {
	initOnce.Do(func() {
		globalLogger = NewLogger(DefaultLoggerConfig())
	})
	return globalLogger
}

// SetGlobalLogger 设置全局日志管理器
func SetGlobalLogger(logger *Logger) {
	globalLogger = logger
}

// 错误检查函数（保持原有逻辑）
func CheckErrs(err error) error {
	if err == nil {
		return nil
	}

	// 已知需要重试的错误列表
	errs := []string{
		"closed by the remote host", "too many connections",
		"EOF", "A connection attempt failed",
		"established connection failed", "connection attempt failed",
		"Unable to read", "is not allowed to connect to this",
		"no pg_hba.conf entry",
		"No connection could be made",
		"invalid packet size",
		"bad connection",
	}

	// 检查错误是否匹配
	errLower := strings.ToLower(err.Error())
	for _, key := range errs {
		if strings.Contains(errLower, strings.ToLower(key)) {
			time.Sleep(1 * time.Second)
			return err
		}
	}

	return nil
}