package dkms

import (
	"fmt"
	"os"
	"os/exec"
	"path"
	"sort"
	"strings"

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

// FindKernelModules returns a map from kernel module names to the paths to
// those modules for each module in a directory tree.
func FindKernelModules(root string) (map[string]string, error) {
	modules := make(map[string]string)
	if err := findKernelModules(root, modules); err != nil {
		return nil, err
	}
	return modules, nil
}

func findKernelModules(root string, modules map[string]string) error {
	entries, err := os.ReadDir(root)
	if err != nil {
		if os.IsNotExist(err) {
			glog.Infof("skipping kernel module search for directory %s because it does not exist: %v", root, err)
			return nil
		}

		return fmt.Errorf("could not read directory %s: %v", root, err)
	}

	for _, entry := range entries {
		name := entry.Name()
		if entry.IsDir() {
			if err := findKernelModules(path.Join(root, name), modules); err != nil {
				return err
			}
		} else if strings.HasSuffix(name, ".ko") {
			modules[name] = path.Join(root, name)
		}
	}

	return nil
}

// ModprobeModules modprobes a list of modules, ignoring ones which have
// already been inserted into the running kernel.
func ModprobeModules(modules []Module) error {
	insertedModules, err := InsertedModules()
	if err != nil {
		return fmt.Errorf("could not get list of inserted modules: %v; skipping modprobe", err)
	}

	for _, module := range modules {
		// Inserted module names only use underscores as separators
		insertedModuleName := strings.ReplaceAll(module.DestName, "-", "_")
		_, inserted := insertedModules[insertedModuleName]
		if inserted {
			glog.Infof("module %s is already inserted; skipping", module.DestName)
			continue
		}

		if err := ModprobeModule(module.DestName); err != nil {
			return err
		}
	}

	return nil
}

// Depmod runs depmod to recalculate module dependencies.
// This is required for modprobe to work properly the first time a module has
// been installed, or if its dependencies change.
func Depmod() error {
	cmd := exec.Command("depmod")
	out, err := cmd.CombinedOutput()
	if err == nil {
		glog.Info(string(out))
	} else {
		glog.Error(string(out))
		return fmt.Errorf("error running depmod to set up module dependencies: %v", err)
	}
	return nil
}

// ModprobeModule runs modprobe on a module to insert it into the running
// kernel.
func ModprobeModule(name string) error {
	cmd := exec.Command("modprobe", name)
	out, err := cmd.CombinedOutput()
	if err == nil {
		glog.Info(string(out))
	} else {
		glog.Error(string(out))
		return fmt.Errorf("error running modprobe with module %s: %v", name, err)
	}
	return nil
}

// InsertedModules returns the list of modules which have been inserted into
// the running kernel.
func InsertedModules() (map[string]bool, error) {
	cmd := exec.Command("lsmod")
	out, err := cmd.CombinedOutput()
	if err != nil {
		glog.Error(string(out))
		return nil, fmt.Errorf("error listing modules with lsmod: %v", err)
	}
	outStr := string(out)

	moduleNames := make(map[string]bool)
	lines := strings.Split(outStr, "\n")[1:]
	for _, line := range lines {
		fields := strings.Fields(line)
		if len(fields) > 0 {
			name := fields[0]
			moduleNames[name] = true
		}
	}

	return moduleNames, nil
}

// Insmod runs insmod on a module to insert it into the running
// kernel. Module params should be a (possibly empty) slice containing
// strings the form param1=value1.
func Insmod(path string, params []string) error {
	insmodArgs := append([]string{path}, params...)
	cmd := exec.Command("insmod", insmodArgs...)
	out, err := cmd.CombinedOutput()
	if err == nil {
		if len(out) > 0 {
			glog.Info(string(out))
		}
	} else {
		glog.Error(string(out))
		return fmt.Errorf("error running insmod for module at %s: %v", path, err)
	}
	return nil
}

// ModuleDependencies gets the list of dependencies for a module
// via a call to modinfo.
func ModuleDependencies(modulePath string) ([]string, error) {
	cmd := exec.Command("modinfo", "-F", "depends", modulePath)
	out, err := cmd.CombinedOutput()
	if err != nil {
		glog.Error(string(out))
		return nil, fmt.Errorf("error getting dependencies for module at %s: %v", modulePath, err)
	}

	outString := strings.TrimSpace(string(out))
	if len(outString) == 0 {
		return []string{}, nil
	}

	modules := strings.Split(outString, ",")
	return modules, nil
}

// InsertModules inserts a list of modules into the running kernel.
//
// The modules will be inserted in an order which respects their dependencies,
// as determined by modinfo, and which maintains the relative order of modules
// in the modules list as well as possible. See ModuleInsertionOrder for more
// detail.
//
// This also takes care not to insert or get dependencies for modules which
// are already inserted, which saves some time.
func InsertModules(modules []string, modulePaths map[string]string, moduleParams modules.ModuleParameters) error {
	insertedModules, err := InsertedModules()
	if err != nil {
		return fmt.Errorf("could not get list of inserted modules: %v; skipping module insertion", err)
	}

	dependencyGraph, err := ModuleDependencyGraph(modules, modulePaths)
	if err != nil {
		return fmt.Errorf("could not insert modules %s: %v", modules, err)
	}
	insertionOrder := ModuleInsertionOrder(modules, dependencyGraph)
	glog.Infof("planned module insertion order %v", insertionOrder)

	for _, module := range insertionOrder {
		// Inserted module names only use underscores as separators
		insertedModuleName := strings.ReplaceAll(module, "-", "_")
		if insertedModules[insertedModuleName] {
			// Some non-dkms modules or modules from a previous run may already be inserted
			glog.Infof("module %s already inserted; skipping", module)
			continue
		}

		if err := Insmod(modulePaths[module+".ko"], moduleParams[module]); err != nil {
			return err
		}

		insertedModules[insertedModuleName] = true
	}

	return nil
}

// ModuleDependencyGraph uses modinfo to compute a mapping from module names
// to the list of modules that must be inserted before a given module.
//
// For example, the output {a: [], b: [d], c: [a, b], d: []} indicates that
// module d must be inserted before module b, and modules a and b must be
// inserted before module c.
func ModuleDependencyGraph(modules []string, modulePaths map[string]string) (map[string][]string, error) {
	var stack = make([]string, len(modules))
	copy(stack, modules)
	graph := make(map[string][]string)

	for len(stack) > 0 {
		module := stack[len(stack)-1]
		stack = stack[:len(stack)-1]

		_, added := graph[module]
		if added {
			continue
		}

		modulePath, exists := modulePaths[module+".ko"]
		if !exists { // built-in modules do not have a path in the kernel modules tree
			modulePath = module
		}
		deps, err := ModuleDependencies(modulePath)
		if err != nil {
			return nil, fmt.Errorf("could not build module dependency graph: %v", err)
		}
		stack = append(stack, deps...)
		graph[module] = deps
	}

	return graph, nil
}

// ModuleInsertionOrder takes a list of modules and a graph of their
// dependencies and returns a valid order in which a list of modules can
// be inserted into the running kernel.
//
// The structure of the dependency graph is described in
// ModuleDependencyGraph.
//
// Where there is not a dependency preventing it, the ordered modules will be
// in the same relative order as in the slice provided. For example, if
// the module list is [a, b, c, d] and the dependency graph is
// {a: [], b: [d], c: [a, b], d: []}, then the returned ordering must be
// [a, d, b, c]}. While a different ordering like [d, b, a, c] would be valid
// with respect to the depencency graph, it would not be output because a
// comes before b in the given module list and there is no dependency forcing
// b to be inserted before a.
func ModuleInsertionOrder(modules []string, dependencyGraph map[string][]string) []string {
	preferredModuleInsertionOrder := make(map[string]int)
	for i, module := range modules {
		preferredModuleInsertionOrder[module] = i
	}

	var order []string
	added := make(map[string]bool)
	for _, module := range modules {
		addAllOrderedDependencies(module, dependencyGraph, preferredModuleInsertionOrder, &order, added)
	}

	return order
}

// addAllOrderedDependencies adds all of the transitive dependencies of an item
// in a dependency graph to a topological ordering.
//
// Ensures that no item is added twice, the ordering remains a topological
// order on the dependency graph, and that ties in the dependency graph are
// broken in favor of the item which is earlier in the preferred order.
func addAllOrderedDependencies(
	item string,
	dependencyGraph map[string][]string,
	preferredOrder map[string]int,
	order *[]string,
	added map[string]bool,
) {
	if added[item] {
		return
	}

	dependencies := dependencyGraph[item]
	sortByPreferredOrder(dependencies, preferredOrder)
	for _, dep := range dependencies {
		addAllOrderedDependencies(dep, dependencyGraph, preferredOrder, order, added)
	}

	*order = append(*order, item)
	added[item] = true
}

func sortByPreferredOrder(items []string, preferredOrder map[string]int) {
	sort.Slice(items, func(i, j int) bool {
		a, inserted := preferredOrder[items[i]]
		if !inserted {
			return false
		}
		b, inserted := preferredOrder[items[j]]
		if !inserted {
			return true
		}
		return a < b
	})
}