package zapcloudwatch

import (
	"encoding/json"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/cloudwatchlogs"
	"go.uber.org/zap/zapcore"
)

//LogMsg structure for messages in cw
type LogMsg struct {
	Level  string `json:"level"`
	Caller string `json:"caller"`
	Msg    string `json:"msg"`
}

// CloudwatchHook is a zap Hook for dispatching messages to the specified
type CloudwatchHook struct {
	// Messages with a log level not contained in this array
	// will not be dispatched. If nil, all messages will be dispatched.
	AcceptedLevels    []zapcore.Level
	GroupName         string
	StreamName        string
	AWSSession        *session.Session
	nextSequenceToken *string
	svc               *cloudwatchlogs.CloudWatchLogs
	// Remove async for now, TODO - research how to implement it
	//Async             bool // if async is true, send a message asynchronously.
	m sync.Mutex
}

// NewCloudwatchHook creates a new zap hook for cloudwatch
func NewCloudwatchHook(groupName, streamName string, session *session.Session) *CloudwatchHook {
	return &CloudwatchHook{
		GroupName:      groupName,
		StreamName:     streamName,
		AWSSession:     session,
		AcceptedLevels: AllLevels,
	}
}

// PrepareCWConfig checks whether log group and log stream exists, create if not
func (ch *CloudwatchHook) PrepareCWConfig() error {
	ch.svc = cloudwatchlogs.New(ch.AWSSession)

	lgresp, err := ch.svc.DescribeLogGroups(&cloudwatchlogs.DescribeLogGroupsInput{LogGroupNamePrefix: aws.String(ch.GroupName), Limit: aws.Int64(1)})
	if err != nil {
		return err
	}

	if len(lgresp.LogGroups) < 1 {
		// we need to create this log group
		_, err := ch.svc.CreateLogGroup(&cloudwatchlogs.CreateLogGroupInput{LogGroupName: aws.String(ch.GroupName)})
		if err != nil {
			return err
		}
	}

	resp, err := ch.svc.DescribeLogStreams(&cloudwatchlogs.DescribeLogStreamsInput{
		LogGroupName:        aws.String(ch.GroupName), // Required
		LogStreamNamePrefix: aws.String(ch.StreamName),
	})
	if err != nil {
		return err
	}

	// grab the next sequence token
	if len(resp.LogStreams) > 0 {
		ch.nextSequenceToken = resp.LogStreams[0].UploadSequenceToken
		return nil
	}

	// create stream if it doesn't exist. the next sequence token will be null
	_, err = ch.svc.CreateLogStream(&cloudwatchlogs.CreateLogStreamInput{
		LogGroupName:  aws.String(ch.GroupName),
		LogStreamName: aws.String(ch.StreamName),
	})

	if err != nil {
		return err
	}
	return nil
}

// Hook function returns hook to zap
func (ch *CloudwatchHook) Hook() (func(zapcore.Entry) error, error) {
	err := ch.PrepareCWConfig()
	if err != nil {
		return nil, err
	}

	var cloudwatchWriter = func(e zapcore.Entry) error {
		if !ch.isAcceptedLevel(e.Level) {
			return nil
		}

		msg := LogMsg{Level: e.Level.String(), Caller: e.Caller.String(), Msg: e.Message}
		marshalledMsg, err := json.Marshal(msg)
		if err != nil {
			return err
		}
		marshalledMsgStr := string(marshalledMsg)
		event := &cloudwatchlogs.InputLogEvent{
			Message:   aws.String(marshalledMsgStr),
			Timestamp: aws.Int64(int64(time.Nanosecond) * time.Now().UnixNano() / int64(time.Millisecond)),
		}
		params := &cloudwatchlogs.PutLogEventsInput{
			LogEvents:     []*cloudwatchlogs.InputLogEvent{event},
			LogGroupName:  aws.String(ch.GroupName),
			LogStreamName: aws.String(ch.StreamName),
			SequenceToken: ch.nextSequenceToken,
		}

		return ch.sendEvent(params)
	}
	return cloudwatchWriter, nil
}

func (ch *CloudwatchHook) sendEvent(params *cloudwatchlogs.PutLogEventsInput) error {
	resp, err := ch.svc.PutLogEvents(params)
	if err != nil {
		return err
	}
	ch.nextSequenceToken = resp.NextSequenceToken
	return nil
}

// Levels sets which levels to sent to cloudwatch
func (ch *CloudwatchHook) Levels() []zapcore.Level {
	if ch.AcceptedLevels == nil {
		return AllLevels
	}
	return ch.AcceptedLevels
}

func (ch *CloudwatchHook) isAcceptedLevel(level zapcore.Level) bool {
	for _, lv := range ch.Levels() {
		if lv == level {
			return true
		}
	}
	return false
}

// AllLevels Supported log levels
var AllLevels = []zapcore.Level{
	zapcore.DebugLevel,
	zapcore.InfoLevel,
	zapcore.WarnLevel,
	zapcore.ErrorLevel,
	zapcore.FatalLevel,
	zapcore.PanicLevel,
}

// LevelThreshold - Returns every logging level above and including the given parameter.
func LevelThreshold(l zapcore.Level) []zapcore.Level {
	for i := range AllLevels {
		if AllLevels[i] == l {
			return AllLevels[i:]
		}
	}
	return []zapcore.Level{}
}
