package kernel

import (
	"context"
	"cos.googlesource.com/cos/tools.git/src/pkg/cos"
	"cos.googlesource.com/cos/tools.git/src/pkg/fs"
	"fmt"
	"github.com/golang/glog"
	"os"
	"os/exec"
	"path"
)

// InstallCompilerToolchain installs the COS compiler toolchain in a target directory.
//
// Once the toolchain has been installed, the toolchain tarball is removed to save
// space, and a toolchain-installed.txt file is created. If the installation fails for
// any reason, toolchain-installed.txt file will not be written, so it can be used as a
// reliable indicator that the installation succeeded.
func InstallCompilerToolchain(ctx context.Context, downloader cos.ArtifactsDownloader, dstDir string) error {
	if err := downloader.DownloadToolchainEnv(ctx, dstDir); err != nil {
		return err
	}

	if err := downloader.DownloadToolchain(ctx, dstDir); err != nil {
		return err
	}

	toolchainDir := path.Join(dstDir, "toolchain")
	if err := os.MkdirAll(toolchainDir, 0777); err != nil {
		return fmt.Errorf("error creating directory %s: %v", toolchainDir, err)
	}

	toolchainPath := path.Join(dstDir, "toolchain.tar.xz")
	if err := UnpackCompilerToolchain(toolchainPath, toolchainDir); err != nil {
		return err
	}

	// the toolchain tarball is not needed after it's extracted, so we remove it
	if err := os.Remove(toolchainPath); err != nil {
		return fmt.Errorf("error removing toolchain tarball: %v", err)
	}

	// write a file to indicate that the installation is finished
	donePath := path.Join(dstDir, "toolchain-installed.txt")
	if _, err := os.Create(donePath); err != nil {
		return fmt.Errorf("error finalizing toolchain installation: %v", err)
	}

	return nil
}

// UnpackCompilerToolchain unpacks a compiler toolchain to a target directory.
// This omits the rust toolchain components since they are unused by COS and use
// a significant amount of space.
func UnpackCompilerToolchain(toolchainPath, dstDir string) error {
	cmd := exec.Command(
		"tar", "xf", toolchainPath, "-C", dstDir,
		"--exclude=./usr/lib64/rustlib*",
		"--exclude=./usr/lib/rustlib*",
		"--exclude=./usr/lib64/libstd-*.so",
		"--exclude=./lib/libstd-*.so",
		"--exclude=./lib/librustc*",
		"--exclude=./usr/lib64/librustc*",
	)
	out, err := cmd.CombinedOutput()
	if err == nil {
		glog.V(1).Infof(string(out))
	} else {
		glog.Errorf(string(out))
		return fmt.Errorf("error unpacking compiler toolchain: %v", err)
	}

	return nil
}

// CompilerToolchainInstalled returns whether or not the toolchain has been
// installed in the target directory.
func CompilerToolchainInstalled(toolchainDir string) bool {
	donePath := path.Join(toolchainDir, "toolchain-installed.txt")
	return fs.IsFile(donePath)
}

// InstallKernelHeaders installs the COS kernel headers in a target directory.
//
// Once the headers have been installed, the header tarball is removed to save
// space, and a kernel-headers-installed.txt file is created. If the installation
// fails for any reason, the kernel-headers-installed.txt file will not be written,
// so it can be used as a reliable indicator that the installation succeeded.
func InstallKernelHeaders(ctx context.Context, downloader cos.ArtifactsDownloader, dstDir string) error {
	kernelHeadersPath := path.Join(dstDir, "kernel-headers.tgz")
	if err := downloader.DownloadKernelHeaders(ctx, dstDir); err != nil {
		return err
	}

	if err := UnpackKernelHeaders(kernelHeadersPath, dstDir); err != nil {
		return err
	}

	// the headers tarball is not needed after it's extracted, so we remove it
	if err := os.Remove(kernelHeadersPath); err != nil {
		return fmt.Errorf("error removing kernel headers tarball: %v", err)
	}

	// write a file to indicate that the installation is finished
	donePath := path.Join(dstDir, "kernel-headers-installed.txt")
	if _, err := os.Create(donePath); err != nil {
		return fmt.Errorf("error finalizing kernel header installation: %v", err)
	}

	return nil
}

// UnpackKernelHeaders unpacks kernel headers to a target directory.
func UnpackKernelHeaders(headersPath, dstDir string) error {
	cmd := exec.Command("tar", "xf", headersPath, "-C", dstDir, "--strip-components=4")
	out, err := cmd.CombinedOutput()
	if err == nil {
		glog.V(1).Infof(string(out))
	} else {
		glog.Errorf(string(out))
		return fmt.Errorf("error unpacking kernel headers: %v", err)
	}

	return nil
}

// KernelHeadersInstalled returns whether or not the kernel headers have been
// installed in the target directory.
func KernelHeadersInstalled(buildDir string) bool {
	donePath := path.Join(buildDir, "kernel-headers-installed.txt")
	return fs.IsFile(donePath)
}

// Installs the compiler toolchain atomically if it is not already installed.
// This function checks for a file which is only installed after an
// installation succeeds, meaning it will retry any partial installations.
func InstallCompilerToolchainAtomic(ctx context.Context, downloader cos.ArtifactsDownloader, dstDir string) error {
	toolchainInstalled := CompilerToolchainInstalled(dstDir)
	if !toolchainInstalled {
		glog.Infof("installing compiler toolchain")
		if err := InstallCompilerToolchain(ctx, downloader, dstDir); err != nil {
			return err
		}
		glog.Infof("done installing compiler toolchain")
	}

	return nil
}

// Installs the kernel headers atomically if they are not already installed.
// This function checks for a file which is only installed after an
// installation succeeds, meaning it will retry any partial installations.
func InstallKernelHeadersAtomic(ctx context.Context, downloader cos.ArtifactsDownloader, dstDir string) error {
	headersInstalled := KernelHeadersInstalled(dstDir)
	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
}