package dkms

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

	"cloud.google.com/go/storage"
	"cos.googlesource.com/cos/tools.git/src/pkg/cos"
	"cos.googlesource.com/cos/tools.git/src/pkg/fs"
	"cos.googlesource.com/cos/tools.git/src/pkg/gcs"
	"cos.googlesource.com/cos/tools.git/src/pkg/utils"
	"github.com/golang/glog"
	"google.golang.org/api/option"
)

// Build compiles a DKMS package using the MAKE command from dkms.conf,
// signs the compiled modules, and saves the results to the package's
// build directory in the DKMS tree.
//
// The function prioritizes signing with a Cloud KMS key if `options.CloudKMSKey` is set,
// else it looks for a local key.
// If the path to a private key and a certificate is specified in the environmental
// variables 'MODULES_SIGN_KEY' and 'MODULES_SIGN_CERT' respectively, or the files
// exist in '/var/lib/dkms/mok.key' and '/var/lib/dkms/mok.pub', then this would
// sign the compiled modules with the provided key and certificate.
//
// If the package is not already added to the DKMS source tree, this will try to
// add it.
//
// If options.InstallBuildDependencies is true, this will install the
// appropriate compiler toolchain and kernel headers for the kernel specified by
// the options fields.
//
// If any patches should be applied before building the package, as specified in
// dkms.conf, this will apply them.
//
// If options.MakeVariables is set to "cos-default", this will apply the values
// used for compiling the COS kernel and append them to the MAKE command. See
// DefaultMakeVariables for more detail.
func Build(pkg *Package, options *Options) error {
	// This must happen first in order to set the package config.
	if err := Add(pkg, options); err != nil {
		return err
	}

	if IsBuilt(pkg) {
		if !options.Force {
			glog.Info("package already built; skipping build")
			return nil
		}

		if err := Unbuild(pkg); err != nil {
			return fmt.Errorf("failed to unbuild package before forcing rebuild: %v", err)
		}
	}

	glog.Infof("building package %s-%s", pkg.Name, pkg.Version)
	if pkg.Config.PreBuild != "" {
		if err := utils.RunCommandString(pkg.BuildDir(), pkg.Config.PreBuild); err != nil {
			return fmt.Errorf("failed to run pre-build script: %v", err)
		}
	}

	if options.InstallBuildDependencies {
		if err := InstallBuildDependencies(pkg); err != nil {
			return err
		}
	}

	sourceDir := pkg.SourceDir()
	buildDir := pkg.BuildDir()

	glog.Infof("copying sources from %s to build dir %s", sourceDir, buildDir)
	if err := fs.CopyDir(sourceDir, buildDir, 0777); err != nil {
		return err
	}

	for _, patch := range pkg.Config.Patches {
		patchPath := path.Join(buildDir, "patches", patch)
		if err := ApplyPatch(buildDir, patchPath); err != nil {
			return fmt.Errorf("could not apply patch: %v", err)
		}
	}

	makeVariables := options.MakeVariables
	if makeVariables == "cos-default" {
		var err error
		makeVariables, err = DefaultMakeVariables(pkg)
		if err != nil {
			return err
		}
	}
	makeCommand := fmt.Sprintf("%s %s", pkg.Config.MakeCommand, makeVariables)
	makeCommand = strings.Trim(makeCommand, " ")

	if err := utils.RunCommandString(buildDir, makeCommand); err != nil {
		return fmt.Errorf("failed to compile modules: %v", err)
	}

	if err := sign(pkg.Config.Modules, options); err != nil {
		glog.Warningf("skipping module signing for package (%s): %v", pkg.Name, err)
	}

	if pkg.Config.PostBuild != "" {
		if err := utils.RunCommandString(buildDir, pkg.Config.PostBuild); err != nil {
			return fmt.Errorf("failed to run post-build script: %v", err)
		}
	}

	return nil
}

