package trivy

import (
	"context"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"

	"golang.org/x/xerrors"

	"github.com/pingcap-inc/security/baseline/pkg/utils"
)

func UpdateDB(opt *BaseOption) error {
	args := []string{
		"--cache-dir", opt.CacheDir,
		"image",
		"--timeout", opt.Timeout.String(),
		"--download-db-only",
	}

	dbDir := filepath.Join(opt.CacheDir, "db")
	if _, err := os.Stat(dbDir); os.IsNotExist(err) {
		err = os.MkdirAll(dbDir, 0o700)
		if err != nil {
			return xerrors.Errorf("Update trivy vulnerability database failed, %+v", err)
		}
	}

	lock := utils.NewDirLock(dbDir)
	if err := lock.Lock(); err != nil {
		return err
	}
	defer lock.Unlock()

	ctx, cancel := context.WithTimeout(context.Background(), opt.Timeout)
	defer cancel()

	cmd := exec.CommandContext(ctx, "trivy", args...)
	cmd.Env = os.Environ()
	cmd.Env = append(cmd.Env, fmt.Sprintf("GITHUB_TOKEN=%s", opt.GITHUB_TOKEN))

	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	err := cmd.Start()
	if err != nil {
		return xerrors.Errorf("Update trivy vulnerability database failed, %+v", err)
	}

	return cmd.Wait()
}

func ExecTrivyImage(opt *ImageOption) ([]byte, error) {
	args := []string{
		"--cache-dir", opt.CacheDir,
		"-q",
		"image",
		"-f", "json",
		"--timeout", opt.Timeout.String(),
		"--severity", opt.Severity,
	}
	args = append(args, opt.Targets...)

	ctx, cancel := context.WithTimeout(context.Background(), opt.Timeout)
	defer cancel()

	cmd := exec.CommandContext(ctx, "trivy", args...)
	cmd.Env = os.Environ()
	cmd.Env = append(cmd.Env, fmt.Sprintf("GITHUB_TOKEN=%s", opt.GITHUB_TOKEN))

	out, err := cmd.Output()
	if err != nil {
		return nil, xerrors.Errorf("Trivy image command execute failed, %+v", err)
	}

	return out, nil
}

func execTrivyImage(opt *ImageOption) error {
	args := []string{
		"--cache-dir", opt.CacheDir,
		"image",
		"--timeout", opt.Timeout.String(),
		"--severity", opt.Severity,
	}
	args = append(args, opt.Targets...)

	ctx, cancel := context.WithTimeout(context.Background(), opt.Timeout)
	defer cancel()

	cmd := exec.CommandContext(ctx, "trivy", args...)
	cmd.Env = os.Environ()
	cmd.Env = append(cmd.Env, fmt.Sprintf("GITHUB_TOKEN=%s", opt.GITHUB_TOKEN))

	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	err := cmd.Start()
	if err != nil {
		return xerrors.Errorf("Trivy image command execute failed, %+v", err)
	}

	return cmd.Wait()
}

func ScanImage(opt *ImageOption) error {
	if err := UpdateDB(opt.BaseOption); err != nil {
		return err
	}

	return execTrivyImage(opt)
}

func ExecTrivyRepo(opt *RepoOption) ([]byte, error) {
	args := []string{
		"--cache-dir", opt.CacheDir,
		"-q",
		"fs",
		"-f", "json",
		"--timeout", opt.Timeout.String(),
		"--severity", opt.Severity,
	}

	for _, target := range opt.Targets {
		lock, repoURL, err := LockRepo(opt.GITHUB_TOKEN, opt.CacheDir, target)
		if err != nil {
			return nil, err
		}
		defer lock.Unlock()

		if err := CloneRepo(opt.GITHUB_TOKEN, opt.CacheDir, repoURL, lock.Dir()); err != nil {
			return nil, err
		}

		if err := UpdateRepo(opt.GITHUB_TOKEN, opt.CacheDir, repoURL, lock.Dir()); err != nil {
			return nil, err
		}

		args = append(args, lock.Dir())
	}

	ctx, cancel := context.WithTimeout(context.Background(), opt.Timeout)
	defer cancel()

	cmd := exec.CommandContext(ctx, "trivy", args...)
	cmd.Env = os.Environ()
	cmd.Env = append(cmd.Env, fmt.Sprintf("GITHUB_TOKEN=%s", opt.GITHUB_TOKEN))

	out, err := cmd.Output()
	if err != nil {
		return nil, xerrors.Errorf("Trivy fs command execute failed, %+v", err)
	}

	return out, nil
}

func execTrivyRepo(opt *RepoOption) error {
	args := []string{
		"--cache-dir", opt.CacheDir,
		"fs",
		"--timeout", opt.Timeout.String(),
		"--severity", opt.Severity,
	}

	for _, target := range opt.Targets {
		lock, repoURL, err := LockRepo(opt.GITHUB_TOKEN, opt.CacheDir, target)
		if err != nil {
			return err
		}
		defer lock.Unlock()

		if err := CloneRepo(opt.GITHUB_TOKEN, opt.CacheDir, repoURL, lock.Dir()); err != nil {
			return err
		}

		if err := UpdateRepo(opt.GITHUB_TOKEN, opt.CacheDir, repoURL, lock.Dir()); err != nil {
			return err
		}

		args = append(args, lock.Dir())
	}

	ctx, cancel := context.WithTimeout(context.Background(), opt.Timeout)
	defer cancel()

	cmd := exec.CommandContext(ctx, "trivy", args...)
	cmd.Env = os.Environ()
	cmd.Env = append(cmd.Env, fmt.Sprintf("GITHUB_TOKEN=%s", opt.GITHUB_TOKEN))

	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	err := cmd.Start()
	if err != nil {
		return xerrors.Errorf("Trivy fs command execute failed, %+v", err)
	}

	return cmd.Wait()
}

func ScanRepo(opt *RepoOption) error {
	if err := UpdateDB(opt.BaseOption); err != nil {
		return err
	}

	return execTrivyRepo(opt)
}
