package main

import (
	"bytes"
	"context"
	"encoding/binary"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io/fs"
	"io/ioutil"
	"net"
	"os"
	"path"

	"git.faercol.me/faercol/http-boot-config/config/efivar"
	"git.faercol.me/faercol/http-boot-config/config/logger"
	"git.faercol.me/faercol/http-boot-config/config/prober"
	"git.faercol.me/faercol/http-boot-config/config/remote"
	"github.com/google/uuid"
)

type action int

const (
	actionList action = iota
	actionGetRemote
	actionEnroll
	actionStatus
	actionUnknown
)

const (
	defaultEFIDest  = "/EFI/httpboot/"
	defaultEFIMount = "/boot"
	defaultEFISrc   = "httpboot.efi"
)

type cliArgs struct {
	debug         bool
	colour        bool
	action        action
	id            uuid.UUID
	remoteAddr    string
	prettyPrint   bool
	name          string
	efiDest       string
	efiMountPoint string
	efiDisk       string
	efiSrc        string
	ifaceName     string
	install       bool
}

var defaultArgs cliArgs = cliArgs{
	debug:  false,
	colour: true,
	action: actionUnknown,
}

func parseArgs() (cliArgs, error) {
	args := defaultArgs
	var firstArg int

	listFlagSet := flag.NewFlagSet("list", flag.ExitOnError)

	statusFlagSet := flag.NewFlagSet("status", flag.ExitOnError)
	hostStatusFlag := statusFlagSet.String("remote-host", "http://localhost:5000", "Address for the remote boot server")

	getRemoteFlagSet := flag.NewFlagSet("get-remote", flag.ExitOnError)
	uuidFlag := getRemoteFlagSet.String("uuid", "", "Client UUID")
	hostGetRemoteFlag := getRemoteFlagSet.String("remote-host", "http://localhost:5000", "Address for the remote boot server")
	jsonFlag := getRemoteFlagSet.Bool("json", false, "Display the result in JSON format")

	enrollFlagSet := flag.NewFlagSet("enroll", flag.ExitOnError)
	hostEnrollFlag := enrollFlagSet.String("remote-host", "http://localhost:5000", "Address for the remote boot server")
	nameFlag := enrollFlagSet.String("name", "default", "Name for the client on the remote boot server")
	efiSrcFlag := enrollFlagSet.String("src", defaultEFISrc, "HTTP EFI app to install")
	efiDestFlag := enrollFlagSet.String("dest", defaultEFIDest, "Directory in which to store the EFI app")
	efiMountPointFlag := enrollFlagSet.String("mountpoint", defaultEFIMount, "EFI partition mountpoint")
	installFlag := enrollFlagSet.Bool("install", false, "Install the EFI app in the system")
	efiDiskFlag := enrollFlagSet.String("disk", "/dev/sda1", "Partition in which to install the EFI app")
	ifaceNameFlag := enrollFlagSet.String("iface", "eth0", "Iface name to use for the boot loader")

	debugFlag := flag.Bool("debug", false, "Display debug logs")
	noColourFlag := flag.Bool("no-colour", false, "Disable colour logs")

	for i, v := range os.Args {
		switch v {
		case "list":
			args.action = actionList
		case "get-remote":
			args.action = actionGetRemote
		case "enroll":
			args.action = actionEnroll
		case "status":
			args.action = actionStatus
		default:
			continue
		}
		firstArg = i + 1
	}

	switch args.action {
	case actionList:
		listFlagSet.Parse(os.Args[firstArg:])
	case actionStatus:
		statusFlagSet.Parse(os.Args[firstArg:])
		args.remoteAddr = *hostStatusFlag
	case actionGetRemote:
		getRemoteFlagSet.Parse(os.Args[firstArg:])
		parsedID, err := uuid.Parse(*uuidFlag)
		if err != nil {
			return args, fmt.Errorf("invalid format for uuid %q", *uuidFlag)
		}
		args.id = parsedID
		args.remoteAddr = *hostGetRemoteFlag
		args.prettyPrint = !*jsonFlag
	case actionEnroll:
		enrollFlagSet.Parse(os.Args[firstArg:])
		args.remoteAddr = *hostEnrollFlag
		args.name = *nameFlag
		args.efiSrc = *efiSrcFlag
		args.efiDest = *efiDestFlag
		args.install = *installFlag
		args.efiDisk = *efiDiskFlag
		args.efiMountPoint = *efiMountPointFlag
		args.ifaceName = *ifaceNameFlag
	default:
		flag.Parse()
		return cliArgs{}, errors.New("missing an action")
	}

	flag.Parse()
	args.debug = *debugFlag
	args.colour = !*noColourFlag
	return args, nil
}

