func assertUnitState(name string, js job.JobState, out io.Writer) (ret bool) {
	u, err := cAPI.Unit(name)
	if err != nil {
		log.Warningf("Error retrieving Unit(%s) from Registry: %v", name, err)
		return
	}
	if u == nil {
		log.Warningf("Unit %s not found", name)
		return
	}
	if job.JobState(u.CurrentState) != js {
		log.Debugf("Waiting for Unit(%s) state(%s) to be %s", name, job.JobState(u.CurrentState), js)
		return
	}

	ret = true
	msg := fmt.Sprintf("Unit %s %s", name, u.CurrentState)

	if u.MachineID != "" {
		ms := cachedMachineState(u.MachineID)
		if ms != nil {
			msg = fmt.Sprintf("%s on %s", msg, machineFullLegend(*ms, false))
		}
	}

	fmt.Fprintln(out, msg)
	return
}

func runRestartUnit(cCmd *cobra.Command, args []string) (exit int) {
	if len(args) == 0 {
		stderr("No units given")
		return 0
	}
	units, err := findUnits(args)
	if err != nil {
		stderr("%v", err)
		return 1
	}

	if err := lazyCreateUnits(cCmd, args); err != nil {
		stderr("Error creating units: %v", err)
		return 1
	}

	globalUnits := make([]schema.Unit, 0)
	for _, unit := range units {
		if suToGlobal(unit) {
			globalUnits = append(globalUnits, unit)
			continue
		}
		if job.JobState(unit.CurrentState) == job.JobStateInactive {
			stderr("Unable to restart unit %s in state %s", unit.Name, job.JobStateInactive)
			continue
		} else if job.JobState(unit.CurrentState) == job.JobStateLoaded {
			log.Infof("Unit(%s) already %s, starting.", unit.Name, job.JobStateLoaded)

			exit = setUnitStateAndWait(unit, job.JobStateLaunched, getBlockAttempts(cCmd))
			if exit == 1 {
				return exit
			}
			continue
		} else {
			//stop and start it
			exit = setUnitStateAndWait(unit, job.JobStateLoaded, getBlockAttempts(cCmd))
			if exit == 1 {
				return exit
			}
			exit = setUnitStateAndWait(unit, job.JobStateLaunched, getBlockAttempts(cCmd))
			if exit == 1 {
				return exit
			}
		}
		log.Infof("Unit(%s) was restarted.", unit.Name)
	}

	if err := cmdGlobalMachineState(cCmd, globalUnits); err != nil {
		stderr("Error restarting global units %v err:%v", globalUnits, err)
		return 1
	}

	return
}

func runJournal(args []string) (exit int) {
	if len(args) != 1 {
		stderr("One unit file must be provided.")
		return 1
	}
	name := unitNameMangle(args[0])

	u, err := cAPI.Unit(name)
	if err != nil {
		stderr("Error retrieving unit %s: %v", name, err)
		return 1
	} else if u == nil {
		stderr("Unit %s does not exist.", name)
		return 1
	} else if suToGlobal(*u) {
		stderr("Unable to retrieve journal of global unit %s.", name)
		return 1
	} else if job.JobState(u.CurrentState) == job.JobStateInactive {
		stderr("Unit %s does not appear to be running.", name)
		return 1
	}

	cmd := []string{"journalctl", "--unit", name, "--no-pager", "-n", strconv.Itoa(flagLines)}

	if flagSudo {
		cmd = append([]string{"sudo"}, cmd...)
	}

	if flagFollow {
		cmd = append(cmd, "-f")
	}

	return runCommand(u.MachineID, cmd[0], cmd[1:]...)
}

func runJournal(args []string) (exit int) {
	if len(args) != 1 {
		fmt.Fprintln(os.Stderr, "One unit file must be provided.")
		return 1
	}
	name := unitNameMangle(args[0])

	u, err := cAPI.Unit(name)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error retrieving unit %s: %v", name, err)
		return 1
	}
	if u == nil {
		fmt.Fprintf(os.Stderr, "Unit %s does not exist.\n", name)
		return 1
	} else if job.JobState(u.CurrentState) == job.JobStateInactive {
		fmt.Fprintf(os.Stderr, "Unit %s does not appear to be running.\n", name)
		return 1
	}

	command := fmt.Sprintf("journalctl --unit %s --no-pager -n %d", name, flagLines)
	if flagFollow {
		command += " -f"
	}

	return runCommand(command, u.MachineID)
}

