// Copyright 2018 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

import (
	"context"
	"errors"
	"fmt"
	"io"
	"log"
	"time"

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

// TODO(IN-977) Clean this up per suggestions in go/fxr/251550

const (
	// Controller machines use 192.168.42.1/24 for swarming bots
	// This will broadcast to that entire subnet.
	botBroadcastAddr = "192.168.42.255:9"

	// Controller machines have multiple interfaces, currently
	// 'eno2' is used for swarming bots.
	botInterface = "eno2"

	// Duration to wait before sending dm reboot-recovery after powercycle
	powercycleWait = 10 * time.Second
)

// PowerClient represents a power management configuration for a particular device.
type PowerClient struct {
	// Type is the type of manager to use.
	Type string `json:"type"`

	// Host is the network hostname of the manager, e.g. fuchsia-tests-pdu-001.
	Host string `json:"host"`

	// HostHwAddr is the ethernet MAC address of the manager,  e.g. 10:10:10:10:10:10
	HostMACAddr string `json:"host_mac_addr"`

	// Username is the username used to log in to the manager.
	Username string `json:"username"`

	// Password is the password used to log in to the manager..
	Password string `json:"password"`

	// PDUPort is the port the PDU uses to connect to this device.
	PDUPort int `json:"pdu_port"`
}

type Rebooter interface {
	reboot() error
}

type SshRebooter struct {
	nodename string
	signers []ssh.Signer
}

type SerialRebooter struct {
	serial io.ReadWriter
}

// RebootDevice attempts to reboot the specified device into recovery, and
// additionally uses the given configuration to reboot the device if specified.
func (c PowerClient) RebootDevice(signers []ssh.Signer, nodename string, serial io.ReadWriter) error {
	var rebooter Rebooter
	log.Printf("Attempting to soft reboot device %s", nodename)
	if serial != nil {
		log.Printf("Using serial to reboot.")
		rebooter = NewSerialRebooter(serial)
	} else {
		log.Printf("Using SSH to reboot.")
		rebooter = NewSSHRebooter(nodename, signers)
	}
	return rebooter.reboot()
}

// Powercycle the device using an out of band method
func (c PowerClient) Powercycle(nodename string, serial io.ReadWriter) error {
	log.Printf("Attempting to powercycle device %s", nodename)
	switch c.Type {
		case "AMT":
			log.Printf("Using AMT to powercycle.")
			return AMTReboot(c.Host, c.Username, c.Password)
		case "WOL":
			if serial == nil {
				return errors.New(fmt.Sprintf("WOL Reboot requires serial connection."))
			}
			log.Printf("Using WOL to powercycle")
			if err := WOLReboot(botBroadcastAddr, botInterface, c.HostMACAddr); err != nil {
				return err
			}
		case "PDU":
			if serial == nil {
				return errors.New(fmt.Sprintf("PDU Reboot requires serial connection."))
			}
			log.Printf("Using PDU to powercycle")
			if err := PDUReboot(c.PDUPort, c.Host, c.Username, c.Password); err != nil {
				return err
			}
		default:
			return errors.New(fmt.Sprintf("%v does not have AMT, WOL, or PDU support. Cannot powercycle.", nodename))
	}
	// Send dm reboot-recovery to the device so that it always boots into zedboot.
	log.Printf("Powercycle complete; using serial to send dm reboot recovery.")
	time.Sleep(powercycleWait)
	rebooter := NewSerialRebooter(serial)
	return rebooter.reboot()
}

func NewSerialRebooter(serial io.ReadWriter) *SerialRebooter {
	return &SerialRebooter{
		serial: serial,
	}
}

func (s *SerialRebooter) reboot() error {
	// TODO(fxb/42761): This doesn't work on x86 due to issues in the serial
	// implementation in tools/net/serial. We will need to solve this if we
	// want to depend on serial to deliver messages reliably.
	_, err := io.WriteString(s.serial, "\ndm reboot-recovery\n")
	return err
}

func NewSSHRebooter(nodename string, signers []ssh.Signer) *SshRebooter {
	return &SshRebooter{
		nodename: nodename,
		signers: signers,
	}
}

func (s *SshRebooter) reboot() error {
	config, err := sshutil.DefaultSSHConfigFromSigners(s.signers...)
	if err != nil {
		return err
	}

	client, err := sshutil.ConnectToNode(context.Background(), s.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` with a 2 second delay in the background, then exit the SSH shell.
	// This prevents the SSH connection from hanging waiting for `dm reboot-recovery` to return.
	err = session.Start("{ sleep 2; dm reboot-recovery; } >/dev/null & exit")
	if err != nil {
		return err
	}

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

	select {
	case err := <-done:
		return err
	}
}