func displayAppList(l *logger.SimpleLogger) {
	l.Info("Checking EFI directory for available boot images...")
	apps, err := prober.GetEFIApps()
	if err != nil {
		if errors.Is(err, fs.ErrPermission) {
			l.Fatal("Permission error, try to run the command as sudo")
		}
		l.Fatalf("Failed to check EFI directory: %s", err.Error())
	}

	l.Info("Found the following EFI applications:")
	for _, a := range apps {
		prefix := " "
		if a.Active {
			prefix = "*"
		}
		l.Infof("\t- %s[%d] %s: %s", prefix, a.ID, a.Name, a.DevicePath)
	}
}

func getRemoteConfig(l *logger.SimpleLogger, host string, id uuid.UUID, pretty bool) {
	if pretty {
		l.Info("Getting config from remote server...")
		if err := remote.DisplayRemoteConfigPretty(context.Background(), host, id, l); err != nil {
			l.Fatal(err.Error())
		}
	} else {
		if err := remote.DisplayRemoteConfigJSON(context.Background(), host, id, l); err != nil {
			l.Fatal(err.Error())
		}
	}
}

func enroll(l *logger.SimpleLogger, host, name, src, dest, destPart, mount, iface string, install bool) {
	if install {
		l.Debug("Installing boot application")
		if err := installApp(l, src, dest, destPart, mount); err != nil {
			l.Fatalf("Failed to install boot application: %s", err.Error())
		}
	}
	l.Info("Enrolling client")
	clientConf, err := remote.Enroll(context.Background(), host, name, l)
	if err != nil {
		l.Fatal(err.Error())
	}

	if install {
		l.Debug("Setting EFI vars")
		if err := installVars(l, clientConf.ID, clientConf.MulticastPort, clientConf.MulticastGroup, iface); err != nil {
			l.Fatalf("Failed to install efi vars: %s", err.Error())
		}
	}
}

func getEFIConf() (clientID uuid.UUID, port int, group net.IP, ip net.IP, err error) {
	idRaw, err := efivar.GetVar(efivar.VendorID, efivar.VarBootUUID)
	if err != nil {
		return
	}
	clientID, err = uuid.ParseBytes(idRaw)
	if err != nil {
		err = fmt.Errorf("invalid uuid %q: %w", string(idRaw), err)
		return
	}

	portRaw, err := efivar.GetVar(efivar.VendorID, efivar.VarBootPort)
	if err != nil {
		return
	}
	if len(portRaw) == 1 {
		portRaw = append(portRaw, 0)
	}
	port = int(binary.LittleEndian.Uint16(portRaw))

	groupRaw, err := efivar.GetVar(efivar.VendorID, efivar.VarBootGroup)
	if err != nil {
		return
	}
	group = net.IP(groupRaw)

	ipRaw, err := efivar.GetVar(efivar.VendorID, efivar.VarBootIP)
	if err != nil {
		return
	}
	ip = net.IP(ipRaw)
	return
}

func getStatus(l *logger.SimpleLogger, host string) {
	l.Debug("Checking EFI vars")
	clientID, port, group, ip, err := getEFIConf()
	if err != nil {
		if errors.Is(err, efivar.ErrNotFound) {
			l.Info("Boot client is not configured")
			return
		}
		l.Fatalf("Failed to get EFI configuration: %s", err.Error())
	}

	l.Infof("Registered client ID %s", clientID.String())
	l.Infof("Registered multicast group %s:%d", group.String(), port)
	l.Infof("HTTP_BOOT client using local ip address %s", ip.String())

	l.Debug("Getting remote config")
	remoteRaw, err := remote.GetRemoteConfig(context.Background(), host, clientID, l)
	if err != nil {
		l.Fatalf("Failed to get remote config: %s", err.Error())
	}

	var remoteConf remote.ClientConfig
	if err := json.Unmarshal(remoteRaw, &remoteConf); err != nil {
		l.Fatalf("Invalid remote config: %s", err.Error())
	}
	l.Infof("Remote HTTP Boot configured and set to boot to %s", remoteConf.EFIConfig.Options[remoteConf.EFIConfig.SelectedOption].Name)

	l.Debug("Checking selected EFI image")
	apps, err := prober.GetEFIApps()
	if err != nil {
		if errors.Is(err, fs.ErrPermission) {
			l.Fatal("Permission error, try to run the command as sudo")
		}
		l.Fatalf("Failed to check EFI directory: %s", err.Error())
	}
	installed := false
	for _, a := range apps {
		if a.Name == prober.HTTPBootLabel {
			installed = true
			l.Info("HTTP_BOOT efi application found in the EFI partitions")
		}
	}
	if !installed {
		l.Warning("HTTP_BOOT efi application is not installed")
		return
	}
	httpNext, err := prober.IsHTTPBootNext()
	if err != nil {
		l.Fatalf("Failed to check EFI app selected for next boot: %s", err.Error())
	}
	if httpNext {
		l.Info("HTTP_BOOT set for next boot")
	} else {
		l.Warning("HTTP_BOOT is not set for next boot")
	}
}

