AIRouter/internal/observability/logger.go

package observability

import (
	"encoding/json"
	"fmt"
	"os"
	"strings"
	"sync"
	"time"
)

// LogLevel represents logging severity levels.
type LogLevel int

const (
	LevelDebug LogLevel = iota
	LevelInfo
	LevelWarn
	LevelError
)

func (l LogLevel) String() string {
	switch l {
	case LevelDebug:
		return "DEBUG"
	case LevelInfo:
		return "INFO"
	case LevelWarn:
		return "WARN"
	case LevelError:
		return "ERROR"
	default:
		return "INFO"
	}
}

// ParseLogLevel parses a string to LogLevel.
func ParseLogLevel(s string) LogLevel {
	switch strings.ToLower(s) {
	case "debug":
		return LevelDebug
	case "info":
		return LevelInfo
	case "warn", "warning":
		return LevelWarn
	case "error":
		return LevelError
	default:
		return LevelInfo
	}
}

// LogEntry is a structured JSON log entry.
type LogEntry struct {
	Timestamp  string `json:"timestamp"`
	Level      string `json:"level"`
	Event      string `json:"event,omitempty"`
	Message    string `json:"message,omitempty"`
	RequestID  string `json:"request_id,omitempty"`
	TaskID     string `json:"task_id,omitempty"`
	SessionID  string `json:"session_id,omitempty"`
	TraceID    string `json:"trace_id,omitempty"`
	Application string `json:"application,omitempty"`
	TenantID   string `json:"tenant_id,omitempty"`
	UserID     string `json:"user_id,omitempty"`
	FromState  string `json:"from_state,omitempty"`
	ToState    string `json:"to_state,omitempty"`
	Reason     string `json:"reason,omitempty"`
	LogicalModel string `json:"logical_model,omitempty"`
	ActualModel  string `json:"actual_model,omitempty"`
	NodeID     string `json:"node_id,omitempty"`
	Degraded   bool   `json:"degraded,omitempty"`
	Extra      map[string]any `json:"extra,omitempty"`
}

// Logger is a structured JSON logger with sensitive field masking.
type Logger struct {
	mu       sync.RWMutex
	level    LogLevel
	output   *os.File
	maskFields []string
	promptLogging string
}

var defaultLogger *Logger

func init() {
	defaultLogger = NewLogger(LevelInfo, os.Stdout, "metadata_only")
}

// NewLogger creates a new Logger instance.
func NewLogger(level LogLevel, out *os.File, promptLogging string) *Logger {
	return &Logger{
		level:    level,
		output:   out,
		maskFields: []string{"api_key", "apikey", "authorization", "jwt", "secret", "password", "token"},
		promptLogging: promptLogging,
	}
}

// GetLogger returns the default logger.
func GetLogger() *Logger {
	return defaultLogger
}

// SetLevel updates the log level (thread-safe).
func (l *Logger) SetLevel(level LogLevel) {
	l.mu.Lock()
	l.level = level
	l.mu.Unlock()
}

// SetPromptLogging updates the prompt logging policy.
func (l *Logger) SetPromptLogging(policy string) {
	l.mu.Lock()
	l.promptLogging = policy
	l.mu.Unlock()
}

func (l *Logger) shouldLog(level LogLevel) bool {
	l.mu.RLock()
	defer l.mu.RUnlock()
	return level >= l.level
}

func (l *Logger) maskSensitive(data map[string]any) map[string]any {
	if data == nil {
		return nil
	}
	masked := make(map[string]any, len(data))
	for k, v := range data {
		if l.isSensitive(k) {
			masked[k] = "***REDACTED***"
		} else if sub, ok := v.(map[string]any); ok {
			masked[k] = l.maskSensitive(sub)
		} else {
			masked[k] = v
		}
	}
	return masked
}

func (l *Logger) isSensitive(key string) bool {
	lk := strings.ToLower(key)
	for _, s := range l.maskFields {
		if strings.Contains(lk, s) {
			return true
		}
	}
	return false
}

func (l *Logger) write(entry LogEntry) {
	if !l.shouldLog(parseLevelFromString(entry.Level)) {
		return
	}
	if entry.Extra != nil {
		entry.Extra = l.maskSensitive(entry.Extra)
	}
	if entry.Timestamp == "" {
		entry.Timestamp = time.Now().UTC().Format(time.RFC3339Nano)
	}
	data, err := json.Marshal(entry)
	if err != nil {
		fmt.Fprintf(os.Stderr, "log marshal error: %v\n", err)
		return
	}
	l.mu.Lock()
	fmt.Fprintln(l.output, string(data))
	l.mu.Unlock()
}

