package agent

import (
	"fmt"
	"os"
	"path"
	"sync"
	"time"

	"code.justin.tv/systems/sandstorm/manager"
	"code.justin.tv/systems/sandstorm/queue"

	log "github.com/Sirupsen/logrus"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/cenkalti/backoff"
)

// Mutex ix the global mutex for when we want to reload agent
var Mutex sync.Mutex

// used to implement blocking on startup errors. receive from this channel
// to unblock the startup process.
var blocker chan struct{}

func init() {
	blocker = make(chan struct{})
}

// Blocker blocks the agent
type Blocker interface {
	Block()
	Unblock()
}

type agentBlocker struct{}

// Block blocks until something calls unblock()
func (a agentBlocker) Block() {
	log.Debug("blocking agent")
	blocker <- struct{}{}
}

// Unblock unblocks the agent
func (a agentBlocker) Unblock() {
	select {
	case <-blocker:
		log.Debug("unblocking blocked agent")
	default:
		log.Debug("unblocking unblocked agent")
	}
}

// Agent tracks the state and configuration of the daemon
type Agent struct {
	TopicARN         string
	Dwell            int
	Splay            int
	Queue            *queue.Queue
	DefaultQueue     *queue.Queue
	ProxyQueue       *queue.Queue
	State            *State
	KeyID            string
	manager          manager.API
	defaultManager   manager.API
	proxyManager     manager.API
	Region           string
	TableName        string
	awsConfig        *aws.Config
	RoleARN          string
	Proxy            string
	useProxy         bool
	Pop              string
	Logging          bool
	Role             string
	configFolderPath string
	Port             int
	// TemplateWhitelist, if not empty, will only allow matching template source
	// file paths to be used by the agent for a given run.
	TemplateWhitelist      []string
	BackoffMultiplier      int
	MaxWaitTime            int
	stopPoll               chan bool
	blocker                Blocker
	sandstormAccountNumber string

	newBackoff func() backoff.BackOff
}

// New returns an initialized but unconfigured agent
func New() *Agent {
	return &Agent{
		State: &State{
			LastUpdate:           []*manager.Secret{},
			LastSuccessfulUpdate: []*manager.Secret{},
			Templates:            []*Template{},
			BlockedSecrets:       make(map[string]string),
		},
		TemplateWhitelist: []string{},
		stopPoll:          make(chan bool, 1),
		blocker:           agentBlocker{},
		newBackoff:        func() backoff.BackOff { return NewExponentialBackoff() },
	}
}

// Start uses the Agent's initial state to perform an initial template
// run.
func (a *Agent) Start(disableRestart, blockOnMissingSecret bool) error {
	var secretNames []string
	if disableRestart && len(a.TemplateWhitelist) > 0 {
		templates := []*Template{}
		for _, template := range a.State.Templates {
			for _, whitelistTemplate := range a.TemplateWhitelist {
				if whitelistTemplate == template.Source {
					templates = append(templates, template)
				}
			}

		}
		secretNames = []string{}
		for _, template := range templates {
			for secretName := range template.Secrets {
				secretNames = append(secretNames, secretName)
			}
		}

	} else {
		secretNames = a.GetAllSecretNames()

	}
	return a.sync(secretNames, disableRestart, blockOnMissingSecret, a.Splay)
}

// Simulate outputs all templates, replacing all values with defaultTestSecretPlaintext instead
// of real secret values. It does not touch any part of sandstorm
func (a *Agent) Simulate() (err error) {
	effectedTemplates := a.getEffectedTemplates(a.GetAllSecrets())
	_, err = outputTemplates(effectedTemplates)
	return
}

