package main

// modify-variant-kernel-parameters adds or removes one entry in the
// "kernel-parameters" array of a Bottlerocket variant's Cargo.toml.
// Run it from the root of a variant repository before building, for example
// to change the kernel command line of an aarch64 build only:
//
//	./modify-variant-kernel-parameters -action add -variant aws-dev -kernel-parameter pci=realloc
//	./modify-variant-kernel-parameters -action remove -variant aws-dev -kernel-parameter quiet
//
// This program is not needed once twoliter supports kernel-parameters per CPU architecture.
// https://github.com/bottlerocket-os/twoliter/issues/587

import (
	"flag"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"slices"
	"strings"

	"github.com/BurntSushi/toml"
)

// KernelParameters represents the kernel-parameters in a variant's Cargo.toml
type KernelParameters struct {
	Package struct {
		Metadata struct {
			BuildVariant struct {
				KernelParameters []string `toml:"kernel-parameters"`
			} `toml:"build-variant"`
		} `toml:"metadata"`
	} `toml:"package"`
}

func main() {
	var variant string
	var action string
	var kernelParameter string

	flag.StringVar(&variant, "variant", "", "Variant name (e.g., aws-dev)")
	flag.StringVar(&action, "action", "", "Action to perform: 'add' or 'remove'")
	flag.StringVar(&kernelParameter, "kernel-parameter", "", "Kernel parameter to add or remove")

	flag.Parse()

	if variant == "" {
		flag.Usage()
		log.Fatal("Error: -variant parameter is required")
	}

	if action == "" {
		flag.Usage()
		log.Fatal("Error: -action parameter is required")
	}

	if action != "add" && action != "remove" {
		flag.Usage()
		log.Fatal("Error: -action must be 'add' or 'remove'")
	}

	if kernelParameter == "" {
		flag.Usage()
		log.Fatal("Error: -kernel-parameter parameter is required")
	}

	cargoPath := filepath.Join("variants", variant, "Cargo.toml")
	content, err := os.ReadFile(cargoPath)
	if err != nil {
		log.Fatalf("Error reading file: %v", err)
	}

	// Parse current kernel parameters to validate the file structure
	var config KernelParameters
	if _, err := toml.Decode(string(content), &config); err != nil {
		log.Fatalf("Error parsing Cargo.toml: %v", err)
	}

	currentParams := config.Package.Metadata.BuildVariant.KernelParameters

	var newParams []string
	var modified bool

	switch action {
	case "add":
		newParams, modified = addKernelParameter(currentParams, kernelParameter)
	case "remove":
		newParams, modified = removeKernelParameter(currentParams, kernelParameter)
	}

	if !modified {
		if action == "add" {
			log.Printf("Kernel parameter '%s' already exists in %s", kernelParameter, cargoPath)
		} else {
			log.Printf("Kernel parameter '%s' not found in %s", kernelParameter, cargoPath)
		}
		return
	}

	newContent := modifyKernelParametersInContent(string(content), newParams)

	// Write the modified content back to the file
	if err := os.WriteFile(cargoPath, []byte(newContent), 0644); err != nil {
		log.Fatalf("Error writing file: %v", err)
	}

	log.Printf("Successfully %s kernel parameter '%s' %s %s",
		map[string]string{"add": "added", "remove": "removed"}[action], kernelParameter,
		map[string]string{"add": "to", "remove": "from"}[action],
		cargoPath)
}

// addKernelParameter adds a kernel parameter if it doesn't already exist.
func addKernelParameter(currentParams []string, parameter string) ([]string, bool) {
	if slices.Contains(currentParams, parameter) {
		return currentParams, false
	}

	newParams := append(currentParams, parameter)
	return newParams, true
}

// removeKernelParameter removes a kernel parameter if it exists
func removeKernelParameter(currentParams []string, parameter string) ([]string, bool) {
	if !slices.Contains(currentParams, parameter) {
		return currentParams, false
	}

	newParams := slices.DeleteFunc(slices.Clone(currentParams), func(param string) bool {
		return param == parameter
	})

	return newParams, true
}

// modifyKernelParametersInContent modifies the kernel-parameters section while preserving formatting
// This is preferred over toml.Unmarshal/Encode because Unmarshal normalizes the TOML, which makes it difficult
// to see the diff of kernel-parameters.
func modifyKernelParametersInContent(content string, newParams []string) string {
	lines := strings.Split(content, "\n")
	var result []string
	inKernelParams := false

	// Regex to match the start of kernel-parameters
	kernelParamsRegex := regexp.MustCompile(`^\s*kernel-parameters\s*=\s*\[`)

	for _, line := range lines {
		if kernelParamsRegex.MatchString(line) {
			// Found the start of kernel-parameters
			inKernelParams = true

			// Check if it's a single-line array
			if strings.Contains(line, "]") {
				// Single-line format - replace the entire line
				newLine := formatKernelParametersMultiLine(newParams)
				result = append(result, newLine...)
				inKernelParams = false
			} else {
				// Multi-line format - start replacement
				newLine := formatKernelParametersMultiLine(newParams)
				result = append(result, newLine...)
				inKernelParams = true
			}
		} else if inKernelParams {
			// Skip lines until we find the closing bracket
			if strings.Contains(line, "]") {
				inKernelParams = false
			}
		} else {
			// Regular line, keep as-is
			result = append(result, line)
		}
	}

	return strings.Join(result, "\n")
}

// formatKernelParametersMultiLine formats kernel parameters in the original multi-line style
func formatKernelParametersMultiLine(params []string) []string {
	if len(params) == 0 {
		return []string{"kernel-parameters = []"}
	}

	var lines []string
	lines = append(lines, "kernel-parameters = [")

	for _, param := range params {
		lines = append(lines, "    \""+param+"\",")
	}

	lines = append(lines, "]")
	return lines
}