/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package containers_test

import (
	"context"
	"io/ioutil"
	"os"
	"sync"
	"testing"

	"github.com/containerd/containerd"
	containerdcntrs "github.com/containerd/containerd/containers"
	"github.com/containerd/containerd/namespaces"
	"github.com/containerd/containerd/oci"
	"github.com/stretchr/testify/suite"

	"github.com/talos-systems/talos/internal/app/init/pkg/system/events"
	"github.com/talos-systems/talos/internal/app/init/pkg/system/runner"
	containerdrunner "github.com/talos-systems/talos/internal/app/init/pkg/system/runner/containerd"
	"github.com/talos-systems/talos/internal/app/init/pkg/system/runner/process"
	"github.com/talos-systems/talos/internal/pkg/constants"
	"github.com/talos-systems/talos/internal/pkg/containers"
	"github.com/talos-systems/talos/pkg/userdata"
)

const (
	containerdNamespace = "inspecttest"
	busyboxImage        = "docker.io/library/busybox:latest"
)

func MockEventSink(state events.ServiceState, message string, args ...interface{}) {
}

// nolint: maligned
type ContainersSuite struct {
	suite.Suite

	tmpDir string

	containerdRunner runner.Runner
	containerdWg     sync.WaitGroup

	client *containerd.Client
	image  containerd.Image

	containerRunners []runner.Runner
	containersWg     sync.WaitGroup
}

// WithAnnotations appends or replaces the annotations on the spec with the
// provided annotations
//
// TODO: taken from containerd/oci > 1.2.6
func WithAnnotations(annotations map[string]string) oci.SpecOpts {
	return func(_ context.Context, _ oci.Client, _ *containerdcntrs.Container, s *oci.Spec) error {
		if s.Annotations == nil {
			s.Annotations = make(map[string]string)
		}
		for k, v := range annotations {
			s.Annotations[k] = v
		}
		return nil
	}
}

// nolint: dupl
func (suite *ContainersSuite) SetupSuite() {
	var err error

	args := &runner.Args{
		ID:          "containerd",
		ProcessArgs: []string{"/rootfs/bin/containerd"},
	}

	suite.tmpDir, err = ioutil.TempDir("", "talos")
	suite.Require().NoError(err)

	suite.containerdRunner = process.NewRunner(
		&userdata.UserData{},
		args,
		runner.WithLogPath(suite.tmpDir),
		runner.WithEnv([]string{"PATH=/rootfs/bin:" + constants.PATH}),
	)
	suite.Require().NoError(suite.containerdRunner.Open(context.Background()))
	suite.containerdWg.Add(1)
	go func() {
		defer suite.containerdWg.Done()
		defer func() { suite.Require().NoError(suite.containerdRunner.Close()) }()
		suite.Require().NoError(suite.containerdRunner.Run(MockEventSink))
	}()

	suite.client, err = containerd.New(constants.ContainerdAddress)
	suite.Require().NoError(err)

	ctx := namespaces.WithNamespace(context.Background(), containerdNamespace)

	suite.image, err = suite.client.Pull(ctx, busyboxImage, containerd.WithPullUnpack)
	suite.Require().NoError(err)
}

func (suite *ContainersSuite) TearDownSuite() {
	suite.Require().NoError(suite.client.Close())

	suite.Require().NoError(suite.containerdRunner.Stop())
	suite.containerdWg.Wait()

	suite.Require().NoError(os.RemoveAll(suite.tmpDir))
}

func (suite *ContainersSuite) SetupTest() {
	suite.containerRunners = nil
}

func (suite *ContainersSuite) run(runners ...runner.Runner) {
	runningCh := make(chan struct{}, len(runners))

	for _, r := range runners {
		suite.containerRunners = append(suite.containerRunners, r)

		suite.Require().NoError(r.Open(context.Background()))

		suite.containersWg.Add(1)
		go func(r runner.Runner) {
			runningSink := func(state events.ServiceState, message string, args ...interface{}) {
				if state == events.StateRunning {
					runningCh <- struct{}{}
				}
			}

			defer suite.containersWg.Done()
			suite.Assert().NoError(r.Run(runningSink))
		}(r)
	}

	// wait for the containers to be started actually
	for range runners {
		<-runningCh
	}
}

func (suite *ContainersSuite) TearDownTest() {
	for _, r := range suite.containerRunners {
		suite.Assert().NoError(r.Stop())
	}
	suite.containersWg.Wait()
	for _, r := range suite.containerRunners {
		suite.Assert().NoError(r.Close())
	}
}

