// newPublisher returns a publishFunc that publishes a single UnitState
// by the given name to the provided Registry, with the given TTL
func newPublisher(reg registry.Registry, ttl time.Duration) publishFunc {
	return func(name string, us *unit.UnitState) {
		if us == nil {
			log.V(1).Infof("Destroying UnitState(%s) in Registry", name)
			err := reg.RemoveUnitState(name)
			if err != nil {
				log.Errorf("Failed to destroy UnitState(%s) in Registry: %v", name, err)
			}
		} else {
			// Sanity check - don't want to publish incomplete UnitStates
			// TODO(jonboulle): consider teasing apart a separate UnitState-like struct
			// so we can rely on a UnitState always being fully hydrated?

			// See https://github.com/coreos/fleet/issues/720
			//if len(us.UnitHash) == 0 {
			//	log.Errorf("Refusing to push UnitState(%s), no UnitHash: %#v", name, us)

			if len(us.MachineID) == 0 {
				log.Errorf("Refusing to push UnitState(%s), no MachineID: %#v", name, us)
			} else {
				log.V(1).Infof("Pushing UnitState(%s) to Registry: %#v", name, us)
				reg.SaveUnitState(name, us, ttl)
			}
		}
	}
}

func findAddressInRunningUnits(r *registry.Registry, lookup string) (string, bool) {
	js := r.GetJobState(lookup)
	if js == nil {
		return "", false
	}
	return fmt.Sprintf("%s:22", js.MachineState.PublicIP), true
}

func getMachinePage(reg registry.Registry, tok PageToken) (*schema.MachinePage, error) {
	all, err := reg.Machines()
	if err != nil {
		return nil, err
	}

	page := extractMachinePage(all, tok)
	return page, nil
}

// checkVersion makes a best-effort attempt to verify that fleetctl is at least as new as the
// latest fleet version found registered in the cluster. If any errors are encountered or fleetctl
// is >= the latest version found, it returns true. If it is < the latest found version, it returns
// false and a scary warning to the user.
func checkVersion(reg registry.Registry) (string, bool) {
	fv := version.SemVersion
	lv, err := reg.LatestVersion()
	if err != nil {
		log.Errorf("error attempting to check latest fleet version in Registry: %v", err)
	} else if lv != nil && fv.LessThan(*lv) {
		return fmt.Sprintf(oldVersionWarning, fv.String(), lv.String()), false
	}
	return "", true
}

func getJobsScheduledToMachine(r registry.Registry, machID string) []job.Job {
	var jobs []job.Job

	for _, j := range r.GetAllJobs() {
		tgt := r.GetJobTarget(j.Name)
		if tgt == "" || tgt != machID {
			continue
		}
		jobs = append(jobs, j)
	}

	return jobs
}

func getJobsScheduledToMachine(r registry.Registry, machID string) []job.Job {
	var jobs []job.Job

	jj, _ := r.Jobs()
	for _, j := range jj {
		if j.TargetMachineID == "" || j.TargetMachineID != machID {
			continue
		}
		jobs = append(jobs, j)
	}

	return jobs
}

func getUnitPage(reg registry.Registry, tok PageToken) (*schema.UnitPage, error) {
	all, err := reg.Jobs()
	if err != nil {
		return nil, err
	}

	page, err := extractUnitPage(reg, all, tok)
	if err != nil {
		return nil, err
	}

	return page, nil
}

func rpcAcquireLeadership(reg registry.Registry, lManager lease.Manager, machID string, ver int, ttl time.Duration) lease.Lease {
	existing, err := lManager.GetLease(engineLeaseName)
	if err != nil {
		log.Errorf("Unable to determine current lease: %v", err)
		return nil
	}

	var l lease.Lease
	if (existing == nil && reg.UseEtcdRegistry()) || (existing == nil && !reg.IsRegistryReady()) {
		l, err = lManager.AcquireLease(engineLeaseName, machID, ver, ttl)
		if err != nil {
			log.Errorf("Engine leadership acquisition failed: %v", err)
			return nil
		} else if l == nil {
			log.Infof("Unable to acquire engine leadership")
			return nil
		}
		log.Infof("Engine leadership acquired")
		return l
	}

	if existing != nil && existing.Version() >= ver {
		log.Debugf("Lease already held by Machine(%s) operating at acceptable version %d", existing.MachineID(), existing.Version())
		return existing
	}

	// TODO(hector): Here we could add a possible SLA to determine when the leader
	// is too busy. In such a case, we can trigger a new leader election
	if (existing != nil && reg.UseEtcdRegistry()) || (existing != nil && !reg.IsRegistryReady()) {
		rem := existing.TimeRemaining()
		l, err = lManager.StealLease(engineLeaseName, machID, ver, ttl+rem, existing.Index())
		if err != nil {
			log.Errorf("Engine leadership steal failed: %v", err)
			return nil
		} else if l == nil {
			log.Infof("Unable to steal engine leadership")
			return nil
		}

		log.Infof("Stole engine leadership from Machine(%s)", existing.MachineID())

		if rem > 0 {
			log.Infof("Waiting %v for previous lease to expire before continuing reconciliation", rem)
			<-time.After(rem)
		}

		return l
	}

	log.Infof("Engine leader is BUSY!")

	return existing

}

