package logger

import (
	"context"
	"time"
)

// contextKey is a custom type for context keys to avoid collisions
type contextKey string

const (
	loggerKey        contextKey = "logger"
	requestIDKey     contextKey = "request_id"
	correlationIDKey contextKey = "correlation_id"
	serviceNameKey   contextKey = "service_name"
)

// ContextWithLogger creates a new context with the logger
func ContextWithLogger(ctx context.Context, logger *Logger) context.Context {
	return context.WithValue(ctx, loggerKey, logger)
}

// LoggerFromContext retrieves the logger from context
func LoggerFromContext(ctx context.Context) *Logger {
	if logger, ok := ctx.Value(loggerKey).(*Logger); ok {
		return logger
	}
	return globalLogger
}

// ContextWithRequestID creates a new context with the request ID
func ContextWithRequestID(ctx context.Context, requestID string) context.Context {
	return context.WithValue(ctx, requestIDKey, requestID)
}

// RequestIDFromContext retrieves the request ID from context
func RequestIDFromContext(ctx context.Context) string {
	if requestID, ok := ctx.Value(requestIDKey).(string); ok {
		return requestID
	}
	return ""
}

// ContextWithCorrelationID creates a new context with the correlation ID
func ContextWithCorrelationID(ctx context.Context, correlationID string) context.Context {
	return context.WithValue(ctx, correlationIDKey, correlationID)
}

// CorrelationIDFromContext retrieves the correlation ID from context
func CorrelationIDFromContext(ctx context.Context) string {
	if correlationID, ok := ctx.Value(correlationIDKey).(string); ok {
		return correlationID
	}
	return ""
}

// ContextWithServiceName creates a new context with the service name
func ContextWithServiceName(ctx context.Context, serviceName string) context.Context {
	return context.WithValue(ctx, serviceNameKey, serviceName)
}

// ServiceNameFromContext retrieves the service name from context
func ServiceNameFromContext(ctx context.Context) string {
	if serviceName, ok := ctx.Value(serviceNameKey).(string); ok {
		return serviceName
	}
	return ""
}

// WithContext returns a new logger with context values
func (l *Logger) WithContext(ctx context.Context) *Logger {
	logger := l

	if requestID := RequestIDFromContext(ctx); requestID != "" {
		logger = logger.WithRequestID(requestID)
	}

	if correlationID := CorrelationIDFromContext(ctx); correlationID != "" {
		logger = logger.WithCorrelationID(correlationID)
	}

	if serviceName := ServiceNameFromContext(ctx); serviceName != "" {
		logger = logger.WithService(serviceName)
	}

	return logger
}

// DebugCtx logs a debug message with context
func DebugCtx(ctx context.Context, msg string, fields ...map[string]interface{}) {
	LoggerFromContext(ctx).WithContext(ctx).Debug(msg, fields...)
}

// DebugfCtx logs a formatted debug message with context
func DebugfCtx(ctx context.Context, format string, args ...interface{}) {
	LoggerFromContext(ctx).WithContext(ctx).Debugf(format, args...)
}

// InfoCtx logs an info message with context
func InfoCtx(ctx context.Context, msg string, fields ...map[string]interface{}) {
	LoggerFromContext(ctx).WithContext(ctx).Info(msg, fields...)
}

// InfofCtx logs a formatted info message with context
func InfofCtx(ctx context.Context, format string, args ...interface{}) {
	LoggerFromContext(ctx).WithContext(ctx).Infof(format, args...)
}

// WarnCtx logs a warning message with context
func WarnCtx(ctx context.Context, msg string, fields ...map[string]interface{}) {
	LoggerFromContext(ctx).WithContext(ctx).Warn(msg, fields...)
}

// WarnfCtx logs a formatted warning message with context
func WarnfCtx(ctx context.Context, format string, args ...interface{}) {
	LoggerFromContext(ctx).WithContext(ctx).Warnf(format, args...)
}

// ErrorCtx logs an error message with context
func ErrorCtx(ctx context.Context, msg string, fields ...map[string]interface{}) {
	LoggerFromContext(ctx).WithContext(ctx).Error(msg, fields...)
}

// ErrorfCtx logs a formatted error message with context
func ErrorfCtx(ctx context.Context, format string, args ...interface{}) {
	LoggerFromContext(ctx).WithContext(ctx).Errorf(format, args...)
}

// FatalCtx logs a fatal message with context and exits the program
func FatalCtx(ctx context.Context, msg string, fields ...map[string]interface{}) {
	LoggerFromContext(ctx).WithContext(ctx).Fatal(msg, fields...)
}

// FatalfCtx logs a formatted fatal message with context and exits the program
func FatalfCtx(ctx context.Context, format string, args ...interface{}) {
	LoggerFromContext(ctx).WithContext(ctx).Fatalf(format, args...)
}

// LogDurationCtx logs the duration of an operation with context
func LogDurationCtx(ctx context.Context, start time.Time, operation string, fields ...map[string]interface{}) {
	LoggerFromContext(ctx).WithContext(ctx).LogDuration(start, operation, fields...)
}