func parseLevelFromString(s string) LogLevel {
	switch strings.ToUpper(s) {
	case "DEBUG":
		return LevelDebug
	case "INFO":
		return LevelInfo
	case "WARN", "WARNING":
		return LevelWarn
	case "ERROR":
		return LevelError
	default:
		return LevelInfo
	}
}

// LogFields is a builder for structured log fields.
type LogFields struct {
	fields map[string]any
}

func F() *LogFields {
	return &LogFields{fields: make(map[string]any)}
}

func (f *LogFields) Set(key string, value any) *LogFields {
	f.fields[key] = value
	return f
}

func (f *LogFields) RequestID(id string) *LogFields    { f.fields["request_id"] = id; return f }
func (f *LogFields) TaskID(id string) *LogFields       { f.fields["task_id"] = id; return f }
func (f *LogFields) SessionID(id string) *LogFields    { f.fields["session_id"] = id; return f }
func (f *LogFields) TraceID(id string) *LogFields      { f.fields["trace_id"] = id; return f }
func (f *LogFields) Application(app string) *LogFields { f.fields["application"] = app; return f }
func (f *LogFields) TenantID(id string) *LogFields     { f.fields["tenant_id"] = id; return f }
func (f *LogFields) UserID(id string) *LogFields       { f.fields["user_id"] = id; return f }
func (f *LogFields) Event(e string) *LogFields        { f.fields["event"] = e; return f }
func (f *LogFields) Reason(r string) *LogFields       { f.fields["reason"] = r; return f }

func (l *Logger) Debug(msg string, fields *LogFields) {
	entry := l.buildEntry("DEBUG", msg, fields)
	l.write(entry)
}

func (l *Logger) Info(msg string, fields *LogFields) {
	entry := l.buildEntry("INFO", msg, fields)
	l.write(entry)
}

func (l *Logger) Warn(msg string, fields *LogFields) {
	entry := l.buildEntry("WARN", msg, fields)
	l.write(entry)
}

func (l *Logger) Error(msg string, fields *LogFields) {
	entry := l.buildEntry("ERROR", msg, fields)
	l.write(entry)
}

func (l *Logger) buildEntry(level, msg string, fields *LogFields) LogEntry {
	entry := LogEntry{
		Level:   level,
		Message: msg,
	}
	if fields != nil {
		for k, v := range fields.fields {
			switch k {
			case "event":
				if s, ok := v.(string); ok {
					entry.Event = s
				}
			case "request_id":
				if s, ok := v.(string); ok {
					entry.RequestID = s
				}
			case "task_id":
				if s, ok := v.(string); ok {
					entry.TaskID = s
				}
			case "session_id":
				if s, ok := v.(string); ok {
					entry.SessionID = s
				}
			case "trace_id":
				if s, ok := v.(string); ok {
					entry.TraceID = s
				}
			case "application":
				if s, ok := v.(string); ok {
					entry.Application = s
				}
			case "tenant_id":
				if s, ok := v.(string); ok {
					entry.TenantID = s
				}
			case "user_id":
				if s, ok := v.(string); ok {
					entry.UserID = s
				}
			case "reason":
				if s, ok := v.(string); ok {
					entry.Reason = s
				}
			case "from_state":
				if s, ok := v.(string); ok {
					entry.FromState = s
				}
			case "to_state":
				if s, ok := v.(string); ok {
					entry.ToState = s
				}
			case "logical_model":
				if s, ok := v.(string); ok {
					entry.LogicalModel = s
				}
			case "actual_model":
				if s, ok := v.(string); ok {
					entry.ActualModel = s
				}
			case "node_id":
				if s, ok := v.(string); ok {
					entry.NodeID = s
				}
			case "degraded":
				if b, ok := v.(bool); ok {
					entry.Degraded = b
				}
			default:
				if entry.Extra == nil {
					entry.Extra = make(map[string]any)
				}
				entry.Extra[k] = v
			}
		}
	}
	return entry
}

// SetLogLevel updates the global log level.
func SetLogLevel(level string) {
	defaultLogger.SetLevel(ParseLogLevel(level))
}

// SetPromptLoggingPolicy updates the global prompt logging policy.
func SetPromptLoggingPolicy(policy string) {
	defaultLogger.SetPromptLogging(policy)
}