// © Broadcom. All Rights Reserved.
// The term "Broadcom" refers to Broadcom Inc. and/or its subsidiaries.
// SPDX-License-Identifier: Apache-2.0

package vm

import (
	"context"
	"flag"
	"fmt"
	"os"
	"strconv"
	"strings"

	"github.com/vmware/govmomi/cli"
	"github.com/vmware/govmomi/cli/flags"
	"github.com/vmware/govmomi/object"
	"github.com/vmware/govmomi/vim25/types"
)

type customize struct {
	*flags.VirtualMachineFlag

	alc       int
	prefix    types.CustomizationPrefixName
	tz        string
	domain    string
	host      types.CustomizationFixedName
	mac       flags.StringList
	ip        flags.StringList
	ip6       flags.StringList
	gateway   flags.StringList
	netmask   flags.StringList
	dnsserver flags.StringList
	dnssuffix flags.StringList
	kind      string
	username  string
	org       string
	pwd       string
	expirePwd bool
	script    string
	sig       string
}

func init() {
	cli.Register("vm.customize", &customize{})
}

func (cmd *customize) Register(ctx context.Context, f *flag.FlagSet) {
	cmd.VirtualMachineFlag, ctx = flags.NewVirtualMachineFlag(ctx)
	cmd.VirtualMachineFlag.Register(ctx, f)

	f.IntVar(&cmd.alc, "auto-login", 0, "Number of times the VM should automatically login as an administrator")
	f.StringVar(&cmd.prefix.Base, "prefix", "", "Host name generator prefix")
	f.StringVar(&cmd.tz, "tz", "", "Time zone")
	f.StringVar(&cmd.domain, "domain", "", "Domain name")
	f.StringVar(&cmd.host.Name, "name", "", "Host name")
	f.Var(&cmd.mac, "mac", "MAC address")
	cmd.mac = nil
	f.Var(&cmd.ip, "ip", "IPv4 address")
	cmd.ip = nil
	f.Var(&cmd.ip6, "ip6", "IPv6 addresses with optional netmask (defaults to /64), separated by comma")
	cmd.ip6 = nil
	f.Var(&cmd.gateway, "gateway", "Gateway")
	cmd.gateway = nil
	f.Var(&cmd.netmask, "netmask", "Netmask")
	cmd.netmask = nil
	f.Var(&cmd.dnsserver, "dns-server", "DNS server list")
	cmd.dnsserver = nil
	f.Var(&cmd.dnssuffix, "dns-suffix", "DNS suffix list")
	cmd.dnssuffix = nil
	f.StringVar(&cmd.kind, "type", "Linux", "Customization type if spec NAME is not specified (Linux|Windows)")
	f.StringVar(&cmd.username, "username", "", "Windows only : full name of the end user in firstname lastname format")
	f.StringVar(&cmd.org, "org", "", "Windows only : name of the org that owns the VM")
	f.StringVar(&cmd.pwd, "pwd", "", "The new administrator (for Windows) or root (for Linux) password")
	f.BoolVar(&cmd.expirePwd, "expire-pwd", false, "Expire administrator (for Windows) or root (for Linux) password after customization")
	f.StringVar(&cmd.script, "script-file", "", "Path to a script to run before and after customization")
	f.StringVar(&cmd.sig, "vcfa-sig", "", "The VCF Automation signature to identify customization request source")
}

func (cmd *customize) Usage() string {
	return "[NAME]"
}

