package loger

import (
	"errors"
	"fmt"
	"gitee.com/jiangsu-lanjiang-intelligent/ljmicro/pkg/conf"
	"github.com/grafana/loki-client-go/loki"
	"github.com/grafana/loki-client-go/pkg/labelutil"
	"github.com/prometheus/common/model"
	"github.com/sirupsen/logrus"
	"strings"
	"time"
)

type LokiWriter struct {
	*loki.Client
}

func (rcv *LokiWriter) Fire(entry *logrus.Entry) error {
	if entry == nil {
		return fmt.Errorf("log entry is nil")
	}

	ls := model.LabelSet{
		"level": model.LabelValue(entry.Level.String()),
	}
	for k, v := range entry.Data {
		ls[model.LabelName(k)] = model.LabelValue(fmt.Sprintf("%v", v))
	}

	return rcv.Handle(ls, time.Now(), entry.Message)
}

func (rcv *LokiWriter) Levels() []logrus.Level {
	return logrus.AllLevels
}

func InitializeLoki() error {
	lokiAddr := conf.GetString("loki.addr")
	if lokiAddr == "" {
		fmt.Println("loki addr is empty skip initialize log.loki")
		return nil
	}

	if !(strings.HasPrefix(lokiAddr, "http://") ||
		strings.HasPrefix(lokiAddr, "https://")) {
		lokiAddr = "http://" + lokiAddr
	}

	lokiConf, err := loki.NewDefaultConfig(lokiAddr + "/loki/api/v1/push")
	if err != nil {
		return errors.New("new log.loki config error: " + err.Error())
	}
	lokiConf.ExternalLabels = labelutil.LabelSet{LabelSet: model.LabelSet{}}
	c, err := loki.New(lokiConf)
	if err != nil {
		return errors.New("initialize log.loki error: " + err.Error())
	}
	logrus.AddHook(&LokiWriter{c})
	fmt.Println("initialize log.loki success")
	return nil
}