// checkReplaceUnitState checks if the unit should be replaced.
// It takes a Unit object as a parameter.
// It returns 0 on success and if the unit should be replaced, 1 if the
// unit should not be replaced; and any error encountered.
func checkReplaceUnitState(unit *schema.Unit) (int, error) {
	// We replace units only for 'submit', 'load' and
	// 'start' commands.
	allowedReplace := map[string][]job.JobState{
		"submit": []job.JobState{
			job.JobStateInactive,
		},
		"load": []job.JobState{
			job.JobStateInactive,
			job.JobStateLoaded,
		},
		"start": []job.JobState{
			job.JobStateInactive,
			job.JobStateLoaded,
			job.JobStateLaunched,
		},
	}

	if allowedJobs, ok := allowedReplace[currentCommand]; ok {
		for _, j := range allowedJobs {
			if job.JobState(unit.DesiredState) == j {
				return 0, nil
			}
		}
		// Report back to caller that we are not allowed to
		// cross unit transition states
		stderr("Warning: can not replace Unit(%s) in state '%s', use the appropriate command", unit.Name, unit.DesiredState)
	} else {
		// This function should only be called from 'submit',
		// 'load' and 'start' upper paths.
		return 1, fmt.Errorf("error: replacing units is not supported in this context")
	}

	return 1, nil
}

func runStatusUnit(cCmd *cobra.Command, args []string) (exit int) {
	for i, arg := range args {
		name := unitNameMangle(arg)
		unit, err := cAPI.Unit(name)
		if err != nil {
			stderr("Error retrieving unit: %v", err)
			return 1
		}

		if unit == nil {
			stderr("Unit %s does not exist.", name)
			return 1
		} else if suToGlobal(*unit) {
			stderr("Unable to determine status of global unit %s.", unit.Name)
			return 1
		} else if job.JobState(unit.CurrentState) == job.JobStateInactive {
			stderr("Unit %s does not appear to be loaded.", unit.Name)
			return 1
		}

		// This extra newline is here to match systemctl status output
		if i != 0 {
			fmt.Printf("\n")
		}

		if exitVal := runCommand(cCmd, unit.MachineID, "systemctl", "status", "-l", unit.Name); exitVal != 0 {
			exit = exitVal
			break
		}
	}

	return
}

func runUnloadUnit(args []string) (exit int) {
	units, err := findUnits(args)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		return 1
	}

	wait := make([]string, 0)
	for _, s := range units {
		if job.JobState(s.CurrentState) == job.JobStateInactive {
			log.V(1).Infof("Unit(%s) already %s, skipping.", s.Name, job.JobStateInactive)
			continue
		}

		log.V(1).Infof("Setting target state of Unit(%s) to %s", s.Name, job.JobStateInactive)
		cAPI.SetUnitTargetState(s.Name, string(job.JobStateInactive))
		wait = append(wait, s.Name)
	}

	if !sharedFlags.NoBlock {
		errchan := waitForUnitStates(wait, job.JobStateInactive, sharedFlags.BlockAttempts, os.Stdout)
		for err := range errchan {
			fmt.Fprintf(os.Stderr, "%v\n", err)
			exit = 1
		}
	}

	return
}

func runUnloadUnit(args []string) (exit int) {
	if len(args) == 0 {
		stderr("No units given")
		return 0
	}

	units, err := findUnits(args)
	if err != nil {
		stderr("%v", err)
		return 1
	}

	wait := make([]string, 0)
	for _, s := range units {
		if !suToGlobal(s) {
			if job.JobState(s.CurrentState) == job.JobStateInactive {
				log.Debugf("Target state of Unit(%s) already %s, skipping.", s.Name, job.JobStateInactive)
				continue
			}
		}

		log.Debugf("Setting target state of Unit(%s) to %s", s.Name, job.JobStateInactive)
		cAPI.SetUnitTargetState(s.Name, string(job.JobStateInactive))
		if suToGlobal(s) {
			stdout("Triggered global unit %s unload", s.Name)
		} else {
			wait = append(wait, s.Name)
		}
	}

	exit = tryWaitForUnitStates(wait, "unload", job.JobStateInactive, getBlockAttempts(), os.Stdout)

	return
}

func runJournal(args []string) (exit int) {
	if len(args) != 1 {
		stderr("One unit file must be provided.")
		return 1
	}
	name := unitNameMangle(args[0])

	u, err := cAPI.Unit(name)
	if err != nil {
		stderr("Error retrieving unit %s: %v", name, err)
		return 1
	} else if u == nil {
		stderr("Unit %s does not exist.", name)
		return 1
	} else if suToGlobal(*u) {
		stderr("Unable to retrieve journal of global unit %s.", name)
		return 1
	} else if job.JobState(u.CurrentState) == job.JobStateInactive {
		stderr("Unit %s does not appear to be running.", name)
		return 1
	}

	command := fmt.Sprintf("journalctl --unit %s --no-pager -n %d", name, flagLines)

	if flagSudo {
		command = "sudo " + command
	}

	if flagFollow {
		command += " -f"
	}

	return runCommand(command, u.MachineID)
}

