package pgxlogrus

import (
	"context"
	"fmt"
	"time"

	"github.com/jackc/pgx/v4"
	"github.com/sirupsen/logrus"
)

type pgxLogger struct {
	logger logrus.FieldLogger
}

func New(logger logrus.FieldLogger) *pgxLogger {
	return &pgxLogger{
		logger: logger,
	}
}

func (lg pgxLogger) Log(ctx context.Context, level pgx.LogLevel, msg string, data map[string]interface{}) {
	lg.logger.WithFields(lg.fieldsData(data)).Info(msg)
}

func (lg pgxLogger) fieldsData(data map[string]interface{}) logrus.Fields {
	fields := make(logrus.Fields)
	if v, ok := data["time"]; ok && v != nil {
		fields["cost"] = v
		fields["costString"] = fmt.Sprintf("%v", time.Duration(v.(time.Duration)))
	}

	if v, ok := data["sql"]; ok {
		fields["SQL"] = v
	}

	if v, ok := data["pid"]; ok {
		fields["PID"] = v
	}

	if v, ok := data["args"]; ok {
		fields["args"] = v
	}

	if v, ok := data["rowCount"]; ok {
		fields["rowCount"] = v
	}

	return fields
}