// fwprune is a buildsystem utility that filters linux-firmware repository
// contents to include only files required by the built-in kernel modules,
// that are specified in modules.builtin.modinfo.
// (see: https://www.kernel.org/doc/Documentation/kbuild/kbuild.txt)
package main

import (
	"bytes"
	"log"
	"os"
	"path/filepath"
	"sort"
	"strings"

	"google.golang.org/protobuf/encoding/prototext"

	"source.monogon.dev/metropolis/node/build/fsspec"
)

// fwPaths returns a slice of filesystem paths relative to the root of the
// linux-firmware repository, pointing at firmware files, according to contents
// of the kernel build side effect: modules.builtin.modinfo.
func fwPaths(mi []byte) []string {
	// Use a map pset to deduplicate firmware paths.
	pset := make(map[string]bool)
	// Get a slice of entries of the form "unix.license=GPL" from mi. Then extract
	// firmware information from it.
	entries := bytes.Split(mi, []byte{0})
	for _, entry := range entries {
		// Skip empty entries.
		if len(entry) == 0 {
			continue
		}
		// Parse the entries. Split each entry into a key-value pair, separated
		// by "=".
		kv := strings.SplitN(string(entry), "=", 2)
		key, value := kv[0], kv[1]
		// Split the key into a module.attribute] pair, such as "unix.license".
		ma := strings.SplitN(key, ".", 2)
		// Skip, if it's not a firmware entry, according to the attribute.
		if ma[1] != "firmware" {
			continue
		}
		// If it is though, value holds a firmware path.
		pset[value] = true
	}
	// Convert the deduplicated pset to a slice.
	pslice := make([]string, 0, len(pset))
	for p, _ := range pset {
		pslice = append(pslice, p)
	}
	sort.Strings(pslice)
	return pslice
}

// fwprune takes a modinfo file from the kernel and extracts a list of all
// firmware files requested by all modules in that file. It then takes all
// available firmware file paths (newline-separated in the firmwareList file)
// and tries to match the requested file paths as a suffix of them.
// For example if a module requests firmware foo/bar.bin and in the firmware list
// there is a file at path build-out/x/y/foo/bar.bin it will use that file.
// Finally it generates an fsspec placing each file under its requested path
// under /lib/firmware.
func main() {
	if len(os.Args) != 4 {
		log.Fatal("Usage: fwprune modules.builtin.modinfo firmwareListPath outSpec")
	}
	modinfo := os.Args[1]
	firmwareListPath := os.Args[2]
	outSpec := os.Args[3]

	allFirmwareData, err := os.ReadFile(firmwareListPath)
	if err != nil {
		log.Fatalf("Failed to read firmware source list: %v", err)
	}
	allFirmwarePaths := strings.Split(string(allFirmwareData), "\n")

	// Create a look-up table of all possible suffixes to their full paths as
	// this is much faster at O(n) than calling strings.HasSuffix for every
	// possible combination which is O(n^2).
	suffixLUT := make(map[string]string)
	for _, firmwarePath := range allFirmwarePaths {
		pathParts := strings.Split(firmwarePath, string(os.PathSeparator))
		for i := range pathParts {
			suffixLUT[filepath.Join(pathParts[i:len(pathParts)]...)] = firmwarePath
		}
	}

	// Get the firmware file paths used by modules according to modinfo data
	mi, err := os.ReadFile(modinfo)
	if err != nil {
		log.Fatalf("While reading modinfo: %v", err)
	}
	fwp := fwPaths(mi)

	var files []*fsspec.File

	for _, p := range fwp {
		sourcePath := suffixLUT[p]
		if sourcePath == "" {
			// This should not be fatal as sometimes linux-firmware cannot
			// ship all firmware usable by the kernel for mostly legal reasons.
			log.Printf("WARNING: Requested firmware %q not found", p)
			continue
		}
		files = append(files, &fsspec.File{
			Path:       filepath.Join("/lib/firmware", p),
			Mode:       0444,
			SourcePath: sourcePath,
		})
	}
	fsspecRaw, err := prototext.Marshal(&fsspec.FSSpec{File: files})
	if err := os.WriteFile(outSpec, fsspecRaw, 0644); err != nil {
		log.Fatalf("failed writing output: %v", err)
	}
}