// Copyright 2019 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

// Package devices models devices as state machines, and exposes capabilities for
// other utilites to work with.
package devices

import (
	"context"
	"fmt"
	"io"
	"net"
	"os/exec"
	"time"

	"go.fuchsia.dev/fuchsia/tools/net/serial"
	"go.fuchsia.dev/fuchsia/tools/net/sshutil"
	"go.fuchsia.dev/tools/netboot"
	"golang.org/x/crypto/ssh"
)

// DeviceState represents a physical state the device can be in.
// It is used when performing device recovery.
type DeviceState string

// Transition specifies an action to perform, a validation function to
// determine if the action succeeded or failed, and the appropriate
// states to put the device in.
type Transition struct {
	PerformAction func(context.Context) error
	Validate      func(context.Context) error
	SuccessState  DeviceState
	FailureState  DeviceState
}

const (
	// Healthy refers to the state in which the device is booted into zedboot.
	Healthy DeviceState = "zedboot"
	// Devices are marked Unrecoverable if all recovery attempts fail.
	Unrecoverable DeviceState = "unrecoverable"
	// All devices are created in the Inital state. This is the entrypoint into the
	// state machine.
	Initial DeviceState = "unknown"
	// Duration to wait before sending dm reboot-recovery after powercycle.
	powercycleWait = 1 * time.Minute
	// netbootTimeout is the duration to wait when creating a new netboot client.
	netbootTimeout = 10 * time.Second
)

// powerManager describes any object that gives a device powercycle functionality.
// Different devices use different power management systems, so each device is expected
// to provide a concrete implementation of this.
type powerManager interface {
	Powercycle(context.Context) error
}

// fuchsiaDevice represents an abstract fuchsia device. It should not be
// directly instantiated.
type fuchsiaDevice struct {
	mac           string
	nodename      string
	networkIf     string
	state         DeviceState
	transitionMap map[DeviceState]*Transition
	signers       []ssh.Signer
	serial        io.ReadWriteCloser
	BootserverCmd []string
	power         powerManager
}

// initFuchsiaDevice initializes the fields in fuchsiaDevice using the given config.
// Concrete implementations of fuchsiaDevice can use this to avoid code duplication.
func initFuchsiaDevice(device *fuchsiaDevice, config DeviceConfig, bootserverCmdStub []string) error {
	var signers []ssh.Signer
	var err error
	if config.SSHKeys != nil {
		signers, err = parseOutSigners(config.SSHKeys)
		if err != nil {
			return fmt.Errorf("could not parse out signers from private keys: %v", err)
		}
	}
	var s io.ReadWriteCloser
	if config.Serial != "" {
		s, err = serial.Open(config.Serial)
		if err != nil {
			return fmt.Errorf("could not open serial line: %s", config.Serial)
		}
	}
	device.nodename = config.Network.Nodename
	device.networkIf = config.Network.Interface
	device.state = Initial
	device.signers = signers
	device.serial = s
	device.mac = config.Network.Mac
	device.BootserverCmd = append(bootserverCmdStub, device.nodename)
	return nil
}

// Mac returns the mac address of this device.
func (f *fuchsiaDevice) Mac() string {
	return f.mac
}

// Interface returns the network interface of this device.
func (f *fuchsiaDevice) Interface() string {
	return f.networkIf
}

// HasSerial returns true if this device has a serial line.
func (f *fuchsiaDevice) HasSerial() bool {
	return f.serial != nil
}

// Powercycle uses an out of band method to reboot a fuchsia device.
// Is a no-op if no out of band is possible.
// TODO(rudymathu): remove the no-op as soon as all devices have an out of band
// reboot method.
func (f *fuchsiaDevice) Powercycle(ctx context.Context) error {
	if f.power != nil {
		return f.power.Powercycle(ctx)
	}
	return nil
}

// SendSSHCommand lets us send a command via SSH to the fuchsia device.
func (f *fuchsiaDevice) SendSSHCommand(ctx context.Context, command string) error {
	if f.signers == nil {
		return fmt.Errorf("device %s does not support SSH", f.nodename)
	}
	config, err := sshutil.DefaultSSHConfigFromSigners(f.signers...)
	if err != nil {
		return err
	}

	client, err := sshutil.ConnectToNode(ctx, f.nodename, config)
	if err != nil {
		return err
	}

	defer client.Close()

	session, err := client.NewSession()
	if err != nil {
		return err
	}

	defer session.Close()

	// Invoke `dm reboot-recovery`
	err = session.Start(command)
	if err != nil {
		return err
	}

	done := make(chan error)
	go func() {
		done <- session.Wait()
	}()

	select {
	case err := <-done:
		return err
	case <-time.After(10 * time.Second):
		return nil
	}
}