func (cmd *customize) Description() string {
	return `Customize VM.

Optionally specify a customization spec NAME.

The '-ip', '-netmask' and '-gateway' flags are for static IP configuration.
If the VM has multiple NICs, an '-ip' and '-netmask' must be specified for each.

The '-dns-server' and '-dns-suffix' flags can be specified multiple times.

Windows -tz value requires the Index (hex): https://support.microsoft.com/en-us/help/973627/microsoft-time-zone-index-values

Examples:
  govc vm.customize -vm VM NAME
  govc vm.customize -vm VM -name my-hostname -ip dhcp
  govc vm.customize -vm VM -gateway GATEWAY -ip NEWIP -netmask NETMASK -dns-server DNS1,DNS2 NAME
  # Multiple -ip without -mac are applied by vCenter in the order in which the NICs appear on the bus
  govc vm.customize -vm VM -ip 10.0.0.178 -netmask 255.255.255.0 -ip 10.0.0.162 -netmask 255.255.255.0
  # Multiple -ip with -mac are applied by vCenter to the NIC with the given MAC address
  govc vm.customize -vm VM -mac 00:50:56:be:dd:f8 -ip 10.0.0.178 -netmask 255.255.255.0 -mac 00:50:56:be:60:cf -ip 10.0.0.162 -netmask 255.255.255.0
  # Dual stack IPv4/IPv6, single NIC
  govc vm.customize -vm VM -ip 10.0.0.1 -netmask 255.255.255.0 -ip6 '2001:db8::1/64' -name my-hostname NAME
  # DHCPv6, single NIC
  govc vm.customize -vm VM -ip6 dhcp6 NAME
  # Static IPv6, three NICs, last one with two addresses
  govc vm.customize -vm VM -ip6 2001:db8::1/64 -ip6 2001:db8::2/64 -ip6 2001:db8::3/64,2001:db8::4/64 NAME
  govc vm.customize -vm VM -auto-login 3 NAME
  govc vm.customize -vm VM -prefix demo NAME
  govc vm.customize -vm VM -tz America/New_York NAME
  # Change password, expire password, run customization script, and specify the VCF Automation signature
  govc vm.customize -vm VM -name my-hostname -ip dhcp -pwd '<your password>' -expire-pwd -script-file './test.bat' -vcfa-sig '<your hash>'`
}

// Parse a string of multiple IPv6 addresses with optional netmask; separated by comma
func parseIPv6Argument(argv string) (ipconf []types.BaseCustomizationIpV6Generator, err error) {
	for _, substring := range strings.Split(argv, ",") {
		// remove leading and trailing white space
		substring = strings.TrimSpace(substring)
		// handle "dhcp6" and lists of static IPv6 addresses
		switch substring {
		case "dhcp6":
			ipconf = append(
				ipconf,
				&types.CustomizationDhcpIpV6Generator{},
			)
		default:
			// check if subnet mask was specified
			switch strings.Count(substring, "/") {
			// no mask, set default
			case 0:
				ipconf = append(ipconf, &types.CustomizationFixedIpV6{
					IpAddress:  substring,
					SubnetMask: 64,
				})
			// a single forward slash was found: parse and use subnet mask
			case 1:
				parts := strings.Split(substring, "/")
				mask, err := strconv.Atoi(parts[1])
				if err != nil {
					return nil, fmt.Errorf("unable to convert subnet mask to int: %w", err)
				}
				ipconf = append(ipconf, &types.CustomizationFixedIpV6{
					IpAddress:  parts[0],
					SubnetMask: int32(mask),
				})
			// too many forward slashes; return error
			default:
				return nil, fmt.Errorf("unable to parse IPv6 address (too many subnet separators): %s", substring)
			}
		}
	}
	return ipconf, nil
}

