package dkms

import (
	"fmt"
	"os"
	"regexp"
	"sort"
	"strconv"
	"strings"

	"cos.googlesource.com/cos/tools.git/src/pkg/utils"
	"github.com/golang/glog"
)

// Config is the set of configurations for a DKMS package.
// See ParseConfig for details on how to write a dkms.conf file to
// populate these values.
type Config struct {
	// The name of the package.
	PackageName string
	// The version of the package.
	PackageVersion string
	// The make command which will be used to compile all of the
	// modules in the package.
	MakeCommand string
	// A list of patches to apply to the package before it is built.
	// This is constructed from the PATCH[#] and PATCH_MATCH[#] arrays.
	Patches []string
	// The list of Modules in the package.
	Modules []Module
	// The make command which will be used to clean the package when
	// it is unbuilt or removed.
	Clean string
	// Whether or not the package should be automatically installed when
	// the kernel is upgraded.
	Autoinstall bool
	// A list of packages which should be built before this one.
	// Currently, this is only an annotation and does not affect the
	// build order.
	BuildDepends []string
	// A regex which specifies the kernels for which this package can
	// be built.
	BuildExclusiveKernel string
	// A regex which specifies which arches for which this package can
	// be built.
	BuildExclusiveArch string
	// The name of a script to run after adding the package..
	PostAdd string
	// The name of a script to run after building the package.
	PostBuild string
	// The name of a script to run after installing the package.
	PostInstall string
	// The name of a script to run after removing the package.
	PostRemove string
	// The name of a script to run before building the package.
	PreBuild string
	// The name of a script to run before installing the package.
	PreInstall string
}

// LoadConfig reads a dkms.conf file for a package and parses its contents.
// See ParseConfig for more details.
func LoadConfig(pkg *Package) (*Config, error) {
	configPath := pkg.ConfigPath()
	var contents []byte
	var err error
	contents, err = os.ReadFile(configPath)
	if err != nil {
		if os.IsNotExist(err) {
			glog.V(2).Infof("could not find dkms.conf at %s; using default values", configPath)
		} else {
			return nil, err
		}
	}

	configText := string(contents)
	config, err := ParseConfig(configText, pkg)
	if err != nil {
		return nil, err
	}

	return config, nil
}