func (suite *ContainersSuite) runK8sContainers() {
	suite.run(containerdrunner.NewRunner(&userdata.UserData{}, &runner.Args{
		ID:          "test1",
		ProcessArgs: []string{"/bin/sh", "-c", "sleep 3600"},
	},
		runner.WithLogPath(suite.tmpDir),
		runner.WithNamespace(containerdNamespace),
		runner.WithContainerImage(busyboxImage),
		runner.WithContainerOpts(containerd.WithContainerLabels(map[string]string{
			"io.kubernetes.pod.name":      "fun",
			"io.kubernetes.pod.namespace": "ns1",
		})),
		runner.WithOCISpecOpts(WithAnnotations(map[string]string{
			"io.kubernetes.cri.sandbox-log-directory": "sandbox",
			"io.kubernetes.cri.sandbox-id":            "c888d69b73b5b444c2b0bd70da28c3da102b0aeb327f3a297626e2558def327f",
		})),
	), containerdrunner.NewRunner(&userdata.UserData{}, &runner.Args{
		ID:          "test2",
		ProcessArgs: []string{"/bin/sh", "-c", "sleep 3600"},
	},
		runner.WithLogPath(suite.tmpDir),
		runner.WithNamespace(containerdNamespace),
		runner.WithContainerImage(busyboxImage),
		runner.WithContainerOpts(containerd.WithContainerLabels(map[string]string{
			"io.kubernetes.pod.name":       "fun",
			"io.kubernetes.pod.namespace":  "ns1",
			"io.kubernetes.container.name": "run",
		})),
		runner.WithOCISpecOpts(WithAnnotations(map[string]string{
			"io.kubernetes.cri.sandbox-id": "c888d69b73b5b444c2b0bd70da28c3da102b0aeb327f3a297626e2558def327f",
		})),
	))
}

func (suite *ContainersSuite) TestPodsNonK8s() {
	suite.run(containerdrunner.NewRunner(&userdata.UserData{}, &runner.Args{
		ID:          "test",
		ProcessArgs: []string{"/bin/sh", "-c", "sleep 3600"},
	},
		runner.WithLogPath(suite.tmpDir),
		runner.WithNamespace(containerdNamespace),
		runner.WithContainerImage(busyboxImage),
	))

	i, err := containers.NewInspector(context.Background(), containerdNamespace)
	suite.Assert().NoError(err)

	pods, err := i.Pods()
	suite.Require().NoError(err)
	suite.Require().Len(pods, 1)
	suite.Assert().Equal("test", pods[0].Name)
	suite.Assert().Equal("", pods[0].Sandbox)
	suite.Require().Len(pods[0].Containers, 1)
	suite.Assert().Equal("test", pods[0].Containers[0].Display)
	suite.Assert().Equal("test", pods[0].Containers[0].Name)
	suite.Assert().Equal("test", pods[0].Containers[0].ID)
	suite.Assert().Equal(busyboxImage, pods[0].Containers[0].Image)
	suite.Assert().Equal(containerd.Running, pods[0].Containers[0].Status.Status)
	suite.Assert().NotNil(pods[0].Containers[0].Metrics)

	suite.Assert().NoError(i.Close())
}

func (suite *ContainersSuite) TestPodsK8s() {
	suite.runK8sContainers()

	i, err := containers.NewInspector(context.Background(), containerdNamespace)
	suite.Assert().NoError(err)

	pods, err := i.Pods()
	suite.Require().NoError(err)
	suite.Require().Len(pods, 1)
	suite.Assert().Equal("ns1/fun", pods[0].Name)
	suite.Assert().Equal("sandbox", pods[0].Sandbox)
	suite.Require().Len(pods[0].Containers, 2)

	suite.Assert().Equal("ns1/fun", pods[0].Containers[0].Display)
	suite.Assert().Equal("test1", pods[0].Containers[0].Name)
	suite.Assert().Equal("test1", pods[0].Containers[0].ID)
	suite.Assert().Equal("sandbox", pods[0].Containers[0].Sandbox)
	suite.Assert().Equal("0", pods[0].Containers[0].RestartCount)
	suite.Assert().Equal("", pods[0].Containers[0].GetLogFile())
	suite.Assert().Equal(busyboxImage, pods[0].Containers[0].Image)
	suite.Assert().Equal(containerd.Running, pods[0].Containers[0].Status.Status)
	suite.Assert().NotNil(pods[0].Containers[0].Metrics)

	suite.Assert().Equal("ns1/fun:run", pods[0].Containers[1].Display)
	suite.Assert().Equal("run", pods[0].Containers[1].Name)
	suite.Assert().Equal("test2", pods[0].Containers[1].ID)
	suite.Assert().Equal("sandbox", pods[0].Containers[1].Sandbox)
	suite.Assert().Equal("sandbox/run/0.log", pods[0].Containers[1].GetLogFile())
	suite.Assert().Equal(busyboxImage, pods[0].Containers[1].Image)
	suite.Assert().Equal(containerd.Running, pods[0].Containers[1].Status.Status)
	suite.Assert().NotNil(pods[0].Containers[1].Metrics)

	suite.Assert().NoError(i.Close())
}

