// AddAgentUpdateHandlers adds the needed update handlers to perform agent
// updates
func AddAgentUpdateHandlers(cs wsclient.ClientServer, cfg *config.Config, saver statemanager.Saver, taskEngine engine.TaskEngine) {
	singleUpdater = &updater{
		acs:        cs,
		config:     cfg,
		fs:         os.Default,
		httpclient: httpclient.New(updateDownloadTimeout, false),
	}
	cs.AddRequestHandler(singleUpdater.stageUpdateHandler())
	cs.AddRequestHandler(singleUpdater.performUpdateHandler(saver, taskEngine))
}

func ackMessageId(cs wsclient.ClientServer, cluster, containerInstanceArn, messageID string) {
	err := cs.MakeRequest(&ecsacs.AckRequest{
		Cluster:           &cluster,
		ContainerInstance: &containerInstanceArn,
		MessageId:         &messageID,
	})
	if err != nil {
		log.Warn("Error 'ack'ing request", "MessageID", messageID)
	}
}

// newDisconnectionTimer creates a new time object, with a callback to
// disconnect from ACS on inactivity
func newDisconnectionTimer(client wsclient.ClientServer, _time ttime.Time, timeout time.Duration, jitter time.Duration) ttime.Timer {
	timer := _time.AfterFunc(utils.AddJitter(timeout, jitter), func() {
		seelog.Warn("ACS Connection hasn't had any activity for too long; closing connection")
		closeErr := client.Close()
		if closeErr != nil {
			seelog.Warnf("Error disconnecting: %v", closeErr)
		}
	})

	return timer
}

// AddAgentUpdateHandlers adds the needed update handlers to perform agent
// updates
func AddAgentUpdateHandlers(cs wsclient.ClientServer, cfg *config.Config, saver statemanager.Saver, taskEngine engine.TaskEngine) {
	if cfg.UpdatesEnabled {
		singleUpdater = &updater{
			acs:        cs,
			config:     cfg,
			fs:         os.Default,
			httpclient: httpclient.New(updateDownloadTimeout, false),
		}
		cs.AddRequestHandler(singleUpdater.stageUpdateHandler())
		cs.AddRequestHandler(singleUpdater.performUpdateHandler(saver, taskEngine))
		log.Debug("Added update handlers")
	} else {
		log.Debug("Updates disabled; no handlers added")
	}
}

// startACSSession starts a session with ACS. It adds request handlers for various
// kinds of messages expected from ACS. It returns on server disconnection or when
// the context is cancelled
func startACSSession(ctx context.Context, client wsclient.ClientServer, timer ttime.Timer, args StartSessionArguments, backoff *utils.SimpleBackoff, acsSessionState sessionState) error {
	// Any message from the server resets the disconnect timeout
	client.SetAnyRequestHandler(anyMessageHandler(timer))
	cfg := args.Config

	refreshCredsHandler := newRefreshCredentialsHandler(ctx, cfg.Cluster, args.ContainerInstanceArn, client, args.CredentialsManager, args.TaskEngine)
	defer refreshCredsHandler.clearAcks()
	refreshCredsHandler.start()
	client.AddRequestHandler(refreshCredsHandler.handlerFunc())

	// Add request handler for handling payload messages from ACS
	payloadHandler := newPayloadRequestHandler(ctx, args.TaskEngine, args.ECSClient, cfg.Cluster, args.ContainerInstanceArn, client, args.StateManager, refreshCredsHandler, args.CredentialsManager)
	// Clear the acks channel on return because acks of messageids don't have any value across sessions
	defer payloadHandler.clearAcks()
	payloadHandler.start()
	client.AddRequestHandler(payloadHandler.handlerFunc())

	// Ignore heartbeat messages; anyMessageHandler gets 'em
	client.AddRequestHandler(func(*ecsacs.HeartbeatMessage) {})

	updater.AddAgentUpdateHandlers(client, cfg, args.StateManager, args.TaskEngine)

	err := client.Connect()
	if err != nil {
		seelog.Errorf("Error connecting to ACS: %v", err)
		return err
	}
	acsSessionState.connectedToACS()

	backoffResetTimer := args.time().AfterFunc(utils.AddJitter(args.heartbeatTimeout(), args.heartbeatJitter()), func() {
		// If we do not have an error connecting and remain connected for at
		// least 5 or so minutes, reset the backoff. This prevents disconnect
		// errors that only happen infrequently from damaging the
		// reconnectability as significantly.
		backoff.Reset()
	})
	defer backoffResetTimer.Stop()

	serveErr := make(chan error, 1)
	go func() {
		serveErr <- client.Serve()
	}()

	for {
		select {
		case <-ctx.Done():
			// Stop receiving and sending messages from and to ACS when
			// the context received from the main function is canceled
			payloadHandler.stop()
			refreshCredsHandler.stop()
			return ctx.Err()
		case err := <-serveErr:
			// Stop receiving and sending messages from and to ACS when
			// client.Serve returns an error. This can happen when the
			// the connection is closed by ACS or the agent
			payloadHandler.stop()
			refreshCredsHandler.stop()
			return err
		}
	}
}