// ParseConfig parses the contents of a dkms.conf file for a package,
// applying defaults where required.
//
// The defaults are assigned as bash variables in a string which is
// prepended to the contents of dkms.conf. The resulting script is then
// sourced, and the shell variables after executing the script are
// collected into a Config.
//
// In addition, the following DKMS variables are available as bash
// variables in dkms.conf:
//   module: The package name (--module from the command line)
//   module_version: The package version (--module-version from the command line)
//   arch: The target architecture for the package.
//   build_id: The build id of the COS kernel for the package.
//   board: The board of the COS kernel for the package.
//   kernelver: The version of the kernel for the package.
//   dkms_tree: The path to the DKMS tree.
//   source_tree: The path to the user source tree.
//   kernel_source_dir: The path to the directory containing the COS kernel sources.
func ParseConfig(contents string, pkg *Package) (*Config, error) {
	trees := pkg.Trees
	buildDir := pkg.BuildDir()

	// To quote the original DKMS authors: 'This is really ugly, but a neat hack'.
	// By declaring the array variables as arrays initially, bash will always
	// treat them as arrays, even if a normal assignment happens to the variable
	// later. Thus, a user can assign to them and read from them as if they were
	// either arrays or normal variables in their dkms.conf, but we can always
	// read their values as array elements.
	defaults := map[string]string{
		"module":                 pkg.Name,
		"module_version":         pkg.Version,
		"arch":                   pkg.Arch,
		"build_id":               pkg.BuildId,
		"board":                  pkg.Board,
		"kernelver":              pkg.KernelVersion,
		"dkms_tree":              trees.Dkms,
		"source_tree":            trees.Source,
		"kernel_source_dir":      trees.Kernel,
		"MAKE":                   fmt.Sprintf("( 'make -C %s M=%s' )", trees.Kernel, buildDir),
		"MAKE_MATCH":             "( '.*' )",
		"CLEAN":                  fmt.Sprintf("'make -C %s M=%s clean'", trees.Kernel, buildDir),
		"PATCH":                  "()",
		"PATCH_MATCH":            "()",
		"BUILT_MODULE_NAME":      fmt.Sprintf("( %s )", pkg.Name),
		"BUILT_MODULE_LOCATION":  "()",
		"DEST_MODULE_NAME":       "()", // if this is still empty later, it will default to BUILT_MODULE_NAME
		"DEST_MODULE_LOCATION":   "( /kernel/updates )",
		"STRIP":                  "( yes )",
		"BUILD_DEPENDS":          "()",
		"PACKAGE_NAME":           pkg.Name,
		"PACKAGE_VERSION":        pkg.Version,
		"AUTOINSTALL":            "no",
		"BUILD_EXCLUSIVE_KERNEL": ".*",
		"BUILD_EXCLUSIVE_ARCH":   ".*",
	}
	var defaultsLines []string
	for variable, value := range defaults {
		defaultsLines = append(defaultsLines, fmt.Sprintf("%s=%s", variable, value))
	}
	sort.Strings(defaultsLines)
	defaultsString := strings.Join(defaultsLines, "\n")

	configString := strings.Join([]string{defaultsString, contents}, "\n")

	vars, err := utils.SourceString(configString)
	if err != nil {
		return nil, err
	}

	// Get the MAKE command.
	makeCommands := utils.ArrayElements(vars["MAKE"])
	makeMatches := utils.ArrayElements(vars["MAKE_MATCH"])
	makeCommand, err := matchMakeCommand(pkg.KernelVersion, makeCommands, makeMatches)
	if err != nil {
		return nil, fmt.Errorf("error while matching MAKE command: %v", err)
	}

	// Get the set of patches to apply.
	patchMap := utils.ArrayElements(vars["PATCH"])
	patchMatches := utils.ArrayElements(vars["PATCH_MATCH"])
	patches, err := selectPatches(pkg.KernelVersion, patchMap, patchMatches)
	if err != nil {
		return nil, fmt.Errorf("error while selecting patches to apply: %v", err)
	}

	// Get the list of build dependencies. These should be kept in the numerical order in
	// which they appear in the BUILD_DEPENDS array.
	buildDependsMap := utils.ArrayElements(vars["BUILD_DEPENDS"])
	var buildDependsKeys []string
	buildDepends := []string{}
	for key := range buildDependsMap {
		buildDependsKeys = append(buildDependsKeys, key)
	}
	err = sortStringsAsInts(buildDependsKeys)
	if err != nil {
		return nil, fmt.Errorf("error while sorting BUILD_DEPENDS keys: %v", err)
	}
	for _, key := range buildDependsKeys {
		buildDepends = append(buildDepends, buildDependsMap[key])
	}

	// Get the list of modules.
	builtModuleNames := utils.ArrayElements(vars["BUILT_MODULE_NAME"])
	builtModuleLocations := utils.ArrayElements(vars["BUILT_MODULE_LOCATION"])
	destModuleNames := utils.ArrayElements(vars["DEST_MODULE_NAME"])
	destModuleLocations := utils.ArrayElements(vars["DEST_MODULE_LOCATION"])
	strips := utils.ArrayElements(vars["STRIP"])
	modules, err := collateModules(pkg, builtModuleNames, builtModuleLocations, destModuleNames, destModuleLocations, strips)
	if err != nil {
		return nil, fmt.Errorf("error while collecting modules from config: %v", err)
	}

	autoinstall := vars["AUTOINSTALL"] == "yes"

	config := Config{
		PackageName:          vars["PACKAGE_NAME"],
		PackageVersion:       vars["PACKAGE_VERSION"],
		MakeCommand:          makeCommand,
		Patches:              patches,
		Modules:              modules,
		Clean:                vars["CLEAN"],
		Autoinstall:          autoinstall,
		BuildDepends:         buildDepends,
		BuildExclusiveKernel: vars["BUILD_EXCLUSIVE_KERNEL"],
		BuildExclusiveArch:   vars["BUILD_EXCLUSIVE_ARCH"],
		PostAdd:              vars["POST_ADD"],
		PostBuild:            vars["POST_BUILD"],
		PostInstall:          vars["POST_INSTALL"],
		PostRemove:           vars["POST_REMOVE"],
		PreBuild:             vars["PRE_BUILD"],
		PreInstall:           vars["PRE_INSTALL"],
	}
	return &config, nil
}

// matchMakeCommand returns the first make command in numerical order which matches
// the provided kernel version.
func matchMakeCommand(kernelVersion string, makeCommands, makeMatches map[string]string) (string, error) {
	if len(makeCommands) == 1 && len(makeMatches) == 0 {
		// return the first make command since there is only one
		for _, makeCommand := range makeCommands {
			return makeCommand, nil
		}
	}

	key, err := firstMatchingKey(kernelVersion, makeMatches)
	if err != nil {
		return "", fmt.Errorf("error while trying to match make commands to kernel version: %v", err)
	}

	makeCommand, ok := makeCommands[key]
	if !ok {
		return "", fmt.Errorf("kernelver matched MAKE_MATCH[%s], but no corresponding MAKE found", key)
	}

	return makeCommand, nil
}