// Watch observes an SNS Queue for changed secrets and updates it's templates
// accordingly if one of it's values change.
func (a *Agent) Watch() {
	if len(a.TemplateWhitelist) > 0 {
		log.Warning("attempting to use template whitelisting in watch mode, ignoring template whitelist")
		a.TemplateWhitelist = []string{}
	}

	// check for state file, if one doesn't exist do an intial run and don't run commands
	var err error
	if !fileExists(path.Join(a.configFolderPath, StateFileName)) {
		err = a.Start(true, true)
		if err != nil {
			log.Fatal(err.Error())
		}
	} else {
		// else do initial run with commands
		err = a.Start(false, true)
		if err != nil {
			log.Fatal(err.Error())
		}
	}

	for {
		err := a.runOnUpdate()
		// log errors, but do not block
		if err != nil {
			a.State.handleError(err)
		}
		// Allow time for HUP Lock to run, avoid race
		time.Sleep(time.Second)
	}
}

// runOnUpdate polls SQS for secret updates, and then performs a template run for
// the newly updated state
func (a *Agent) runOnUpdate() error {
	// Read new secrets to global state
	Mutex.Lock()
	secretNames := a.pollForSecrets(a.GetAllSecretNames())
	if len(secretNames) > 0 {
		// Compare and update against templates
		err := a.sync(secretNames, false, true, a.Splay)
		if err != nil {
			Mutex.Unlock()
			log.Error(err.Error())
			return err
		}
	}
	Mutex.Unlock()
	return nil
}

// sync takes a list of secrets as input, fetches the most recent versions of them,
// and if any changes have occurred, synchronized any dependent templates with
// the updated states.
func (a *Agent) sync(secretNames []string, disableRestart, blockOnMissingSecret bool, splay int) error {
	var err error
	var newSecrets []*manager.Secret
	newSecrets, err = a.FetchSecretsByName(secretNames)
	if err != nil {
		a.State.SetRecentError(err)
		if e, ok := err.(MissingSecretError); ok {
			if !blockOnMissingSecret {
				return err
			}
			log.Error(e.Error())
		} else {
			return fmt.Errorf("error fetching secrets: %s", err.Error())
		}
	}

	// if we have missing secrets, we should keep trying to fetch them and log
	// that we're doing so until we resolve all of the issues.
	if blockOnMissingSecret && len(a.State.BlockedSecrets) > 0 {
		log.Infof("listening for SQS messages for %d missing secret(s)", len(a.State.BlockedSecrets))
		for len(a.State.BlockedSecrets) > 0 {
			log.Debugf("blocked by these secrets: %+v", a.State.BlockedSecrets)
			updatedSecrets := a.pollForSecrets(secretNames)
			for _, s := range updatedSecrets {
				if _, ok := a.State.BlockedSecrets[s]; ok {
					log.Infof("received update for blocked secret %s, removing from BlockedSecrets", s)
					delete(a.State.BlockedSecrets, s)
				}
			}
		}
		log.Info("received updates for all blocked secrets, running sync")
		return a.sync(secretNames, disableRestart, blockOnMissingSecret, splay)
	}

	// Update health check
	a.State.LastUpdate = newSecrets

	effectedTemplates := a.getEffectedTemplates(newSecrets)

	// update Agent state secrets/templates with changes
	a.updateState(effectedTemplates, newSecrets)

	dirtyTemplates := []*Template{}
	for _, template := range a.State.Templates {
		if template.Dirty {
			if disableRestart && !a.isTemplateWhitelisted(template) {
				continue
			}
			dirtyTemplates = append(dirtyTemplates, template)
		}
	}

	dirtyTemplates, err = outputTemplates(dirtyTemplates)
	if err != nil {
		if e, ok := err.(*TemplatingError); ok {
			log.Error(e.Error())
			err = nil
		} else {
			return err
		}
	}

	if !disableRestart {
		err = runRestartCommands(dirtyTemplates, splay)
		if err != nil {
			return err
		}
	}

	if blockOnMissingSecret {
		stateFilePath := path.Join(a.configFolderPath, StateFileName)
		err = a.writeStateFile(stateFilePath)
		if err != nil {
			log.Errorf("error writing statefile: %s", err)
			err = nil
		}
	}

	// If we made it through the gauntlet, update health check last success
	a.State.LastSuccessfulUpdate = newSecrets
	a.State.SetRecentError(nil)
	return nil
}