// SendSerialCommand lets us send a command via serial to the fuchsia device.
func (f *fuchsiaDevice) SendSerialCommand(ctx context.Context, command string) error {
	if f.serial == nil {
		return fmt.Errorf("device %s does not support serial", f.nodename)
	}
	_, err := io.WriteString(f.serial, fmt.Sprintf("\n%s\n", command))
	return err
}

// ReadSerialData fills the given buffer with bytes from the serial line.
func (f *fuchsiaDevice) ReadSerialData(ctx context.Context, buf []byte) error {
	_, err := io.ReadAtLeast(f.serial, buf, len(buf))
	return err
}

// SoftReboot is a convenience function that reboots into the specified partition
// using the specified method.
func (f *fuchsiaDevice) SoftReboot(ctx context.Context, partition string, method string) error {
	command := ""
	if partition == "A" {
		command = "dm reboot"
	} else if partition == "R" {
		command = "dm reboot-recovery"
	} else {
		return fmt.Errorf("rebooting into partition %s is not supported", partition)
	}

	if method == "ssh" {
		return f.SendSSHCommand(ctx, command)
	} else if method == "serial" {
		return f.SendSerialCommand(ctx, command)
	}
	return fmt.Errorf("reboot method %s is not supported", method)
}

// State returns the current state of the fuchsia device.
func (f *fuchsiaDevice) State() DeviceState {
	return f.state
}

// Nodename returns the nodename of the fuchsia device.
func (f *fuchsiaDevice) Nodename() string {
	return f.nodename
}

// Transition tries to move the device into the next state. It does so by:
// 1) Getting the transition associated with the current state.
// 2) Performing the action specified by the transition.
// 3) Validating if the action worked, and moving the device into the corresponding
//    failure or success state.
func (f *fuchsiaDevice) Transition(ctx context.Context) error {
	transition := f.transitionMap[f.state]
	if err := transition.PerformAction(ctx); err != nil {
		f.state = transition.FailureState
		return err
	}
	if err := transition.Validate(ctx); err != nil {
		f.state = transition.FailureState
		return err
	}
	f.state = transition.SuccessState
	return nil
}

func (f *fuchsiaDevice) serialRebootRecovery(ctx context.Context) error {
	return f.SoftReboot(ctx, "R", "serial")
}

func (f *fuchsiaDevice) sshRebootRecovery(ctx context.Context) error {
	return f.SoftReboot(ctx, "R", "ssh")
}

// pingZedboot creates a netboot client and attempts to ping the zedboot ipv6
// address of the device with the given nodename.
func pingZedboot(n *netboot.Client, nodename string) error {
	netsvcAddr, err := n.Discover(nodename, false)
	if err != nil {
		return fmt.Errorf("Failed to discover netsvc addr: %v.", err)
	}
	netsvcIpAddr := &net.IPAddr{IP: netsvcAddr.IP, Zone: netsvcAddr.Zone}
	cmd := exec.Command("ping", "-6", netsvcIpAddr.String(), "-c", "1")
	if _, err = cmd.Output(); err != nil {
		return fmt.Errorf("Failed to ping netsvc addr %s: %v.", netsvcIpAddr, err)
	}
	return nil
}

// ensureNotFuchsia creates a netboot client and ensures it cannot ping the
// fuchsia ipv6 address of the device with the given nodename.
func ensureNotFuchsia(n *netboot.Client, nodename string) error {
	fuchsiaAddr, err := n.Discover(nodename, true)
	if err != nil {
		return fmt.Errorf("Failed to discover fuchsia addr: %v.", err)
	}
	fuchsiaIpAddr := &net.IPAddr{IP: fuchsiaAddr.IP, Zone: fuchsiaAddr.Zone}
	cmd := exec.Command("ping", "-6", fuchsiaIpAddr.String(), "-c", "1")
	if _, err = cmd.Output(); err == nil {
		return fmt.Errorf("Device is in Fuchsia, should be in Zedboot.")
	}
	return nil
}

// deviceInZedboot ensures that the device with the given nodename is in zedboot.
func deviceInZedboot(n *netboot.Client, nodename string) error {
	if err := pingZedboot(n, nodename); err != nil {
		return err
	}
	if err := ensureNotFuchsia(n, nodename); err != nil {
		return err
	}
	return nil
}