// Copyright 2023 The Inspektor Gadget authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package image

import (
	"context"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"slices"
	"strconv"
	"strings"
	"time"

	"github.com/moby/moby/api/types/container"
	"github.com/moby/moby/api/types/mount"
	"github.com/moby/moby/client"
	"github.com/moby/moby/client/pkg/jsonmessage"
	"github.com/spf13/cobra"
	"golang.org/x/term"
	"gopkg.in/yaml.v3"

	"github.com/inspektor-gadget/inspektor-gadget/cmd/common"
	"github.com/inspektor-gadget/inspektor-gadget/pkg/oci"
)

// It can be overridden at build time
var builderImage = "ghcr.io/inspektor-gadget/gadget-builder:main"

const (
	DEFAULT_EBPF_SOURCE = "program.bpf.c"
	DEFAULT_WASM        = "" // Wasm is optional; unset by default
	DEFAULT_METADATA    = "gadget.yaml"
)

type buildFile struct {
	EBPFSource string `yaml:"ebpfsource"`
	Wasm       string `yaml:"wasm"`
	Metadata   string `yaml:"metadata"`
}

type cmdOpts struct {
	path             string
	file             string
	fileChanged      bool
	image            string
	local            bool
	outputDir        string
	builderImage     string
	builderImagePull string
	updateMetadata   bool
	validateMetadata bool
	btfgen           bool
	btfhubarchive    string
}

func NewBuildCmd() *cobra.Command {
	opts := &cmdOpts{}

	cmd := &cobra.Command{
		Use:          "build PATH",
		Short:        "Build a gadget image",
		SilenceUsage: true,
		Args:         cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			if opts.local && opts.builderImage != builderImage {
				return fmt.Errorf("--local and --builder-image cannot be used at the same time")
			}

			// Validate builderImagePull flag by checking it against a list of valid values
			validValues := []string{"always", "missing", "never"}
			if !slices.Contains(validValues, opts.builderImagePull) {
				return fmt.Errorf("invalid value for --builder-image-pull: %s. Valid values are %s", opts.builderImagePull, strings.Join(validValues, ","))
			}

			fFlag := cmd.Flags().Lookup("file")
			opts.fileChanged = fFlag.Changed

			opts.path = args[0]

			return runBuild(cmd, opts)
		},
	}

	cmd.Flags().StringVarP(&opts.file, "file", "f", "build.yaml", "Path to build.yaml")
	cmd.Flags().BoolVarP(&opts.local, "local", "l", false, "Build using local tools")
	cmd.Flags().StringVarP(&opts.outputDir, "output", "o", "", "Path to a folder to store generated files while building")
	cmd.Flags().StringVarP(&opts.image, "tag", "t", "", "Name for the built image (format name:tag)")
	cmd.Flags().StringVar(&opts.builderImage, "builder-image", builderImage, "Builder image to use")
	cmd.Flags().StringVar(&opts.builderImagePull, "builder-image-pull", "always", "Specify when the builder image should be pulled [always, missing, never]")
	cmd.Flags().BoolVar(&opts.updateMetadata, "update-metadata", false, "Update the metadata according to the eBPF code")
	cmd.Flags().BoolVar(&opts.validateMetadata, "validate-metadata", true, "Validate the metadata file before building the gadget image")

	cmd.Flags().BoolVar(&opts.btfgen, "btfgen", false, "Enable btfgen")
	cmd.Flags().StringVar(&opts.btfhubarchive, "btfhub-archive", "", "Path to the location of the btfhub-archive files")

	return cmd
}

type buildOptions struct {
	outputDir         string
	ebpfSourcePath    string
	wasmSourcePath    string
	btfHubArchivePath string
	btfgen            bool
	forceColorsFlag   bool
	useInTreeHeaders  bool
}

func isFlag(str string) bool {
	return strings.HasPrefix(str, "-")
}

func isResponseFile(str string) bool {
	// Response file starts with @, see:
	// https://gcc.gnu.org/wiki/Response_Files
	return strings.HasPrefix(str, "@")
}

// Ensure the string build options are path to avoid flags injection.
func checkBuildOptions(buildOpts buildOptions) error {
	for _, opt := range []string{buildOpts.outputDir, buildOpts.ebpfSourcePath, buildOpts.wasmSourcePath, buildOpts.btfHubArchivePath} {
		opt = strings.TrimSpace(opt)
		if isFlag(opt) {
			return fmt.Errorf("flags are not accepted: %q", opt)
		}

		if isResponseFile(opt) {
			return fmt.Errorf("response file are not accepted: %q", opt)
		}
	}

	return nil
}