func getIPAddr(ifaceName string) (net.IP, error) {
	iface, err := net.InterfaceByName(ifaceName)
	if err != nil {
		return nil, fmt.Errorf("failed to get iface: %w", err)
	}

	addrs, err := iface.Addrs()
	if err != nil {
		return nil, fmt.Errorf("failed to get addresses for iface: %w", err)
	}
	for _, a := range addrs {
		parsedIP, _, err := net.ParseCIDR(a.String())
		if err != nil {
			return nil, fmt.Errorf("invalid addr %s: %w", a.String(), err)
		}
		if parsedIP.IsLinkLocalUnicast() && parsedIP.To4() == nil {
			return parsedIP, nil
		}
	}
	return nil, errors.New("no link-local ipv6 found")
}

func installApp(l *logger.SimpleLogger, srcPath, destDir, destPart, mount string) error {
	destFullDir := path.Join(mount, destDir)
	l.Debug("Creating dest directory")
	if err := os.MkdirAll(destFullDir, 0o755); err != nil {
		return fmt.Errorf("failed to create dest directory: %w", err)
	}
	l.Debugf("Opening source file %s for copy", srcPath)
	input, err := ioutil.ReadFile(srcPath)
	if err != nil {
		return fmt.Errorf("failed to open source file for copy: %w", err)
	}
	dest := path.Join(destFullDir, path.Base(srcPath))
	l.Debugf("Copying efi image to %s", dest)
	if err := ioutil.WriteFile(dest, input, 0o755); err != nil {
		return fmt.Errorf("failed to copy efi image: %w", err)
	}
	l.Debug("Deleting source file after copy")
	if err := os.Remove(srcPath); err != nil {
		return fmt.Errorf("failed to remove source file after copy: %w", err)
	}
	l.Debug("Installing app in efibootmgr")
	if err := prober.Install(path.Join(destDir, path.Base(srcPath)), destPart); err != nil {
		return fmt.Errorf("failed to install EFI app in efibootmgr: %w", err)
	}
	return nil
}

func installVars(l *logger.SimpleLogger, clientID uuid.UUID, port int, group net.IP, ifaceName string) error {
	l.Debug("Setting multicast port")
	buf := new(bytes.Buffer)
	if err := binary.Write(buf, binary.LittleEndian, uint16(port)); err != nil {
		return fmt.Errorf("invalid port format: %w", err)
	}
	if err := efivar.SetVar(efivar.VendorID, efivar.VarBootPort, buf.Bytes()); err != nil {
		return fmt.Errorf("failed to set port number: %w", err)
	}

	l.Debug("Setting multicast group")
	if err := efivar.SetVar(efivar.VendorID, efivar.VarBootGroup, group); err != nil {
		return fmt.Errorf("failed to set multicast group: %w", err)
	}

	l.Debug("Setting UUID")
	if err := efivar.SetVar(efivar.VendorID, efivar.VarBootUUID, []byte(clientID.String())); err != nil {
		return fmt.Errorf("failed to set client UUID: %w", err)
	}

	l.Debug("Setting link local IPv6")
	ipAddr, err := getIPAddr(ifaceName)
	if err != nil {
		return fmt.Errorf("failed to get ipv6 link local address: %w", err)
	}
	if err := efivar.SetVar(efivar.VendorID, efivar.VarBootIP, ipAddr); err != nil {
		return fmt.Errorf("failed to set ip address: %w", err)
	}
	return nil
}

func main() {
	args, err := parseArgs()
	if err != nil {
		l := logger.New(true, false)
		l.Fatalf("Invalid command: %s", err.Error())
	}
	l := logger.New(args.colour, args.debug)
	fmt.Print("")

	switch args.action {
	case actionList:
		displayAppList(l)
	case actionStatus:
		getStatus(l, args.remoteAddr)
	case actionGetRemote:
		getRemoteConfig(l, args.remoteAddr, args.id, args.prettyPrint)
	case actionEnroll:
		enroll(l, args.remoteAddr, args.name, args.efiSrc, args.efiDest, args.efiDisk, args.efiMountPoint, args.ifaceName, args.install)
	default:
		l.Fatal("Unknown action")
	}
}