func mapUnitToJob(entity *schema.Unit, mm map[string]*machine.MachineState) (*job.Job, error) {
	contents, err := base64.StdEncoding.DecodeString(entity.FileContents)
	if err != nil {
		return nil, err
	}
	u, err := unit.NewUnit(string(contents))
	if err != nil {
		return nil, err
	}

	js := job.JobState(entity.CurrentState)
	j := job.Job{
		Name:  entity.Name,
		State: &js,
		Unit:  *u,
	}

	// populate a UnitState object only if the entity
	// is actually reporting relevant data
	if entity.Systemd != nil {
		j.UnitState = &unit.UnitState{
			LoadState:   entity.Systemd.LoadState,
			ActiveState: entity.Systemd.ActiveState,
			SubState:    entity.Systemd.SubState,
		}
		if len(entity.Systemd.MachineID) > 0 {
			j.UnitState.MachineID = entity.Systemd.MachineID
		}
	}

	return &j, nil
}

func (ur *unitsResource) set(rw http.ResponseWriter, req *http.Request, item string) {
	if err := validateContentType(req); err != nil {
		sendError(rw, http.StatusUnsupportedMediaType, err)
		return
	}

	var su schema.Unit
	dec := json.NewDecoder(req.Body)
	err := dec.Decode(&su)
	if err != nil {
		sendError(rw, http.StatusBadRequest, fmt.Errorf("unable to decode body: %v", err))
		return
	}
	if su.Name == "" {
		su.Name = item
	}
	if item != su.Name {
		sendError(rw, http.StatusBadRequest, fmt.Errorf("name in URL %q differs from unit name in request body %q", item, su.Name))
		return
	}
	if err := ValidateName(su.Name); err != nil {
		sendError(rw, http.StatusBadRequest, err)
		return
	}

	eu, err := ur.cAPI.Unit(su.Name)
	if err != nil {
		log.Errorf("Failed fetching Unit(%s) from Registry: %v", su.Name, err)
		sendError(rw, http.StatusInternalServerError, nil)
		return
	}

	if eu == nil {
		if len(su.Options) == 0 {
			err := errors.New("unit does not exist and options field empty")
			sendError(rw, http.StatusConflict, err)
		} else if err := ValidateOptions(su.Options); err != nil {
			sendError(rw, http.StatusBadRequest, err)
		} else {
			ur.create(rw, su.Name, &su)
		}
		return
	}

	if len(su.DesiredState) == 0 {
		err := errors.New("must provide DesiredState to update existing unit")
		sendError(rw, http.StatusConflict, err)
		return
	}

	un := unit.NewUnitNameInfo(su.Name)
	if un.IsTemplate() && job.JobState(su.DesiredState) != job.JobStateInactive {
		err := fmt.Errorf("cannot activate template %q", su.Name)
		sendError(rw, http.StatusBadRequest, err)
		return
	}

	ur.update(rw, su.Name, su.DesiredState)
}

func TestAgentLoadUnloadJob(t *testing.T) {
	uManager := unit.NewFakeUnitManager()
	usGenerator := unit.NewUnitStateGenerator(uManager)
	fReg := registry.NewFakeRegistry()
	mach := &machine.FakeMachine{machine.MachineState{ID: "XXX"}}
	a, err := New(uManager, usGenerator, fReg, mach, DefaultTTL)
	if err != nil {
		t.Fatalf("Failed creating Agent: %v", err)
	}

	j := newTestJobFromUnitContents(t, "foo.service", "")
	err = a.loadJob(j)
	if err != nil {
		t.Fatalf("Failed calling Agent.loadJob: %v", err)
	}

	jobs, err := a.jobs()
	if err != nil {
		t.Fatalf("Failed calling Agent.jobs: %v", err)
	}

	jsLoaded := job.JobStateLoaded
	expectJobs := map[string]*job.Job{
		"foo.service": &job.Job{
			Name: "foo.service",
			UnitState: &unit.UnitState{
				LoadState:   "loaded",
				ActiveState: "active",
				SubState:    "running",
				MachineID:   "",
			},
			State: &jsLoaded,

			Unit:            unit.Unit{},
			TargetState:     job.JobState(""),
			TargetMachineID: "",
		},
	}

	if !reflect.DeepEqual(expectJobs, jobs) {
		t.Fatalf("Received unexpected collection of Jobs: %#v\nExpected: %#v", jobs, expectJobs)
	}

	a.unloadJob("foo.service")

	// This sucks, but we have to do it if Agent.unloadJob is going to spin
	// off the real work that matters in a goroutine
	time.Sleep(200)

	jobs, err = a.jobs()
	if err != nil {
		t.Fatalf("Failed calling Agent.jobs: %v", err)
	}

	expectJobs = map[string]*job.Job{}
	if !reflect.DeepEqual(expectJobs, jobs) {
		t.Fatalf("Received unexpected collection of Jobs: %#v\nExpected: %#v", jobs, expectJobs)
	}
}