func (suite *ContainersSuite) TestContainerNonK8s() {
	suite.run(containerdrunner.NewRunner(&userdata.UserData{}, &runner.Args{
		ID:          "shelltest",
		ProcessArgs: []string{"/bin/sh", "-c", "sleep 3600"},
	},
		runner.WithLogPath(suite.tmpDir),
		runner.WithNamespace(containerdNamespace),
		runner.WithContainerImage(busyboxImage),
	))

	i, err := containers.NewInspector(context.Background(), containerdNamespace)
	suite.Assert().NoError(err)

	cntr, err := i.Container("shelltest")
	suite.Require().NoError(err)
	suite.Require().NotNil(cntr)
	suite.Assert().Equal("shelltest", cntr.Name)
	suite.Assert().Equal("shelltest", cntr.Display)
	suite.Assert().Equal("shelltest", cntr.ID)
	suite.Assert().Equal("sha256:c888d69b73b5b444c2b0bd70da28c3da102b0aeb327f3a297626e2558def327f", cntr.Image) // image is not resolved
	suite.Assert().Equal(containerd.Running, cntr.Status.Status)

	cntr, err = i.Container("nosuchcontainer")
	suite.Require().NoError(err)
	suite.Require().Nil(cntr)

	suite.Assert().NoError(i.Close())
}

func (suite *ContainersSuite) TestContainerK8s() {
	suite.runK8sContainers()

	i, err := containers.NewInspector(context.Background(), containerdNamespace)
	suite.Assert().NoError(err)

	cntr, err := i.Container("ns1/fun")
	suite.Require().NoError(err)
	suite.Require().NotNil(cntr)
	suite.Assert().Equal("test1", cntr.Name)
	suite.Assert().Equal("ns1/fun", cntr.Display)
	suite.Assert().Equal("test1", cntr.ID)
	suite.Assert().Equal("sandbox", cntr.Sandbox)
	suite.Assert().Equal("", cntr.GetLogFile())
	suite.Assert().Equal("sha256:c888d69b73b5b444c2b0bd70da28c3da102b0aeb327f3a297626e2558def327f", cntr.Image) // image is not resolved
	suite.Assert().Equal(containerd.Running, cntr.Status.Status)

	cntr, err = i.Container("ns1/fun:run")
	suite.Require().NoError(err)
	suite.Require().NotNil(cntr)
	suite.Assert().Equal("run", cntr.Name)
	suite.Assert().Equal("ns1/fun:run", cntr.Display)
	suite.Assert().Equal("test2", cntr.ID)
	suite.Assert().Equal("sandbox", cntr.Sandbox)
	suite.Assert().Equal("sandbox/run/0.log", cntr.GetLogFile())
	suite.Assert().Equal("sha256:c888d69b73b5b444c2b0bd70da28c3da102b0aeb327f3a297626e2558def327f", cntr.Image) // image is not resolved
	suite.Assert().Equal(containerd.Running, cntr.Status.Status)

	cntr, err = i.Container("ns2/fun:run")
	suite.Require().NoError(err)
	suite.Require().Nil(cntr)

	cntr, err = i.Container("ns1/run:run")
	suite.Require().NoError(err)
	suite.Require().Nil(cntr)

	cntr, err = i.Container("ns1/fun:go")
	suite.Require().NoError(err)
	suite.Require().Nil(cntr)

	suite.Assert().NoError(i.Close())
}

func TestContainersSuite(t *testing.T) {
	if os.Getuid() != 0 {
		t.Skip("can't run the test as non-root")
	}
	_, err := os.Stat("/rootfs/bin/containerd")
	if err != nil {
		t.Skip("containerd binary is not available, skipping the test")
	}

	suite.Run(t, new(ContainersSuite))
}