// selectPatches returns the list of patches in numerical order which match the 
// provided kernel version.
func selectPatches(kernelVersion string, patchMap, patchMatches map[string]string) ([]string, error) {
	var keys []string
	for key := range patchMap {
		keys = append(keys, key)
	}
	err := sortStringsAsInts(keys)
	if err != nil {
		return nil, fmt.Errorf("invalid keys in patterns array: %v", err)
	}

	patches := []string{}
	for _, key := range keys {
		pattern, ok := patchMatches[key]
		// If there's a pattern, check the kernel version and skip the patch
		// if it doesn't match. If there's no pattern, we always apply the patch.
		if ok {
			regex, err := regexp.Compile(pattern)
			if err != nil {
				return nil, fmt.Errorf("could not compile regexp %s for key %s", pattern, key)
			}
			if !regex.MatchString(kernelVersion) {
				continue
			}
		}

		patch, ok := patchMap[key]
		if !ok {
			return nil, fmt.Errorf("kernelver matched PATCH_MATCH[%s], but no corresponding PATCH found", key)
		}
		patches = append(patches, patch)
	}

	return patches, nil
}

// firstMatchingKey takes a mapping from keys to patterns and returns the first key
// in ascending numerical order whose corresponding pattern matches a target string.
// Errors if any key cannot be converted to an integer or if any of the patterns
// cannot be compiled as a regexp.
func firstMatchingKey(text string, patterns map[string]string) (string, error) {
	var keys []string
	for key := range patterns {
		keys = append(keys, key)
	}
	err := sortStringsAsInts(keys)
	if err != nil {
		return "", fmt.Errorf("invalid keys in patterns array: %v", err)
	}

	var regexes []*regexp.Regexp
	for _, key := range keys {
		pattern := patterns[key]
		regex, err := regexp.Compile(pattern)
		if err != nil {
			return "", fmt.Errorf("could not compile regexp %s for key %s", pattern, key)
		}
		regexes = append(regexes, regex)
	}

	for i, key := range keys {
		if regexes[i].MatchString(text) {
			return key, nil
		}
	}

	return "", fmt.Errorf("no matching patterns found for text %s", text)
}

// sortStringsAsInts sorts a list of strings in ascending numerical order.
// Returns an error if any of the strings cannot be converted to an int.
func sortStringsAsInts(strings []string) error {
	var ints []int
	for _, s := range strings {
		intValue, err := strconv.Atoi(s)
		if err != nil {
			return fmt.Errorf("error converting %s to int when sorting strings: %v", s, err)
		}

		ints = append(ints, intValue)
	}

	sort.Ints(ints)
	for i, intValue := range ints {
		strings[i] = strconv.Itoa(intValue)
	}

	return nil
}

// collateModules takes the module-related arrays from dkms.conf and collates them to
// construct a list of Modules.
// This gathers the full set of keys from all of the arrays and sorts them in
// numeric order before iterating, which ensures that the output is deterministic.
func collateModules(pkg *Package, builtNames, builtLocations, destNames, destLocations, strips map[string]string) ([]Module, error) {
	moduleKeys := allKeys([]map[string]string{
		builtNames,
		builtLocations,
		destNames,
		destLocations,
		strips,
	})
	err := sortStringsAsInts(moduleKeys)
	if err != nil {
		return nil, fmt.Errorf("invalid keys in module arrays: %v", err)
	}

	var modules []Module
	for _, key := range moduleKeys {
		builtName := builtNames[key]
		builtLocation := builtLocations[key]

		destName := destNames[key]
		if destName == "" {
			destName = builtName
		}

		destLocation := destLocations[key]
		if destLocation == "" {
			destLocation = destLocations[moduleKeys[0]]
		}
		if !strings.HasPrefix(destLocation, "/kernel") {
			return nil, fmt.Errorf("DEST_MODULE_LOCATION %s must begin with /kernel", destLocation)
		}

		strip := strips[key]
		if strip == "" {
			strip = strips[moduleKeys[0]]
		}

		modules = append(modules, Module{
			Package:       pkg,
			BuiltName:     builtName,
			BuiltLocation: builtLocation,
			DestName:      destName,
			DestLocation:  destLocation,
			Strip:         strip != "no",
		})
	}

	if len(modules) == 0 {
		return nil, fmt.Errorf("package must contain at least one module, found 0")
	}

	return modules, nil
}

// allKeys returns the set of all keys from a list of maps.
// For example, allKeys([]map[string]int {{"a": "1", "b": "2"}, {"a": "3", "c": "4"}})
// will output []string{"a", "b", "c"} (though the order of keys is not guaranteed).
func allKeys(maps []map[string]string) []string {
	keySet := make(map[string]bool)
	for _, items := range maps {
		for key := range items {
			keySet[key] = true
		}
	}

	var keys []string
	for key := range keySet {
		keys = append(keys, key)
	}

	return keys
}