// DefaultMakeVariables returns the default make variables for a package.
//
// The values of CC and CXX are expected to be provided by sourcing the toolchain_env
// file in the kernel source tree. The values set in toolchain_env take precedence over
// the default CC and CXX values.
//
// The default values are:
//
//	ARCH=${arch}
//	CC=toolchain/bin/${arch}-cros-linux-gnu-clang
//	CXX=toolchain/bin/${arch}-cros-linux-gnu-clang++
//	LD=toolchain/bin/${arch}-cros-linux-gnu-ld.lld
//	STRIP=toolchain/bin/llvm-strip
//	OBJCOPY=toolchain/bin/llvm-objcopy
//	HOSTCC=${arch}-pc-linux-gnu-clang
//	HOSTCXX=${arch}-pc-linux-gnu-clang++
//	HOSTLD=${arch}-pc-linux-gnu-clang
func DefaultMakeVariables(pkg *Package) (string, error) {
	makeVariablesMap := map[string]string{
		// paths are relative to the kernel source tree
		"ARCH":    pkg.Arch,
		"CC":      fmt.Sprintf("toolchain/bin/%s-cros-linux-gnu-clang", pkg.Arch),
		"CXX":     fmt.Sprintf("toolchain/bin/%s-cros-linux-gnu-clang++", pkg.Arch),
		"LD":      fmt.Sprintf("toolchain/bin/%s-cros-linux-gnu-ld.lld", pkg.Arch),
		"STRIP":   "toolchain/bin/llvm-strip",
		"OBJCOPY": "toolchain/bin/llvm-objcopy",
		"HOSTCC":  fmt.Sprintf("%s-pc-linux-gnu-clang", pkg.Arch),
		"HOSTCXX": fmt.Sprintf("%s-pc-linux-gnu-clang++", pkg.Arch),
		"HOSTLD":  fmt.Sprintf("%s-pc-linux-gnu-clang", pkg.Arch),
	}

	toolchainEnvPath := path.Join(pkg.Trees.Kernel, "toolchain_env")
	toolchainEnv, err := utils.SourceFile(toolchainEnvPath)
	if err != nil {
		return "", fmt.Errorf("could not source toolchain_env to determine build toolchain: %v", err)
	}

	var makeVariablesList []string
	for key, value := range makeVariablesMap {
		toolchainValue, ok := toolchainEnv[key]
		if ok {
			value = fmt.Sprintf("toolchain/bin/%s", toolchainValue)
		}

		makeVariablesList = append(makeVariablesList, fmt.Sprintf("%s=%s", key, value))
	}
	sort.Strings(makeVariablesList)

	return strings.Join(makeVariablesList, " "), nil
}

// ApplyPatch applies a patch in p1 format to the files in a directory.
func ApplyPatch(dir, patchPath string) error {
	file, err := os.Open(patchPath)
	if err != nil {
		return fmt.Errorf("could not open patch file: %v", err)
	}

	cmd := exec.Command("patch", "-p1")
	cmd.Dir = dir
	cmd.Stdin = file
	out, err := cmd.CombinedOutput()
	if err == nil {
		glog.Info(string(out))
	} else {
		glog.Error(string(out))
		return err
	}

	return nil
}

// InstallBuildDependencies downloads and installs the kernel headers and
// compiler toolchain for a package, if they are not already present.
func InstallBuildDependencies(pkg *Package) error {
	ctx := context.Background()
	client, err := storage.NewClient(ctx, option.WithoutAuthentication())
	if err != nil {
		return err
	}

	// Artifact paths look like gs://cos-tools/18244.151.84/lakitu/${artifact-path}
	prefix := path.Join(pkg.BuildId, pkg.Board)
	downloader := cos.NewGCSDownloader(client, nil, "", prefix, "", "")

	return installBuildDependencies(ctx, downloader, pkg.Trees.Kernel)
}

