package commonctx

import (
	"context"

	"github.com/sirupsen/logrus"
)

// AddMainLogger will add the provided logrus entry as the main logger in context.
func AddMainLogger(ctx context.Context, logger *logrus.Entry) context.Context {
	return context.WithValue(ctx, ContextKey_MainLogger, logger)
}

// GetLogger retrieve logger from the context.
func GetLogger(ctx context.Context) *logrus.Entry {
	if log := ctx.Value(ContextKey_MainLogger); log != nil {
		return log.(*logrus.Entry)
	}
	return nil
}

// GetContextAndLoggerWithField will add a key and value field to the logger, update the context, and return the new logger.
func GetContextAndLoggerWithField(ctx context.Context, key string, value interface{}) (context.Context, *logrus.Entry) {
	if log := ctx.Value(ContextKey_MainLogger); log != nil {
		log, ok := log.(*logrus.Entry)
		if !ok {
			panic("invalid logger in context")
		}
		log = log.WithField(key, value)
		ctx = AddMainLogger(ctx, log)
		return ctx, log
	}
	return ctx, nil
}

// GetContextAndLoggerWithFields will add provided fields to the logger, update the context, and return the new logger.
func GetContextAndLoggerWithFields(ctx context.Context, fields logrus.Fields) (context.Context, *logrus.Entry) {
	if log := ctx.Value(ContextKey_MainLogger); log != nil {
		log, ok := log.(*logrus.Entry)
		if !ok {
			panic("invalid logger in context")
		}
		log = log.WithFields(fields)
		ctx = AddMainLogger(ctx, log)
		return ctx, log
	}
	return ctx, nil
}