package kernel

import (
	"encoding/json"
	"fmt"
	"io/fs"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"

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

// A KernelVariantConfig is a configuration representing a COS kernel variant.
//
// The Build must be given as a build number, e.g., 19804.0.0.
// If unspecified, the ToolchainArch is the same as the Arch.
type KernelVariantConfig struct {
	Build         string
	Board         string
	Arch          string
	Toolchain     string
	ToolchainArch string
	Configs       []string
	Patches       []string
}

// ReadConfig reads a KernelVariantConfig from a .json file.
func ReadConfig(path string) (KernelVariantConfig, error) {
	var config KernelVariantConfig

	contents, err := os.ReadFile(path)
	if err != nil {
		return config, fmt.Errorf("could not read config from %s: %v", path, err)
	}

	if err := json.Unmarshal(contents, &config); err != nil {
		return config, fmt.Errorf("failed to unmarshal config %s: %v", path, err)
	}

	return config, nil
}

// kmakeArgs returns the list of arguments that should be passed to
// COS kernel make commands.
//
// This will return the COS kernel's default make arguments for the given
// kernel config, followed by any overrides passed via makeArgs.
func kmakeArgs(conf KernelVariantConfig, makeArgs []string) ([]string, error) {
	kernelMakeFlags, err := cos.KernelMakeFlags(conf.Arch, conf.ToolchainArch, conf.Toolchain)
	if err != nil {
		return nil, fmt.Errorf("could not determine kernel compilation environment: %v", err)
	}

	var defaultMakeArgs []string
	for key, value := range kernelMakeFlags {
		defaultMakeArgs = append(defaultMakeArgs, fmt.Sprintf("%s=%s", key, value))
	}
	// Order matters; user-provided makeArgs take precedence over defaults.
	makeArgs = append(defaultMakeArgs, makeArgs...)

	return makeArgs, nil
}

// Kmake will run the make command from the provided makeArgs using
// the default make arguments corresponding to the provided kernel
// variant config.
//
// If an argument is passed in makeArgs, it will take precedence over
// the corresponding default arg.
func Kmake(conf KernelVariantConfig, makeArgs []string) error {
	makeArgs, err := kmakeArgs(conf, makeArgs)
	if err != nil {
		return fmt.Errorf("failed to prepare make args: %v", err)
	}

	makeCmd := exec.Command("make", makeArgs...)
	makeCmd.Stdin = os.Stdin
	makeCmd.Stdout = os.Stdout
	makeCmd.Stderr = os.Stderr
	glog.Infof("[ running make command ] %v\n", makeCmd)
	if err := makeCmd.Run(); err != nil {
		return fmt.Errorf("error running make command: %v", err)
	}

	return nil
}

// kernelReleaseVersion gets the release string for the kernel in the current directory.
func kernelReleaseVersion(conf KernelVariantConfig) (string, error) {
	kernelReleaseArgs, err := kmakeArgs(conf, []string{"-s", "kernelrelease"})
	if err != nil {
		return "", fmt.Errorf("failed to prepare kernelrelease args: %v", err)
	}
	kernelReleaseCmd := exec.Command("make", kernelReleaseArgs...)
	out, err := kernelReleaseCmd.Output()
	if err != nil {
		return "", fmt.Errorf("could not determine kernel release: %v", err)
	}
	return strings.Trim(string(out), "\n"), nil
}

// BuildKernelTarball compiles the kernel and modules in the current directory and
// packages them into a tarball suitable for use with cos-kernel-deploy.
//
// This will also strip debug symbols from the modules.
func BuildKernelTarball(conf KernelVariantConfig, extraMakeArgs []string) (string, error) {
	imageTarget := "bzImage"
	if conf.Arch == "arm64" {
		imageTarget = "Image"
	}

	makeArgs := append([]string{"-j", imageTarget, "modules"}, extraMakeArgs...)
	if err := Kmake(conf, makeArgs); err != nil {
		return "", fmt.Errorf("failed to compile %s and modules: %v", imageTarget, err)
	}

	version, err := kernelReleaseVersion(conf)
	if err != nil {
		return "", err
	}

	buildRoot, err := os.MkdirTemp("", "build-root")
	if err != nil {
		return "", fmt.Errorf("failed to create temporary build root dir: %v", err)
	}
	defer os.RemoveAll(buildRoot)

	if err := os.Mkdir(path.Join(buildRoot, "boot"), 0777); err != nil {
		return "", fmt.Errorf("failed to create temporary build root boot dir: %v", err)
	}

	bootConfigPath := path.Join(buildRoot, "boot", "config-"+version)
	if err := utils.CopyFile(".config", bootConfigPath); err != nil {
		return "", err
	}

	imageTargetSrcPath := path.Join("arch", conf.Arch, "boot", imageTarget)
	imageTargetDstPath := path.Join(buildRoot, "boot", "vmlinuz-"+version)
	if err := utils.CopyFile(imageTargetSrcPath, imageTargetDstPath); err != nil {
		return "", err
	}

	installModulesArgs := []string{"INSTALL_MOD_PATH=" + buildRoot, "modules_install", "-j"}
	if err := Kmake(conf, installModulesArgs); err != nil {
		return "", fmt.Errorf("could not install modules in build root dir: %v", err)
	}

	kernelMakeFlags, err := cos.KernelMakeFlags(conf.Arch, conf.ToolchainArch, conf.Toolchain)
	if err != nil {
		return "", fmt.Errorf("could not determine kernel compilation environment: %v", err)
	}
	objcopy := kernelMakeFlags["OBJCOPY"]
	if err := stripModuleDebugSymbols(buildRoot, objcopy); err != nil {
		return "", fmt.Errorf("failed to strip module debug symbols: %v", err)
	}

	tarName := fmt.Sprintf("cos-kernel-%s-%s.txz", version, conf.Arch)
	tarCmd := exec.Command(
		"tar", "-c", "-J", "-v",
		"-f", tarName, "-C", buildRoot,
		"boot/", "lib/",
	)
	tarCmd.Stdout = os.Stdout
	tarCmd.Stderr = os.Stderr
	if err := tarCmd.Run(); err != nil {
		return "", fmt.Errorf("failed to make kernel tarball: %v", err)
	}

	return tarName, nil
}

// stripModuleDebugSymbols removes the debug symbols from all kernel modules under
// the provided build root.
//
// Takes a path to an objcopy binary which will be used to do the stripping.
func stripModuleDebugSymbols(buildRoot, objcopy string) error {
	modulesDir := path.Join(buildRoot, "lib", "modules")
	err := filepath.Walk(modulesDir, func(path string, info fs.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if strings.HasSuffix(path, ".ko") {
			stripDebugCmd := exec.Command(objcopy, "--strip-debug", path)
			stripDebugCmd.Stdout = os.Stdout
			stripDebugCmd.Stderr = os.Stderr
			if err := stripDebugCmd.Run(); err != nil {
				return fmt.Errorf("failed to strip debug symbols from %s: %v", path, err)
			}
		}
		return nil
	})

	if err != nil {
		return fmt.Errorf("could not find built kernel modules: %v", err)
	}

	return nil
}