func installBuildDependencies(ctx context.Context, downloader cos.ArtifactsDownloader, dstDir string) error {
	toolchainInstalled := CompilerToolchainInstalled(dstDir)
	headersInstalled := KernelHeadersInstalled(dstDir)
	if !toolchainInstalled || !headersInstalled {
		if err := os.MkdirAll(dstDir, 0777); err != nil {
			return err
		}

		if !toolchainInstalled {
			glog.Infof("installing compiler toolchain")
			if err := InstallCompilerToolchain(ctx, downloader, dstDir); err != nil {
				return err
			}
			glog.Infof("done installing compiler toolchain")
		}

		if !headersInstalled {
			glog.Infof("installing kernel headers")
			if err := InstallKernelHeaders(ctx, downloader, dstDir); err != nil {
				return err
			}
			glog.Infof("done installing kernel headers")
		}
	}

	return nil
}

func sign(modules []Module, options *Options) error {
	var signer ModuleSigner
	var err error
	if options.CloudKMSKey != "" {
		kmsSigner, err := newKmsSigner(context.Background(), options.CloudKMSKey)
		if err != nil {
			return fmt.Errorf("failed to retrieve kms signer: %v", err)
		}
		defer kmsSigner.client.Close()
		signer = kmsSigner
	} else {
		signer, err = newLocalSigner(options.PrivateKeyPath)
		if err != nil {
			return fmt.Errorf("failed to retrieve local signer: %v", err)
		}
	}

	if err := SignModules(modules, options.CertificatePath, options.Hash, signer); err != nil {
		return err
	}
	return nil
}

// CachedBuild tries to download all of a package's built modules from the cache,
// and falls back to building them locally if they are not present in the cache.
//
// If options.Upload is specified and the package's built modules are not present
// in the cache, then this will upload them after they are built.
//
// See Build for more information on how modules are built locally.
func CachedBuild(ctx context.Context, pkg *Package, cache *gcs.GCSBucket, options *Options) error {
	err := CachedAdd(ctx, pkg, cache, options)
	if err != nil {
		return err
	}

	// Get the list of modules which are not available locally
	var moduleDownloads []gcs.ObjectDownload
	for _, module := range pkg.Config.Modules {
		builtPath := module.BuiltPath()
		if !fs.IsFile(builtPath) {
			glog.Infof("could not find compiled module %s locally", module.BuiltName)
			cacheBuiltPath := module.CacheBuiltPath()
			download := gcs.NewObjectDownload(cacheBuiltPath, builtPath)
			moduleDownloads = append(moduleDownloads, download)
		}
	}

	if len(moduleDownloads) > 0 {
		glog.Info("downloading missing modules from cache")

		if options.DownloadWorkers == 1 {
			err = cache.DownloadObjects(ctx, moduleDownloads)
		} else {
			err = cache.DownloadObjectsParallel(ctx, moduleDownloads, options.DownloadWorkers)
		}

		if err != nil {
			glog.Infof("error while downloading modules: %v", err)
			glog.Info("could not download some modules; building from scratch")
		}
	}

	if err := Build(pkg, options); err != nil {
		return err
	}

	// Upload the modules to the cache after they have all been built.
	if options.Upload && !IsBuiltInCache(ctx, pkg, cache) {
		for _, module := range pkg.Config.Modules {
			builtPath := module.BuiltPath()
			cacheBuiltPath := module.CacheBuiltPath()
			err := cache.UploadObjectFromFile(ctx, builtPath, cacheBuiltPath)
			if err != nil {
				return err
			}
		}
	}

	return nil
}

// IsBuilt returns whether or not all of a package's modules have been built.
func IsBuilt(pkg *Package) bool {
	if pkg.Config == nil {
		return false
	}

	for _, module := range pkg.Config.Modules {
		if !fs.IsFile(module.BuiltPath()) {
			return false
		}
	}

	return true
}

// IsBuiltInCache returns whether or not all of a package's built modules are
// present in a cache.
//
// If the package does not have its Config set, this returns false.
func IsBuiltInCache(ctx context.Context, pkg *Package, cache *gcs.GCSBucket) bool {
	if pkg.Config == nil {
		return false
	}

	for _, module := range pkg.Config.Modules {
		builtPath := module.CacheBuiltPath()
		exists, err := cache.Exists(ctx, builtPath)
		if !exists || err != nil {
			return false
		}
	}

	return true
}