1package log
2
3import (
4	"bytes"
5	"encoding/json"
6	"fmt"
7	"io"
8	"os"
9	"path"
10	"runtime"
11	"strconv"
12	"sync"
13	"sync/atomic"
14	"time"
15
16	"github.com/mattn/go-isatty"
17	"github.com/valyala/fasttemplate"
18
19	"github.com/labstack/gommon/color"
20)
21
22type (
23	Logger struct {
24		prefix     string
25		level      uint32
26		skip       int
27		output     io.Writer
28		template   *fasttemplate.Template
29		levels     []string
30		color      *color.Color
31		bufferPool sync.Pool
32		mutex      sync.Mutex
33	}
34
35	Lvl uint8
36
37	JSON map[string]interface{}
38)
39
40const (
41	DEBUG Lvl = iota + 1
42	INFO
43	WARN
44	ERROR
45	OFF
46	panicLevel
47	fatalLevel
48)
49
50var (
51	global        = New("-")
52	defaultHeader = `{"time":"${time_rfc3339_nano}","level":"${level}","prefix":"${prefix}",` +
53		`"file":"${short_file}","line":"${line}"}`
54)
55
56func init() {
57	global.skip = 3
58}
59
60func New(prefix string) (l *Logger) {
61	l = &Logger{
62		level:    uint32(INFO),
63		skip:     2,
64		prefix:   prefix,
65		template: l.newTemplate(defaultHeader),
66		color:    color.New(),
67		bufferPool: sync.Pool{
68			New: func() interface{} {
69				return bytes.NewBuffer(make([]byte, 256))
70			},
71		},
72	}
73	l.initLevels()
74	l.SetOutput(output())
75	return
76}
77
78func (l *Logger) initLevels() {
79	l.levels = []string{
80		"-",
81		l.color.Blue("DEBUG"),
82		l.color.Green("INFO"),
83		l.color.Yellow("WARN"),
84		l.color.Red("ERROR"),
85		"",
86		l.color.Yellow("PANIC", color.U),
87		l.color.Red("FATAL", color.U),
88	}
89}
90
91func (l *Logger) newTemplate(format string) *fasttemplate.Template {
92	return fasttemplate.New(format, "${", "}")
93}
94
95func (l *Logger) DisableColor() {
96	l.color.Disable()
97	l.initLevels()
98}
99
100func (l *Logger) EnableColor() {
101	l.color.Enable()
102	l.initLevels()
103}
104
105func (l *Logger) Prefix() string {
106	return l.prefix
107}
108
109func (l *Logger) SetPrefix(p string) {
110	l.prefix = p
111}
112
113func (l *Logger) Level() Lvl {
114	return Lvl(atomic.LoadUint32(&l.level))
115}
116
117func (l *Logger) SetLevel(level Lvl) {
118	atomic.StoreUint32(&l.level, uint32(level))
119}
120
121func (l *Logger) Output() io.Writer {
122	return l.output
123}
124
125func (l *Logger) SetOutput(w io.Writer) {
126	l.output = w
127	if w, ok := w.(*os.File); !ok || !isatty.IsTerminal(w.Fd()) {
128		l.DisableColor()
129	}
130}
131
132func (l *Logger) Color() *color.Color {
133	return l.color
134}
135
136func (l *Logger) SetHeader(h string) {
137	l.template = l.newTemplate(h)
138}
139
140func (l *Logger) Print(i ...interface{}) {
141	l.log(0, "", i...)
142	// fmt.Fprintln(l.output, i...)
143}
144
145func (l *Logger) Printf(format string, args ...interface{}) {
146	l.log(0, format, args...)
147}
148
149func (l *Logger) Printj(j JSON) {
150	l.log(0, "json", j)
151}
152
153func (l *Logger) Debug(i ...interface{}) {
154	l.log(DEBUG, "", i...)
155}
156
157func (l *Logger) Debugf(format string, args ...interface{}) {
158	l.log(DEBUG, format, args...)
159}
160
161func (l *Logger) Debugj(j JSON) {
162	l.log(DEBUG, "json", j)
163}
164
165func (l *Logger) Info(i ...interface{}) {
166	l.log(INFO, "", i...)
167}
168
169func (l *Logger) Infof(format string, args ...interface{}) {
170	l.log(INFO, format, args...)
171}
172
173func (l *Logger) Infoj(j JSON) {
174	l.log(INFO, "json", j)
175}
176
177func (l *Logger) Warn(i ...interface{}) {
178	l.log(WARN, "", i...)
179}
180
181func (l *Logger) Warnf(format string, args ...interface{}) {
182	l.log(WARN, format, args...)
183}
184
185func (l *Logger) Warnj(j JSON) {
186	l.log(WARN, "json", j)
187}
188
189func (l *Logger) Error(i ...interface{}) {
190	l.log(ERROR, "", i...)
191}
192
193func (l *Logger) Errorf(format string, args ...interface{}) {
194	l.log(ERROR, format, args...)
195}
196
197func (l *Logger) Errorj(j JSON) {
198	l.log(ERROR, "json", j)
199}
200
201func (l *Logger) Fatal(i ...interface{}) {
202	l.log(fatalLevel, "", i...)
203	os.Exit(1)
204}
205
206func (l *Logger) Fatalf(format string, args ...interface{}) {
207	l.log(fatalLevel, format, args...)
208	os.Exit(1)
209}
210
211func (l *Logger) Fatalj(j JSON) {
212	l.log(fatalLevel, "json", j)
213	os.Exit(1)
214}
215
216func (l *Logger) Panic(i ...interface{}) {
217	l.log(panicLevel, "", i...)
218	panic(fmt.Sprint(i...))
219}
220
221func (l *Logger) Panicf(format string, args ...interface{}) {
222	l.log(panicLevel, format, args...)
223	panic(fmt.Sprintf(format, args...))
224}
225
226func (l *Logger) Panicj(j JSON) {
227	l.log(panicLevel, "json", j)
228	panic(j)
229}
230
231func DisableColor() {
232	global.DisableColor()
233}
234
235func EnableColor() {
236	global.EnableColor()
237}
238
239func Prefix() string {
240	return global.Prefix()
241}
242
243func SetPrefix(p string) {
244	global.SetPrefix(p)
245}
246
247func Level() Lvl {
248	return global.Level()
249}
250
251func SetLevel(level Lvl) {
252	global.SetLevel(level)
253}
254
255func Output() io.Writer {
256	return global.Output()
257}
258
259func SetOutput(w io.Writer) {
260	global.SetOutput(w)
261}
262
263func SetHeader(h string) {
264	global.SetHeader(h)
265}
266
267func Print(i ...interface{}) {
268	global.Print(i...)
269}
270
271func Printf(format string, args ...interface{}) {
272	global.Printf(format, args...)
273}
274
275func Printj(j JSON) {
276	global.Printj(j)
277}
278
279func Debug(i ...interface{}) {
280	global.Debug(i...)
281}
282
283func Debugf(format string, args ...interface{}) {
284	global.Debugf(format, args...)
285}
286
287func Debugj(j JSON) {
288	global.Debugj(j)
289}
290
291func Info(i ...interface{}) {
292	global.Info(i...)
293}
294
295func Infof(format string, args ...interface{}) {
296	global.Infof(format, args...)
297}
298
299func Infoj(j JSON) {
300	global.Infoj(j)
301}
302
303func Warn(i ...interface{}) {
304	global.Warn(i...)
305}
306
307func Warnf(format string, args ...interface{}) {
308	global.Warnf(format, args...)
309}
310
311func Warnj(j JSON) {
312	global.Warnj(j)
313}
314
315func Error(i ...interface{}) {
316	global.Error(i...)
317}
318
319func Errorf(format string, args ...interface{}) {
320	global.Errorf(format, args...)
321}
322
323func Errorj(j JSON) {
324	global.Errorj(j)
325}
326
327func Fatal(i ...interface{}) {
328	global.Fatal(i...)
329}
330
331func Fatalf(format string, args ...interface{}) {
332	global.Fatalf(format, args...)
333}
334
335func Fatalj(j JSON) {
336	global.Fatalj(j)
337}
338
339func Panic(i ...interface{}) {
340	global.Panic(i...)
341}
342
343func Panicf(format string, args ...interface{}) {
344	global.Panicf(format, args...)
345}
346
347func Panicj(j JSON) {
348	global.Panicj(j)
349}
350
351func (l *Logger) log(level Lvl, format string, args ...interface{}) {
352	if level >= l.Level() || level == 0 {
353		buf := l.bufferPool.Get().(*bytes.Buffer)
354		buf.Reset()
355		defer l.bufferPool.Put(buf)
356		_, file, line, _ := runtime.Caller(l.skip)
357		message := ""
358
359		if format == "" {
360			message = fmt.Sprint(args...)
361		} else if format == "json" {
362			b, err := json.Marshal(args[0])
363			if err != nil {
364				panic(err)
365			}
366			message = string(b)
367		} else {
368			message = fmt.Sprintf(format, args...)
369		}
370
371		_, err := l.template.ExecuteFunc(buf, func(w io.Writer, tag string) (int, error) {
372			switch tag {
373			case "time_rfc3339":
374				return w.Write([]byte(time.Now().Format(time.RFC3339)))
375			case "time_rfc3339_nano":
376				return w.Write([]byte(time.Now().Format(time.RFC3339Nano)))
377			case "level":
378				return w.Write([]byte(l.levels[level]))
379			case "prefix":
380				return w.Write([]byte(l.prefix))
381			case "long_file":
382				return w.Write([]byte(file))
383			case "short_file":
384				return w.Write([]byte(path.Base(file)))
385			case "line":
386				return w.Write([]byte(strconv.Itoa(line)))
387			}
388			return 0, nil
389		})
390
391		if err == nil {
392			s := buf.String()
393			i := buf.Len() - 1
394			if s[i] == '}' {
395				// JSON header
396				buf.Truncate(i)
397				buf.WriteByte(',')
398				if format == "json" {
399					buf.WriteString(message[1:])
400				} else {
401					buf.WriteString(`"message":`)
402					buf.WriteString(strconv.Quote(message))
403					buf.WriteString(`}`)
404				}
405			} else {
406				// Text header
407				buf.WriteByte(' ')
408				buf.WriteString(message)
409			}
410			buf.WriteByte('\n')
411			l.mutex.Lock()
412			defer l.mutex.Unlock()
413			l.output.Write(buf.Bytes())
414		}
415	}
416}
417