func MapSchemaUnitToScheduledUnit(entity *Unit) *job.ScheduledUnit {
	cs := job.JobState(entity.CurrentState)
	return &job.ScheduledUnit{
		Name:            entity.Name,
		State:           &cs,
		TargetMachineID: entity.MachineID,
	}
}

func runStopUnit(args []string) (exit int) {
	units, err := findUnits(args)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		return 1
	}

	stopping := make([]string, 0)
	for _, u := range units {
		if !suToGlobal(u) {
			if job.JobState(u.CurrentState) == job.JobStateInactive {
				fmt.Fprintf(os.Stderr, "Unable to stop unit %s in state %s\n", u.Name, job.JobStateInactive)
				return 1
			} else if job.JobState(u.CurrentState) == job.JobStateLoaded {
				log.V(1).Infof("Unit(%s) already %s, skipping.", u.Name, job.JobStateLoaded)
				continue
			}
		}

		log.V(1).Infof("Setting target state of Unit(%s) to %s", u.Name, job.JobStateLoaded)
		cAPI.SetUnitTargetState(u.Name, string(job.JobStateLoaded))
		if suToGlobal(u) {
			fmt.Printf("Triggered global unit %s stop\n", u.Name)
		} else {
			stopping = append(stopping, u.Name)
		}
	}

	if !sharedFlags.NoBlock {
		errchan := waitForUnitStates(stopping, job.JobStateLoaded, sharedFlags.BlockAttempts, os.Stdout)
		for err := range errchan {
			fmt.Fprintf(os.Stderr, "Error waiting for units: %v\n", err)
			exit = 1
		}
	} else {
		for _, name := range stopping {
			fmt.Printf("Triggered unit %s stop\n", name)
		}
	}

	return
}

func assertUnitState(name string, js job.JobState, out io.Writer) bool {
	fetchUnitState := func() error {
		var state string

		u, err := cAPI.Unit(name)
		if err != nil {
			return fmt.Errorf("Error retrieving Unit(%s) from Registry: %v", name, err)
		}
		if u == nil {
			return fmt.Errorf("Unit %s not found", name)
		}

		// If this is a global unit, CurrentState will never be set. Instead, wait for DesiredState.
		if suToGlobal(*u) {
			state = u.DesiredState
		} else {
			state = u.CurrentState
		}

		if job.JobState(state) != js {
			return fmt.Errorf("Waiting for Unit(%s) state(%s) to be %s", name, job.JobState(state), js)
		}

		msg := fmt.Sprintf("Unit %s %s", name, u.CurrentState)

		if u.MachineID != "" {
			ms := cachedMachineState(u.MachineID)
			if ms != nil {
				msg = fmt.Sprintf("%s on %s", msg, machineFullLegend(*ms, false))
			}
		}

		fmt.Fprintln(out, msg)
		return nil
	}
	_, err := waitForState(fetchUnitState)
	if err != nil {
		return false
	}

	return true
}

func checkLoadUnitState(unit schema.Unit, loadRet int, errchan chan error) {
	if loadRet == 0 {
		if job.JobState(unit.DesiredState) != job.JobStateLoaded {
			errchan <- fmt.Errorf("Error: unit %s was not loaded as requested", unit.Name)
		}
	} else if unit.DesiredState != "" {
		// if the whole load operation failed, then no unit
		// should have a DesiredState set
		errchan <- fmt.Errorf("Error: Unit(%s) DesiredState was set to (%s)", unit.Name, unit.DesiredState)
	}
}