func (cmd *customize) Run(ctx context.Context, f *flag.FlagSet) error {
	vm, err := cmd.VirtualMachineFlag.VirtualMachine()
	if err != nil {
		return err
	}

	if vm == nil {
		return flag.ErrHelp
	}

	var spec *types.CustomizationSpec

	name := f.Arg(0)
	if name == "" {
		spec = &types.CustomizationSpec{
			NicSettingMap: make([]types.CustomizationAdapterMapping, len(cmd.ip)),
		}

		switch cmd.kind {
		case "Linux":
			spec.Identity = &types.CustomizationLinuxPrep{
				HostName: new(types.CustomizationVirtualMachineName),
			}
		case "Windows":
			spec.Identity = &types.CustomizationSysprep{
				UserData: types.CustomizationUserData{
					ComputerName: new(types.CustomizationVirtualMachineName),
				},
			}
		default:
			return flag.ErrHelp
		}
	} else {
		m := object.NewCustomizationSpecManager(vm.Client())

		exists, err := m.DoesCustomizationSpecExist(ctx, name)
		if err != nil {
			return err
		}
		if !exists {
			return fmt.Errorf("specification %q does not exist", name)
		}

		item, err := m.GetCustomizationSpec(ctx, name)
		if err != nil {
			return err
		}

		spec = &item.Spec
	}

	if len(cmd.ip) > len(spec.NicSettingMap) {
		return fmt.Errorf("%d -ip specified, spec %q has %d", len(cmd.ip), name, len(spec.NicSettingMap))
	}

	sysprep, isWindows := spec.Identity.(*types.CustomizationSysprep)
	linprep, _ := spec.Identity.(*types.CustomizationLinuxPrep)

	if isWindows {
		sysprep.Identification.JoinDomain = cmd.domain
		sysprep.UserData.FullName = cmd.username
		sysprep.UserData.OrgName = cmd.org
	} else {
		linprep.Domain = cmd.domain
	}

	if len(cmd.dnsserver) != 0 {
		if !isWindows {
			for _, s := range cmd.dnsserver {
				spec.GlobalIPSettings.DnsServerList =
					append(spec.GlobalIPSettings.DnsServerList, strings.Split(s, ",")...)
			}
		}
	}

	spec.GlobalIPSettings.DnsSuffixList = cmd.dnssuffix

	if cmd.prefix.Base != "" {
		if isWindows {
			sysprep.UserData.ComputerName = &cmd.prefix
		} else {
			linprep.HostName = &cmd.prefix
		}
	}

	if cmd.host.Name != "" {
		if isWindows {
			sysprep.UserData.ComputerName = &cmd.host
		} else {
			linprep.HostName = &cmd.host
		}
	}

	if cmd.sig != "" {
		opt := &types.OptionValue{
			Key:   "vcfaVmHashKey",
			Value: cmd.sig,
		}

		if isWindows {
			sysprep.ExtraConfig = append(sysprep.ExtraConfig, opt)
		} else {
			linprep.ExtraConfig = append(linprep.ExtraConfig, opt)
		}
	}

	if cmd.script != "" {
		data, err := os.ReadFile(cmd.script)
		if err != nil {
			return fmt.Errorf("failed to read script file '%s': %v", cmd.script, err)
		}

		if isWindows {
			sysprep.ScriptText = string(data)
		} else {
			linprep.ScriptText = string(data)
		}
	}

	if cmd.pwd != "" {
		password := &types.CustomizationPassword{
			Value:     cmd.pwd,
			PlainText: true,
		}

		if isWindows {
			sysprep.GuiUnattended.Password = password
		} else {
			linprep.Password = password
		}
	}

	if cmd.expirePwd {
		b := true
		if isWindows {
			sysprep.ResetPassword = &b
		} else {
			linprep.ResetPassword = &b
		}
	}

	if cmd.alc != 0 {
		if !isWindows {
			return fmt.Errorf("option '-auto-login' is Windows only")
		}
		if cmd.pwd == "" {
			return fmt.Errorf("option '-auto-login' requires '-pwd' to be set")
		}
		sysprep.GuiUnattended.AutoLogon = true
		sysprep.GuiUnattended.AutoLogonCount = int32(cmd.alc)
	}

	if cmd.tz != "" {
		if isWindows {
			tz, err := strconv.ParseInt(cmd.tz, 16, 32)
			if err != nil {
				return fmt.Errorf("converting -tz=%q: %s", cmd.tz, err)
			}
			sysprep.GuiUnattended.TimeZone = int32(tz)
		} else {
			linprep.TimeZone = cmd.tz
		}
	}

	for i, ip := range cmd.ip {
		nic := &spec.NicSettingMap[i]
		switch ip {
		case "dhcp":
			nic.Adapter.Ip = new(types.CustomizationDhcpIpGenerator)
		default:
			nic.Adapter.Ip = &types.CustomizationFixedIp{IpAddress: ip}
		}

		if i < len(cmd.netmask) {
			nic.Adapter.SubnetMask = cmd.netmask[i]
		}
		if i < len(cmd.mac) {
			nic.MacAddress = cmd.mac[i]
		}
		if i < len(cmd.gateway) {
			nic.Adapter.Gateway = strings.Split(cmd.gateway[i], ",")
		}
		if isWindows {
			if i < len(cmd.dnsserver) {
				nic.Adapter.DnsServerList = strings.Split(cmd.dnsserver[i], ",")
			}
		}
	}

	for i, ip6 := range cmd.ip6 {
		ipconfig, err := parseIPv6Argument(ip6)
		if err != nil {
			return err
		}
		// use the same logic as the ip switch: the first occurrence of the ip6 switch is assigned to the first nic,
		// the second to the second nic and so forth.
		if spec.NicSettingMap == nil || len(spec.NicSettingMap) < i {
			return fmt.Errorf("unable to find a network adapter for IPv6 settings %d (%s)", i, ip6)
		}
		nic := &spec.NicSettingMap[i]
		if nic.Adapter.IpV6Spec == nil {
			nic.Adapter.IpV6Spec = new(types.CustomizationIPSettingsIpV6AddressSpec)
		}
		nic.Adapter.IpV6Spec.Ip = append(nic.Adapter.IpV6Spec.Ip, ipconfig...)
	}

	task, err := vm.Customize(ctx, *spec)
	if err != nil {
		return err
	}

	return task.Wait(ctx)
}
