package elastic

import (
	"bytes"
	"context"
	"encoding/json"
	"time"

	"github.com/elastic/go-elasticsearch"
	"github.com/sirupsen/logrus"
)

type ElasticHook struct {
	Client *elasticsearch.Client
	Index  string
}

func NewElasticHook(client *elasticsearch.Client, index string) *ElasticHook {
	return &ElasticHook{
		Client: client,
		Index:  index,
	}
}

func (hook *ElasticHook) Fire(entry *logrus.Entry) error {
	doc := map[string]interface{}{
		"@timestamp": entry.Time.Format(time.RFC3339),
		"timestamp":  entry.Time.Format(time.RFC3339),
		"level":      entry.Level.String(),
		"message":    entry.Message,
		"module": func() interface{} {
			if val, ok := entry.Data["module"]; ok {
				return val
			}
			return "default"
		}(),
		"file": entry.Caller.File,
		"line": entry.Caller.Line,
	}

	var buf bytes.Buffer
	if err := json.NewEncoder(&buf).Encode(doc); err != nil {
		return err
	}

	_, err := hook.Client.Index(
		hook.Index,
		&buf,
		hook.Client.Index.WithContext(context.Background()),
		hook.Client.Index.WithRefresh("true"),
	)
	return err
}

func (hook *ElasticHook) Levels() []logrus.Level {
	return logrus.AllLevels
}