package runner

import (
	"bytes"
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"strings"
	"testing"
)

// Debian is the debian based image used by the runner.
const Debian = "debian:qbee"

// OpenWRT is the OpenWRT based image used by the runner.
const OpenWRT = "openwrt:qbee"

// RHEL is the RHEL based image used by the runner.
const RHEL = "rhel:qbee"

// UnprivilegedUser is the unprivileged user used by the runner if the UNPRIVILEGED_USER environment variable is set to "true".
const UnprivilegedUser = "qbee-agent"

// New creates a new runner for the given test.
func New(t *testing.T) *Runner {
	return NewWithImage(t, Debian, false)
}

// NewOpenWRTRunner creates a new runner for the given test using the openwrt:qbee image.
func NewOpenWRTRunner(t *testing.T) *Runner {
	runner := NewWithImage(t, OpenWRT, false)
	runner.MustExec("mkdir", "-p", "/var/lock")
	return runner
}

// NewRHELRunner creates a new runner for the given test using the rhel:qbee image.
func NewRHELRunner(t *testing.T) *Runner {
	return NewWithImage(t, RHEL, false)
}

// NewPodmanRunner creates a new runner for the given test using the debian:qbee-podman image.
func NewPodmanRunner(t *testing.T) *Runner {
	return NewWithImage(t, Debian, true)
}

// NewWithImage creates a new runner for the given test using the given image.
func NewWithImage(t *testing.T, image string, privileged bool) *Runner {
	cmdArgs := []string{
		"run",
		"--rm",
	}

	if privileged {
		cmdArgs = append(cmdArgs, "--privileged")
	} else {
		cmdArgs = append(cmdArgs, []string{
			"-v", "/var/run/docker.sock:/var/run/docker.sock", // mount docker socket
			"-v", "/sys/fs/cgroup:/sys/fs/cgroup:ro", // mount cgroup for docker
			"--tmpfs", "/tmp",
			"--tmpfs", "/run",
			"--tmpfs", "/run/lock",
			"--cap-add=NET_ADMIN", // allow control of firewall
		}...)
	}

	cmdArgs = append(cmdArgs, []string{
		"--detach", // launch in background
		image,
		"sleep", "600", // force exit container after 10 minutes
	}...)

	output, err := exec.Command("docker", cmdArgs...).Output()
	if err != nil {
		panic(err)
	}

	container := strings.TrimSpace(string(output))

	runner := &Runner{
		t:         t,
		container: container,
		image:     image,
	}

	if os.Getenv("UNPRIVILEGED_USER") == "true" {
		runner.WithUnprivileged()
	}

	t.Cleanup(runner.Close)

	return runner
}

// Runner provides a convenient way to run the agent in a container.
type Runner struct {
	t                *testing.T
	container        string
	image            string
	execUnPrivileged bool
}

// WithUnprivileged sets whether commands should be executed unprivileged.
func (runner *Runner) WithUnprivileged() *Runner {
	runner.execUnPrivileged = true
	return runner
}

// GetStateDirectory returns the state directory path for the runner.
func (runner *Runner) GetStateDirectory() string {
	if runner.execUnPrivileged {
		return "/home/" + UnprivilegedUser + "/var"
	}
	return "/var/lib/qbee"
}

// GetUnprivileged returns whether commands are executed unprivileged.
func (runner *Runner) GetUnprivileged() bool {
	return runner.execUnPrivileged
}

// Close kills the container.
func (runner *Runner) Close() {
	_ = exec.Command("docker", "kill", runner.container).Run()
}

// Pause processes within the runner container.
func (runner *Runner) Pause() {
	_ = exec.Command("docker", "pause", runner.container).Run()
}

// Unpause processes within the runner container.
func (runner *Runner) Unpause() {
	_ = exec.Command("docker", "unpause", runner.container).Run()
}

// PackageInstallCommand returns the package manager install command for the runner image.
func (runner *Runner) PackageInstallCommand(pkgName, version string) []string {
	switch runner.image {
	case Debian:
		if version != "" {
			pkgName = fmt.Sprintf("%s=%s", pkgName, version)
		}

		return []string{"apt-get", "install", "-y", pkgName}
	case OpenWRT:
		if version != "" {
			pkgName = fmt.Sprintf("%s=%s", pkgName, version)
		}

		return []string{"opkg", "install", pkgName}
	case RHEL:
		if version != "" {
			pkgName = fmt.Sprintf("%s-%s", pkgName, version)
		}

		return []string{"yum", "install", "-y", pkgName}
	default:
		panic("unsupported image")
	}
}