func (a *Agent) getEffectedTemplates(secrets []*manager.Secret) []*Template {
	templates := []*Template{}

	for _, template := range a.State.Templates {
		if template.ContainsSecrets(secrets) && a.isTemplateWhitelisted(template) {
			templates = append(templates, template)
		}
	}

	return templates
}

// updateState compares the input list of secrets to that of the agent's
// current state. If a newer version is provided, it synchronizes each provided
// template with it. Finally it returns a list of all secrets that were deemed to be new
// versions so we can track the updated state.
func (a *Agent) updateState(effectedTemplates []*Template, newSecrets []*manager.Secret) {

	for _, newSecret := range newSecrets {

		// Templates should be updated in two cases:
		// 1) The secret is an updated version.
		// 2) This is the initialization run. In this case, when secrets are
		// synchronized from a state file, they will be tracked, but won't have
		// a Plaintext value, so we should add it.
		for _, template := range effectedTemplates {
			if _, containsSecret := template.Secrets[newSecret.Name]; containsSecret {
				template.UpsertSecret(newSecret)
				log.Infof("Updated template '%s'.", template.Source)
			}
		}

		log.Infof("Updated secret '%s'.", newSecret.Name)
	}
}

// FetchSecretsByName takes a list of secret names and fetches the most recent
// version of each from KMS.
func (a *Agent) FetchSecretsByName(secretNames []string) ([]*manager.Secret, error) {

	secrets := make([]*manager.Secret, len(secretNames))

	// Now finally update all changed secrets
	for i, secretName := range secretNames {

		// checks local cache, retrieves using KMS if not found, allows
		// us to use this logic for both Watch and Run since a single shot
		// Run should be the only time fresh secrets are fetched from Dynamo
		newSecret, err := a.FetchSecret(secretName)
		if err != nil {
			return nil, err
		}

		secrets[i] = newSecret
	}

	return secrets, nil
}

// FetchSecret is a wrapper around manager.Get that does a random
// exponential backoff. Each secret fetch is guaranteed to  return
// within MaxWaitTime seconds. After each failed attempt, the agent will wait
// rand(2^n) seconds where n = attempt number before trying again.
func (a *Agent) FetchSecret(secretName string) (fetchedSecret *manager.Secret, err error) {

	getSecret := func() (getSecretErr error) {
		fetchedSecret, getSecretErr = a.manager.Get(secretName)
		if getSecretErr != nil {
			getSecretErr = SecretFetchError{
				Secret: secretName,
				Err:    getSecretErr,
			}
			if isSystemError(getSecretErr) {
				a.setProxyUse(true)
			}
			err = getSecretErr
			log.Error(getSecretErr)
			return
		}

		if fetchedSecret == nil {
			err = MissingSecretError{secret: secretName}
			a.State.BlockedSecrets[secretName] = err.Error()
			return
		}

		delete(a.State.BlockedSecrets, secretName)
		return
	}

	retryErr := backoff.Retry(getSecret, a.newBackoff())
	if retryErr != nil {
		err = retryErr
		return
	}
	return
}

// pollForSecrets watches an AWS Queue for any changed secrets over a period of time
// and returns a list of secret names that have since been updated
func (a *Agent) pollForSecrets(whitelistedNames []string) []string {

	whitelistMap := map[string]struct{}{}
	for _, name := range whitelistedNames {
		whitelistMap[name] = struct{}{}
	}

	// Using a map essentially as a set to avoid adding duplicate secrets
	uniqueSecrets := make(map[string]struct{})

	log.Debug("starting to poll for secrets")

	// for length of dwell, watch queue for any secrets that are changed
	for dwell := 0; dwell < a.Dwell; dwell++ {

		// keep looping until all secrets are processed, then dwell
		for {
			select {
			case <-a.stopPoll:
				return []string{}
			default:
			}

			// if no templates are loaded in agent, dont bother listening for
			// updates since there are no secrets
			if len(a.State.Templates) == 0 {
				break
			}
			queueSecret, err := a.Queue.PopQueueSecret()
			if err != nil {
				log.Errorf("error checking SQS for messages: %s", err.Error())
				if isSystemError(err) && a.useProxy {
					a.setProxyUse(true)
				} else {
					break
				}
			}
			if queueSecret == nil {
				// no secret got popped
				break
			}

			// only proceed if we actually popped something
			// skip over any secrets that templates aren't dependent on
			secretName := queueSecret.Name.S
			_, whitelisted := whitelistMap[secretName]

			if whitelisted {
				// reset dwell since we've received a new secret
				dwell = 0
				uniqueSecrets[secretName] = struct{}{}
				log.Infof("dependent secret '%s' changed, resetting dwell to %d", secretName, a.Dwell-dwell)
			}
		}

		// sleep for one second at a time so we don't spam SQS so badly
		time.Sleep(time.Second)
	}

	secretNames := []string{}
	for name := range uniqueSecrets {
		secretNames = append(secretNames, name)
	}

	return secretNames
}