func buildPipeline(opts buildOptions) ([]buildStep, error) {
	if err := checkBuildOptions(opts); err != nil {
		return nil, fmt.Errorf("wrong build option detected: %w", err)
	}

	var steps []buildStep

	if opts.ebpfSourcePath != "" {
		for _, arch := range []string{oci.ArchAmd64, oci.ArchArm64} {
			steps = append(steps, newClangCompileStep(arch, opts))
		}
	}

	if opts.wasmSourcePath != "" {
		switch {
		case strings.HasSuffix(opts.wasmSourcePath, ".go"):
			steps = append(steps, newGoBuildStep(opts))
		case strings.HasSuffix(opts.wasmSourcePath, ".rs"):
			steps = append(steps, newCargoBuildStep(opts))
		case strings.HasSuffix(opts.wasmSourcePath, ".wasm"):
			// Already compiled, nothing to do.
		default:
			return nil, fmt.Errorf("unsupported wasm file type: %s", opts.wasmSourcePath)
		}
	}

	if opts.btfgen {
		for _, arch := range []string{oci.ArchAmd64, oci.ArchArm64} {
			steps = append(steps, newBtfgenStep(arch, opts))
		}
	}

	return steps, nil
}

func runPipeline(runner commandRunner, steps []buildStep) error {
	for _, step := range steps {
		if err := step.run(runner); err != nil {
			return err
		}
	}
	return nil
}

func runBuild(cmd *cobra.Command, opts *cmdOpts) error {
	conf := &buildFile{
		EBPFSource: DEFAULT_EBPF_SOURCE,
		Wasm:       DEFAULT_WASM,
		Metadata:   DEFAULT_METADATA,
	}

	var buildContent []byte
	var err error

	if opts.fileChanged {
		buildContent, err = os.ReadFile(opts.file)
		if err != nil {
			return fmt.Errorf("reading build file: %w", err)
		}
	} else {
		// The user specified the path but not the file. Use the default file build.yaml
		buildContent, err = os.ReadFile(filepath.Join(opts.path, opts.file))
		if err != nil && !errors.Is(err, os.ErrNotExist) {
			return fmt.Errorf("reading build file: %w", err)
		}
	}

	if opts.btfgen && opts.btfhubarchive == "" {
		return errors.New("btfgen requires --btfhub-archive")
	}

	if opts.btfgen {
		cmd.Printf("btfgen is enabled, building will take a while...\n")
	}

	if err := yaml.Unmarshal(buildContent, conf); err != nil {
		return fmt.Errorf("unmarshaling build.yaml: %w", err)
	}

	if opts.outputDir != "" {
		if _, err := os.Stat(opts.outputDir); err != nil {
			return err
		}
	} else {
		// make a temp folder to store the build results
		tmpDir, err := os.MkdirTemp("", "gadget-build-")
		if err != nil {
			return fmt.Errorf("creating temp dir: %w", err)
		}
		defer os.RemoveAll(tmpDir)

		opts.outputDir = tmpDir
	}

	if opts.path != "." {
		cwd, err := os.Getwd()
		if err != nil {
			return fmt.Errorf("getting current directory: %w", err)
		}
		defer os.Chdir(cwd)

		if err := os.Chdir(opts.path); err != nil {
			return fmt.Errorf("changing directory: %w", err)
		}
	}

	if _, err := os.Stat(conf.EBPFSource); errors.Is(err, os.ErrNotExist) {
		if conf.EBPFSource != DEFAULT_EBPF_SOURCE {
			return fmt.Errorf("source file %q not found", conf.EBPFSource)
		}
		conf.EBPFSource = ""
	}

	var hasEBPFSource, hasMetadata, hasWasm, hasGo bool
	if conf.EBPFSource != "" {
		if _, err := os.Stat(conf.EBPFSource); err == nil {
			hasEBPFSource = true
		}
	}
	if conf.Metadata != "" {
		if _, err := os.Stat(conf.Metadata); err == nil {
			hasMetadata = true
		}
	}
	if conf.Wasm != "" {
		if _, err := os.Stat(conf.Wasm); err == nil {
			hasWasm = true
		}
	}

	goFolder := "./go"

	if _, err := os.Stat(goFolder); err == nil {
		_ = filepath.Walk(goFolder, func(path string, info os.FileInfo, err error) error {
			if err != nil {
				return err
			}
			if !info.IsDir() && strings.HasSuffix(info.Name(), ".go") {
				hasGo = true
				return filepath.SkipDir
			}
			return nil
		})
	}

	if !hasEBPFSource && !hasMetadata && !hasWasm && !hasGo {
		return fmt.Errorf("at least one of ebpf source (program.bpf.c), metadata (gadget.yaml), .go files (present in go folder) or wasm module is required")
	}

	if conf.EBPFSource != "" || conf.Wasm != "" {
		if opts.local {
			steps, err := buildPipeline(buildOptions{
				outputDir:         opts.outputDir,
				ebpfSourcePath:    conf.EBPFSource,
				wasmSourcePath:    conf.Wasm,
				btfHubArchivePath: opts.btfhubarchive,
				btfgen:            opts.btfgen,
			})
			if err != nil {
				return fmt.Errorf("building build pipeline: %w", err)
			}

			runner := &localRunner{verbose: common.Verbose}
			if common.Verbose {
				fmt.Printf("Build logs start:\n")
			}
			if err := runPipeline(runner, steps); err != nil {
				return fmt.Errorf("local build: %w", err)
			}
			if common.Verbose {
				fmt.Printf("Build logs end\n")
			}
		} else {
			if err := buildInContainer(opts, conf); err != nil {
				return err
			}
		}
	}

	// TODO: make this configurable?
	archs := []string{oci.ArchAmd64, oci.ArchArm64}
	objectsPaths := map[string]*oci.ObjectPath{}

	for _, arch := range archs {
		obj := &oci.ObjectPath{}

		if conf.EBPFSource != "" {
			obj.EBPF = filepath.Join(opts.outputDir, arch+".bpf.o")
		}

		// TODO: the same wasm file is provided for all architectures. Should we allow per-arch
		// wasm files?
		if strings.HasSuffix(conf.Wasm, ".wasm") {
			// User provided an already-built wasm file
			obj.Wasm = conf.Wasm
		} else if conf.Wasm != "" {
			// User provided a source file to build wasm from
			obj.Wasm = filepath.Join(opts.outputDir, "program.wasm")
		}

		if opts.btfgen {
			archClean := arch
			if arch == oci.ArchAmd64 {
				archClean = "x86_64"
			}

			obj.Btfgen = filepath.Join(opts.outputDir, fmt.Sprintf("btfs-%s.tar.gz", archClean))
		}

		objectsPaths[arch] = obj
	}

	buildOpts := &oci.BuildGadgetImageOpts{
		EBPFSourcePath:   conf.EBPFSource,
		ObjectPaths:      objectsPaths,
		MetadataPath:     conf.Metadata,
		UpdateMetadata:   opts.updateMetadata,
		ValidateMetadata: opts.validateMetadata,
	}

	if sourceDateEpoch, ok := os.LookupEnv("SOURCE_DATE_EPOCH"); ok {
		sde, err := strconv.ParseInt(sourceDateEpoch, 10, 64)
		if err != nil {
			return fmt.Errorf("invalid SOURCE_DATE_EPOCH: %w", err)
		}
		buildOpts.CreatedDate = time.Unix(sde, 0).UTC().Format(time.RFC3339)
	} else {
		buildOpts.CreatedDate = time.Now().Format(time.RFC3339)
	}

	desc, err := oci.BuildGadgetImage(context.TODO(), buildOpts, opts.image)
	if err != nil {
		return err
	}

	cmd.Printf("Successfully built %s\n", desc.String())

	return nil
}