// desiredAgentState builds an *AgentState object that represents what the
// provided Agent should currently be doing.
func desiredAgentState(a *Agent, reg registry.Registry) (*AgentState, error) {
	units, err := reg.Units()
	if err != nil {
		log.Errorf("Failed fetching Units from Registry: %v", err)
		return nil, err
	}

	sUnits, err := reg.Schedule()
	if err != nil {
		log.Errorf("Failed fetching schedule from Registry: %v", err)
		return nil, err
	}

	ms := a.Machine.State()
	as := AgentState{
		MState: &ms,
		Jobs:   make(map[string]*job.Job),
	}

	sUnitMap := make(map[string]*job.ScheduledUnit)
	for _, sUnit := range sUnits {
		sUnit := sUnit
		sUnitMap[sUnit.Name] = &sUnit
	}

	for _, u := range units {
		j := &job.Job{
			Name:        u.Name,
			Unit:        u.Unit,
			TargetState: u.TargetState,
		}
		if !u.IsGlobal() {
			sUnit, ok := sUnitMap[u.Name]
			if !ok || sUnit.TargetMachineID == "" || sUnit.TargetMachineID != ms.ID {
				continue
			}
		}
		as.Jobs[u.Name] = j
	}

	return &as, nil
}

func findAddressInMachineList(r *registry.Registry, lookup string) (string, bool) {
	states := r.GetActiveMachines()
	var match *machine.MachineState

	for i, _ := range states {
		machState := states[i]
		if !strings.HasPrefix(machState.BootId, lookup) {
			continue
		} else if match != nil {
			log.Fatalf("Found more than one Machine, be more specfic")
		}
		match = &machState
	}

	if match == nil {
		return "", false
	}

	return fmt.Sprintf("%s:22", match.PublicIP), true
}

func printUnitStatus(c *cli.Context, r *registry.Registry, jobName string) {
	js := r.GetJobState(jobName)

	if js == nil {
		fmt.Printf("%s does not appear to be running\n", jobName)
		syscall.Exit(1)
	}

	addr := fmt.Sprintf("%s:22", js.MachineState.PublicIP)

	var err error
	var sshClient *gossh.ClientConn
	if tun := getTunnelFlag(c); tun != "" {
		sshClient, err = ssh.NewTunnelledSSHClient("core", tun, addr)
	} else {
		sshClient, err = ssh.NewSSHClient("core", addr)
	}
	if err != nil {
		log.Fatalf("Unable to establish SSH connection: %v", err)
	}

	defer sshClient.Close()

	cmd := fmt.Sprintf("systemctl status -l %s", jobName)
	stdout, err := ssh.Execute(sshClient, cmd)
	if err != nil {
		log.Fatalf("Unable to execute command over SSH: %s", err.Error())
	}

	for true {
		bytes, prefix, err := stdout.ReadLine()
		if err != nil {
			break
		}

		print(string(bytes))
		if !prefix {
			print("\n")
		}
	}
}

// ensureLeader will attempt to renew a non-nil Lease, falling back to
// acquiring a new Lease on the lead engine role.
func ensureLeader(prev registry.Lease, reg registry.Registry, machID string) (cur registry.Lease) {
	if prev != nil {
		err := prev.Renew(engineRoleLeasePeriod)
		if err == nil {
			cur = prev
			return
		} else {
			log.Errorf("Engine leadership could not be renewed: %v", err)
		}
	}

	var err error
	cur, err = reg.LeaseRole(engineRoleName, machID, engineRoleLeasePeriod)
	if err != nil {
		log.Errorf("Failed acquiring engine leadership: %v", err)
	} else if cur == nil {
		log.V(1).Infof("Unable to acquire engine leadership")
	} else {
		log.Infof("Acquired engine leadership")
	}

	return
}

// ensureLeader will attempt to renew a non-nil Lease, falling back to
// acquiring a new Lease on the lead engine role.
func ensureLeader(prev registry.Lease, reg registry.Registry, machID string, ttl time.Duration) (cur registry.Lease) {
	if prev != nil {
		err := prev.Renew(ttl)
		if err == nil {
			log.V(1).Infof("Engine leadership renewed")
			cur = prev
			return
		} else {
			log.Errorf("Engine leadership lost, renewal failed: %v", err)
		}
	}

	var err error
	cur, err = reg.LeaseRole(engineRoleName, machID, ttl)
	if err != nil {
		log.Errorf("Engine leadership acquisition failed: %v", err)
	} else if cur == nil {
		log.V(1).Infof("Unable to acquire engine leadership")
	} else {
		log.Infof("Engine leadership acquired")
	}

	return
}