// AddTemplate attempts to add a new template to be tracked by the agent. Before
// approving it performs a few checks to make sure the template is valid.
func (a *Agent) AddTemplate(template *Template) error {

	if template.Source == template.Destination {
		return fmt.Errorf("template source: %s is the same as template destination: %s",
			template.Source, template.Destination)
	}

	for _, existingTemplate := range a.State.Templates {
		// Check for duplicate sources
		if template.Source == existingTemplate.Source {
			return fmt.Errorf(
				"cannot have two templates with the same source: %s",
				template.Source,
			)
		}

		// Check for duplicate destinations
		if template.Destination == existingTemplate.Destination {
			return fmt.Errorf(
				"cannot have two templates with the same destination: %s",
				template.Destination,
			)
		}
	}

	a.State.Templates = append(a.State.Templates, template)

	return nil
}

/*
Checks that the specified template has been whitelisted by the agent for use. Returns
true if the whitelist is empty, or it's source is contained in the list.
*/
func (a *Agent) isTemplateWhitelisted(template *Template) bool {

	if len(a.TemplateWhitelist) == 0 {
		return true
	}

	for _, filter := range a.TemplateWhitelist {
		if template.Source == filter {
			return true
		}
	}

	return false
}

// GetAllSecretNames returns a list of secret names based on those in the current
// agent state.
func (a *Agent) GetAllSecretNames() []string {
	uniqueSecrets := map[string]struct{}{}

	// find unique secrets
	for _, template := range a.State.Templates {
		for _, secret := range template.Secrets {
			uniqueSecrets[secret.Name] = struct{}{}
		}
	}

	// convert to a slice
	secrets := []string{}
	for name := range uniqueSecrets {
		secrets = append(secrets, name)
	}

	return secrets
}

// GetAllSecrets currently tracked within templates, doesn't current account for
// different secret versions that might occur in run mode with template whitelisting.
func (a *Agent) GetAllSecrets() []*manager.Secret {
	uniqueSecrets := map[string]*manager.Secret{}

	// first find uniques
	for _, template := range a.State.Templates {
		for _, secret := range template.Secrets {
			uniqueSecrets[secret.Name] = secret
		}
	}

	secrets := []*manager.Secret{}
	for _, secret := range uniqueSecrets {
		secrets = append(secrets, secret)
	}

	return secrets
}

// writeStateFile writes the current state of the agent to disk using the current
// Agent template state's
func (a *Agent) writeStateFile(path string) error {
	stateFile := newStateFile(a.State.Templates)
	return stateFile.ToFile(path)
}