func pullImage(ctx context.Context, cli *client.Client, imageReference string) error {
	fmt.Printf("Pulling builder image %s\n", imageReference)
	reader, err := cli.ImagePull(ctx, imageReference, client.ImagePullOptions{})
	if err != nil {
		return fmt.Errorf("pulling builder image: %w", err)
	}
	defer reader.Close()

	out := os.Stdout
	outFd := out.Fd()
	isTTY := term.IsTerminal(int(outFd))
	return jsonmessage.DisplayJSONMessagesStream(reader, out, outFd, isTTY, nil)
}

func isImageLocallyAvailable(ctx context.Context, cli *client.Client, imageReference string) (bool, error) {
	f := make(client.Filters).Add("reference", imageReference)

	result, err := cli.ImageList(ctx, client.ImageListOptions{Filters: f})
	if err != nil {
		return false, fmt.Errorf("listing images: %w", err)
	}

	for _, img := range result.Items {
		for _, tag := range img.RepoTags {
			if tag == imageReference {
				return true, nil
			}
		}
		for _, digest := range img.RepoDigests {
			if digest == imageReference {
				return true, nil
			}
		}
	}
	return false, nil
}

func ensureBuilderImage(ctx context.Context, cli *client.Client, builderImage string, builderImagePull string) error {
	switch builderImagePull {
	case "always":
		return pullImage(ctx, cli, builderImage)
	case "missing":
		localAvailable, err := isImageLocallyAvailable(ctx, cli, builderImage)
		if err != nil {
			return err
		}
		if !localAvailable {
			return pullImage(ctx, cli, builderImage)
		}
	case "never":
		localAvailable, err := isImageLocallyAvailable(ctx, cli, builderImage)
		if err != nil {
			return err
		}
		if !localAvailable {
			return fmt.Errorf("image %s is not available locally and pull is disabled", builderImage)
		}
	default:
		return fmt.Errorf("invalid --builder-image-pull value: %s", builderImagePull)
	}
	return nil
}