func (f *FleetTunnel) Stop(events chan Event, unitNames ...string) (StopStats, error) {
	log.Debugf("stopping %v", unitNames)

	units, err := f.findUnits(unitNames)
	if err != nil {
		return StopStats{}, maskAny(err)
	}

	stopping := make([]string, 0)
	stats := StopStats{}
	for _, u := range units {
		if !suToGlobal(u) {
			if job.JobState(u.CurrentState) == job.JobStateInactive {
				return StopStats{}, maskAny(fmt.Errorf("Unable to stop unit %s in state %s", u.Name, job.JobStateInactive))
			} else if job.JobState(u.CurrentState) == job.JobStateLoaded {
				log.Debugf("Unit(%s) already %s, skipping.", u.Name, job.JobStateLoaded)
				continue
			}
		}

		log.Debugf("Setting target state of Unit(%s) to %s", u.Name, job.JobStateLoaded)
		if err := f.setUnitTargetStateWithRetry(u.Name, string(job.JobStateLoaded)); err != nil {
			return StopStats{}, maskAny(err)
		}
		if suToGlobal(u) {
			stats.StoppedGlobalUnits++
			events <- newEvent(u.Name, "triggered global unit stop")
		} else {
			stats.StoppedUnits++
			events <- newEvent(u.Name, "triggered unit stop")
			stopping = append(stopping, u.Name)
		}
	}

	if err := f.tryWaitForUnitStates(stopping, "stop", job.JobStateLoaded, f.BlockAttempts, events); err != nil {
		return StopStats{}, maskAny(err)
	}

	return stats, nil
}

func (f *FleetTunnel) assertUnitState(name string, js job.JobState, events chan Event) (ret bool) {
	var state string

	u, err := f.unitWithRetry(name)
	if err != nil {
		log.Warningf("Error retrieving Unit(%s) from Registry: %v", name, err)
		return
	}
	if u == nil {
		events <- newEvent(name, "unit not found")
		log.Warningf("Unit %s not found", name)
		return
	}

	// If this is a global unit, CurrentState will never be set. Instead, wait for DesiredState.
	if suToGlobal(*u) {
		state = u.DesiredState
	} else {
		state = u.CurrentState
	}

	if job.JobState(state) != js {
		log.Debugf("Waiting for Unit(%s) state(%s) to be %s", name, job.JobState(state), js)
		return
	}

	ret = true
	msg := u.CurrentState

	if u.MachineID != "" {
		ms := f.cachedMachineState(u.MachineID)
		if ms != nil {
			msg = fmt.Sprintf("%s on %s", msg, machineFullLegend(*ms, false))
		}
	}

	events <- newEvent(name, msg)
	return
}

// jobs returns a collection of all Jobs that the Agent has either loaded
// or launched. The Unit, TargetState and TargetMachineID fields of the
// returned *job.Job objects are not properly hydrated.
func (a *Agent) jobs() (map[string]*job.Job, error) {
	launched := pkg.NewUnsafeSet()
	for _, jName := range a.cache.launchedJobs() {
		launched.Add(jName)
	}

	loaded := pkg.NewUnsafeSet()
	for _, jName := range a.cache.loadedJobs() {
		loaded.Add(jName)
	}

	units, err := a.um.Units()
	if err != nil {
		return nil, fmt.Errorf("failed fetching loaded units from UnitManager: %v", err)
	}

	filter := pkg.NewUnsafeSet()
	for _, u := range units {
		filter.Add(u)
	}
	states, err := a.um.GetUnitStates(filter)
	if err != nil {
		return nil, fmt.Errorf("failed fetching unit states: %v", err)
	}

	jobs := make(map[string]*job.Job)
	for _, uName := range units {
		jobs[uName] = &job.Job{
			Name:      uName,
			UnitState: states[uName],
			State:     nil,

			// The following fields are not properly populated
			// and should not be used in the calling code
			Unit:            unit.Unit{},
			TargetState:     job.JobState(""),
			TargetMachineID: "",
		}

		js := job.JobStateInactive
		if loaded.Contains(uName) {
			js = job.JobStateLoaded
		} else if launched.Contains(uName) {
			js = job.JobStateLaunched
		}
		jobs[uName].State = &js
	}

	return jobs, nil
}

func printUnitStatus(name string) int {
	u, err := cAPI.Unit(name)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error retrieving unit %s: %v", name, err)
		return 1
	}
	if u == nil {
		fmt.Fprintf(os.Stderr, "Unit %s does not exist.\n", name)
		return 1
	} else if job.JobState(u.CurrentState) == job.JobStateInactive {
		fmt.Fprintf(os.Stderr, "Unit %s does not appear to be running.\n", name)
		return 1
	}

	cmd := fmt.Sprintf("systemctl status -l %s", name)
	return runCommand(cmd, u.MachineID)
}