// HUPSignalHandler handles os.SIGHUP and tells the agent to rebaseTemplates
// and then calls sync with splay = 0. This process marks all templates as dirty
// and causes all templates to be outputted and restart commands to be
// immediately run. In case of an error, we gracefully exit and get
// restarted by monit.
func (a *Agent) HUPSignalHandler(c chan os.Signal, pidPath string) {
	for sig := range c {
		log.Info("Received signal: " + sig.String())
		// send signal to stop pollForSecrets
		a.stopPoll <- true
		Mutex.Lock()
		// reload agent
		newAgent, err := Build(a.configFolderPath, false, false)
		if err != nil {
			log.Error(err)
			ExitandCleanUp(pidPath)
		}
		log.Info("Reloaded Agent")
		newSecretsArr, err := a.rebaseTemplates(newAgent.State.Templates)
		if err != nil {
			log.Error(err)
			ExitandCleanUp(pidPath)
		}
		*a = *newAgent
		err = a.sync(newSecretsArr, false, false, 0)
		if err != nil {
			log.Error(err)
			if _, ok := err.(MissingSecretError); !ok {
				ExitandCleanUp(pidPath)
			}
		}
		stateFilePath := path.Join(a.configFolderPath, StateFileName)
		err = a.writeStateFile(stateFilePath)
		if err != nil {
			log.Errorf("error writing statefile: %s", err)
			err = nil
		}
		Mutex.Unlock()

		a.blocker.Unblock()
	}
}

// rebaseTemplates copies all old secrets in memory in a.State.Templates
// that are still relied on to replacementTemplate, and calculates any new (never seen before) secrets it needs to fetch.
func (a *Agent) rebaseTemplates(replacementTemplates []*Template) ([]string, error) {
	// Check for new secrets, copy old secrets from memory to new set of templates otherwise
	newSecretsSet := map[string]bool{}
	for _, rTemplate := range replacementTemplates {
		for _, oldTemplate := range a.State.Templates {
			if oldTemplate.Source == rTemplate.Source {
				if rTemplate.IsEqual(oldTemplate) {
					*rTemplate = *oldTemplate
					// We need to set it to dirty to mark it to be templated out and for the restart
					// command to be run.
					rTemplate.Dirty = true
				} else {
					rTemplate.synchronizeSecrets(oldTemplate, newSecretsSet)
				}
			}
		}
		for _, secret := range rTemplate.Secrets {
			if len(secret.Plaintext) == 0 {
				rTemplate.Dirty = true
				newSecretsSet[secret.Name] = true
			}
		}
	}
	// Set to Array
	newSecretsArr := make([]string, 0, len(newSecretsSet))
	for secretName := range newSecretsSet {
		newSecretsArr = append(newSecretsArr, secretName)
	}
	return newSecretsArr, nil
}

// outputTemplates attempts to populate template secrets and then output each
// template to it's output destination.
// If an error is encountered, it is tracked, but then continues attempting to output
// templates so that as many templates as possible make it into a valid state.
func outputTemplates(templates []*Template) (successfulTemplates []*Template, err error) {
	successfulTemplates = make([]*Template, 0, len(templates))
	outputErrors := NewTemplatingError()

	// Output all templates
	for _, template := range templates {

		outputErr := template.Output()
		if outputErr != nil {
			outputErrors.Add(template, fmt.Errorf("Error outputting '%s' template: %s", template.Source, outputErr.Error()))
			continue
		}

		successfulTemplates = append(successfulTemplates, template)
		template.Dirty = false
	}

	if outputErrors.Count() > 0 {
		err = outputErrors
		return
	}

	return
}

func runRestartCommands(templates []*Template, splay int) error {

	// track commands here so we don't run a duplicate one
	executedCommands := make(map[string]struct{})
	failedCommands := make(map[string]struct{})

	restartErrors := NewTemplatingError()

	for _, template := range templates {
		if _, alreadyRun := executedCommands[template.Command]; !alreadyRun {

			err := template.AttemptCommand(splay)
			if err != nil {
				failedCommands[template.Command] = struct{}{}
				restartErrors.Add(
					template,
					fmt.Errorf("Restart command execution failed for: %s with: %s", template.Command, err.Error()))
				continue
			}

			log.Infof("Successfully ran command: %s ", template.Command)

			executedCommands[template.Command] = struct{}{}
		} else {
			if _, failed := failedCommands[template.Command]; !failed {
				template.CommandStatus.Status = statusSuccess
			}
		}
	}

	if restartErrors.Count() > 0 {
		return restartErrors
	}

	return nil
}
