package logger

import (
	"context"
	"log"
	"os"

	"github.com/go-logr/logr"
	"github.com/go-logr/stdr"
	"github.com/go-logr/zapr"
	"go.opentelemetry.io/otel/trace"
	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
	"gopkg.in/natefinch/lumberjack.v2"

	"github.com/tidbcloud/aisupport/packages/sdk-go/infra/config"
)

var (
	Logger       logr.Logger = stdr.New(log.New(os.Stderr, "STD_LOG: ", log.LstdFlags))
	ExtraLoggers map[string]logr.Logger
)

func InitLogger(conf *config.Loggers) error {
	Logger = initLogger(conf.Logger)

	ExtraLoggers = make(map[string]logr.Logger, len(conf.ExtraLoggers))
	for name, loggerConf := range conf.ExtraLoggers {
		ExtraLoggers[name] = initLogger(loggerConf)
	}

	return nil
}

func initLogger(conf config.Logger) logr.Logger {
	zapConf := zap.NewProductionEncoderConfig()
	zapConf.EncodeTime = zapcore.RFC3339TimeEncoder
	zapConf.EncodeLevel = zapcore.CapitalLevelEncoder

	var encoder zapcore.Encoder = zapcore.NewConsoleEncoder(zapConf)
	if conf.IsJSON {
		encoder = zapcore.NewJSONEncoder(zapConf)
	}

	writer := zapcore.AddSync(os.Stderr)
	if conf.Filepath != "" {
		writer = zapcore.AddSync(&lumberjack.Logger{
			Filename:   conf.Filepath,
			MaxSize:    conf.FileMaxSize,
			MaxBackups: conf.FileMaxBackups,
			MaxAge:     conf.FileMaxAge,
			Compress:   false,
		})
	}

	core := zapcore.NewCore(
		encoder,
		writer,
		zap.NewAtomicLevelAt(zapcore.Level(-1*conf.Level)),
	)
	opts := []zap.Option{}
	if !conf.DisableCaller {
		opts = append(opts, zap.AddCaller())
	}
	if !conf.DisableStacktrace {
		opts = append(opts, zap.AddStacktrace(zapcore.ErrorLevel))
	}

	return zapr.NewLogger(zap.New(core, opts...))
}

// NewContext returns a new Context, derived from ctx, which carries the
// provided Logger.
// It is same as `logr.NewContext`.
func NewContext(ctx context.Context, logger logr.Logger, keysAndValues ...any) context.Context {
	return logr.NewContext(ctx, logger.WithValues(keysAndValues...))
}

func WrapContext(ctx context.Context, keysAndValues ...any) context.Context {
	logger := FromContext(ctx)
	return logr.NewContext(ctx, logger.WithValues(keysAndValues...))
}

// FromContextOrDef returns a Logger from ctx.  If no Logger is found, this
// returns the default Logger.
//
// `logr.FromContextOrDiscard` may return a zero logger, it is not we want.
func FromContextOrDef(ctx context.Context, def logr.Logger, keysAndValues ...any) logr.Logger {
	logger := logr.FromContextOrDiscard(ctx)
	if logger.IsZero() {
		logger = def
	}

	return WithTraceID(ctx, logger).WithValues(keysAndValues...)
}

// FromContext returns a Logger from ctx.  If no Logger is found, this
// returns the global Logger.
//
// `logr.FromContextOrDiscard` may return a zero logger, it is not we want.
func FromContext(ctx context.Context, keysAndValues ...any) logr.Logger {
	return FromContextOrDef(ctx, Logger, keysAndValues...)
}

// WithTraceID returns a Logger with trace_id field.
func WithTraceID(ctx context.Context, logger logr.Logger) logr.Logger {
	spanCtx := trace.SpanContextFromContext(ctx)
	if spanCtx.HasTraceID() {
		logger = logger.WithValues("trace_id", spanCtx.TraceID().String())
	} else {
		logger = logger.WithValues("trace_id", "none")
	}
	return logger
}

type LoggerToWriter struct {
	logger logr.Logger
}

func NewLogrToWriter(logger logr.Logger) *LoggerToWriter {
	return &LoggerToWriter{
		logger: logger,
	}
}

func (l *LoggerToWriter) Write(p []byte) (n int, err error) {
	l.logger.Info(string(p))
	return len(p), nil
}