func buildInContainer(opts *cmdOpts, conf *buildFile) error {
	cwd, err := os.Getwd()
	if err != nil {
		return fmt.Errorf("getting current directory: %w", err)
	}

	ctx := context.TODO()
	cli, err := client.New(client.FromEnv)
	if err != nil {
		return fmt.Errorf("creating docker client: %w", err)
	}
	defer cli.Close()

	if err := ensureBuilderImage(ctx, cli, opts.builderImage, opts.builderImagePull); err != nil {
		return err
	}

	// where the gadget source code is mounted in the container
	gadgetSourcePath := "/work"
	pathHost := cwd

	inspektorGadetSrcPath := os.Getenv("IG_SOURCE_PATH")
	useInTreeHeaders := false
	if inspektorGadetSrcPath != "" {
		pathHost = inspektorGadetSrcPath
		// find the gadget relative path to the inspektor-gadget source
		if !strings.HasPrefix(cwd, inspektorGadetSrcPath) {
			return fmt.Errorf("the current directory %q is not under the inspektor-gadget source path %q", cwd, inspektorGadetSrcPath)
		}
		gadgetRelativePath := strings.TrimPrefix(cwd, inspektorGadetSrcPath)
		gadgetSourcePath = filepath.Join("/work", gadgetRelativePath)

		// use in-tree headers too
		useInTreeHeaders = true
	}

	buildOpts := buildOptions{
		outputDir:         "/out",
		btfgen:            opts.btfgen,
		btfHubArchivePath: "/btfhub-archive",
		useInTreeHeaders:  useInTreeHeaders,
	}

	if conf.Wasm != "" {
		buildOpts.wasmSourcePath = filepath.Join(gadgetSourcePath, conf.Wasm)
	}

	if conf.EBPFSource != "" {
		buildOpts.ebpfSourcePath = filepath.Join(gadgetSourcePath, conf.EBPFSource)
	}

	if term.IsTerminal(int(os.Stdout.Fd())) && term.IsTerminal(int(os.Stderr.Fd())) {
		buildOpts.forceColorsFlag = true
	}

	steps, err := buildPipeline(buildOpts)
	if err != nil {
		return fmt.Errorf("building build pipeline: %w", err)
	}

	// The work mount ReadOnly field is updated as false, to allow Cargo.lock to compiled in /work folder for rust source code.
	mounts := []mount.Mount{
		{
			Type:     mount.TypeBind,
			Target:   "/work",
			Source:   pathHost,
			ReadOnly: false,
		},
		{
			Type:   mount.TypeBind,
			Target: "/out",
			Source: opts.outputDir,
		},
	}

	if opts.btfgen {
		mounts = append(mounts, mount.Mount{
			Type:   mount.TypeBind,
			Target: "/btfhub-archive",
			Source: opts.btfhubarchive,
		})
	}

	resp, err := cli.ContainerCreate(
		ctx,
		client.ContainerCreateOptions{
			Config: &container.Config{
				Image:      opts.builderImage,
				Cmd:        []string{"sleep", "inf"},
				WorkingDir: gadgetSourcePath,
				User:       fmt.Sprintf("%d:%d", os.Getuid(), os.Getgid()),
			},
			HostConfig: &container.HostConfig{
				Mounts: mounts,
			},
		},
	)
	if err != nil {
		return fmt.Errorf("creating builder container: %w", err)
	}
	defer func() {
		if _, err := cli.ContainerRemove(ctx, resp.ID, client.ContainerRemoveOptions{Force: true}); err != nil {
			fmt.Printf("Failed to remove builder container: %s\n", err)
		}
	}()

	if _, err := cli.ContainerStart(ctx, resp.ID, client.ContainerStartOptions{}); err != nil {
		return fmt.Errorf("starting builder container: %w", err)
	}

	runner := &containerRunner{
		ctx:         ctx,
		cli:         cli,
		containerID: resp.ID,
		verbose:     common.Verbose,
	}

	if common.Verbose {
		fmt.Println("Build logs start:")
	}
	if err := runPipeline(runner, steps); err != nil {
		return fmt.Errorf("container build: %w", err)
	}
	if common.Verbose {
		fmt.Println("Build logs end")
	}

	return nil
}