// FullUpdateCommand returns the package manager full update command for the runner image.
func (runner *Runner) FullUpdateCommand() [][]string {
	switch runner.image {
	case Debian:
		return [][]string{
			{"apt-get", "update"},
			{"apt-get", "upgrade", "-y"},
		}
	case OpenWRT:
		return [][]string{
			{"opkg", "update"},
			{"opkg", "upgrade"},
		}
	case RHEL:
		return [][]string{
			{"yum", "update", "-y"},
		}
	default:
		panic("unsupported image")
	}
}

// Exec executes the given command in the runner container.
// It returns the output of the command and an error if the command failed.
func (runner *Runner) Exec(cmd ...string) ([]byte, error) {
	execCommand := append([]string{"exec", runner.container}, cmd...)

	execCmd := exec.Command("docker", execCommand...)
	stderr := new(bytes.Buffer)
	execCmd.Stderr = stderr
	output, err := execCmd.Output()

	for _, line := range strings.Split(stderr.String(), "\n") {
		if line = strings.TrimSpace(line); line != "" {
			fmt.Println(line)
		}
	}

	return bytes.TrimSpace(output), err
}

// MustExec executes the given command and panics if it fails.
func (runner *Runner) MustExec(cmd ...string) []byte {
	output, err := runner.Exec(cmd...)
	if err != nil {
		if len(output) > 0 {
			fmt.Println("stdout:", string(output))
		}

		panic(err)
	}

	return output
}

// MustExecUnprivileged executes the given command unprivileged if the runner is set to unprivileged mode.
func (runner *Runner) MustExecUnprivileged(cmd ...string) []byte {
	if runner.execUnPrivileged {
		suCmd := append([]string{"su", "-s", "/bin/sh", UnprivilegedUser, "-c"}, strings.Join(cmd, " "))
		return runner.MustExec(suCmd...)
	}

	return runner.MustExec(cmd...)
}

// CreateFile creates a file with the given path and contents.
func (runner *Runner) CreateFile(path string, contents []byte) {
	fd, err := os.CreateTemp(runner.t.TempDir(), "qbee-test-*")
	if err != nil {
		panic(err)
	}

	defer func() { _ = os.Remove(fd.Name()) }()

	_, err = fd.Write(contents)

	_ = fd.Close()

	if err != nil {
		panic(err)
	}

	containerPath := fmt.Sprintf("%s:%s", runner.container, path)

	if err = exec.Command("docker", "cp", fd.Name(), containerPath).Run(); err != nil {
		panic(err)
	}
}

// CreateFileWithPerms creates a file with the given path and contents that is world-readable.
func (runner *Runner) CreateFileWithPerms(path string, contents []byte, perms os.FileMode) {
	runner.CreateFile(path, contents)

	// make sure the file can be read by non-root users
	if err := exec.Command("docker", "exec", runner.container, "chmod", fmt.Sprintf("%o", perms), path).Run(); err != nil {
		panic(err)
	}
}

// ReadFile reads the contents of the given file from the runner.
func (runner *Runner) ReadFile(path string) []byte {
	fd, err := os.CreateTemp(runner.t.TempDir(), "qbee-test-*")
	if err != nil {
		panic(err)
	}
	_ = fd.Close()

	defer func() { _ = os.Remove(fd.Name()) }()

	containerPath := fmt.Sprintf("%s:%s", runner.container, path)

	if err = exec.Command("docker", "cp", containerPath, fd.Name()).Run(); err != nil {
		panic(err)
	}

	var output []byte
	if output, err = os.ReadFile(fd.Name()); err != nil {
		panic(err)
	}

	return output
}

// CreateJSON creates a file with the given path and JSON-encodes the given document.
func (runner *Runner) CreateJSON(path string, doc any) {
	docBytes, err := json.Marshal(doc)
	if err != nil {
		panic(err)
	}

	runner.CreateFile(path, docBytes)
}
