// -*- Mode: Go; indent-tabs-mode: t -*-

/*
 * Copyright (C) 2024 Canonical Ltd
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License version 3 as
 * published by the Free Software Foundation.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package ctlcmd_test

import (
	"context"
	"time"

	. "gopkg.in/check.v1"

	"github.com/snapcore/snapd/asserts"
	"github.com/snapcore/snapd/asserts/assertstest"
	"github.com/snapcore/snapd/asserts/snapasserts"
	"github.com/snapcore/snapd/dirs"
	"github.com/snapcore/snapd/interfaces"
	"github.com/snapcore/snapd/overlord/hookstate"
	"github.com/snapcore/snapd/overlord/hookstate/ctlcmd"
	"github.com/snapcore/snapd/overlord/hookstate/hooktest"
	"github.com/snapcore/snapd/overlord/ifacestate/ifacerepo"
	"github.com/snapcore/snapd/overlord/snapstate"
	"github.com/snapcore/snapd/overlord/snapstate/sequence"
	"github.com/snapcore/snapd/overlord/snapstate/snapstatetest"
	"github.com/snapcore/snapd/overlord/state"
	"github.com/snapcore/snapd/snap"
	"github.com/snapcore/snapd/snap/naming"
	"github.com/snapcore/snapd/snap/snaptest"
	"github.com/snapcore/snapd/testutil"
)

type installSuite struct {
	testutil.BaseTest
	st          *state.State
	chg         *state.Change
	mockContext *hookstate.Context
	mockHandler *hooktest.MockHandler
}

var _ = Suite(&installSuite{})

const snapWithCompsYaml = `name: test-snap
version: 1.0
summary: test-snap
components:
  comp1:
    type: standard
  comp2:
    type: standard
`

func (s *installSuite) SetUpTest(c *C) {
	s.BaseTest.SetUpTest(c)
	oldRoot := dirs.GlobalRootDir
	dirs.SetRootDir(c.MkDir())

	s.BaseTest.AddCleanup(func() {
		dirs.SetRootDir(oldRoot)
	})
	s.BaseTest.AddCleanup(snap.MockSanitizePlugsSlots(func(snapInfo *snap.Info) {}))

	s.mockHandler = hooktest.NewMockHandler()

	s.st = state.New(nil)
	s.st.Lock()
	defer s.st.Unlock()

	repo := interfaces.NewRepository()
	ifacerepo.Replace(s.st, repo)

	// mock installed snaps
	info1 := snaptest.MockSnapCurrent(c, string(snapWithCompsYaml), &snap.SideInfo{
		Revision: snap.R(1),
	})
	snapstate.Set(s.st, info1.InstanceName(), &snapstate.SnapState{
		Active: true,
		Sequence: snapstatetest.NewSequenceFromSnapSideInfos([]*snap.SideInfo{
			{
				RealName: info1.SnapName(),
				Revision: info1.Revision,
				SnapID:   "test-snap-id",
			},
		}),
		Current: info1.Revision,
	})

	s.chg = s.st.NewChange("install change", "install change")
	task := s.st.NewTask("test-task", "my test task")
	s.chg.AddTask(task)
	setup := &hookstate.HookSetup{Snap: "test-snap", Revision: snap.R(1), Hook: "test-hook"}

	var err error
	s.mockContext, err = hookstate.NewContext(task, task.State(), setup, s.mockHandler, "")
	c.Assert(err, IsNil)
}

func (s *installSuite) testMngmtCommand(c *C, cmd string) {
	s.st.Lock()
	task := s.st.NewTask("queued", "queued task")
	s.st.Unlock()

	switch cmd {
	case "install":
		ctlcmd.MockSnapstateInstallComponentsFunc(func(ctx context.Context, st *state.State, names []string, info *snap.Info, vsets *snapasserts.ValidationSets, opts snapstate.Options) ([]*state.TaskSet, error) {
			c.Check(names, DeepEquals, []string{"comp1", "comp2"})
			c.Check(opts, DeepEquals, snapstate.Options{ExpectOneSnap: true,
				ConflictOptions: snapstate.ConflictOptions{FromChange: s.mockContext.ChangeID()}})
			var ts state.TaskSet
			ts.AddTask(task)
			return []*state.TaskSet{&ts}, nil
		})
	case "remove":
		ctlcmd.MockSnapstateRemoveComponentsFunc(func(st *state.State, snapName string, compNames []string, opts snapstate.RemoveComponentsOpts) ([]*state.TaskSet, error) {
			c.Check(compNames, DeepEquals, []string{"comp1", "comp2"})
			c.Check(opts, DeepEquals,
				snapstate.RemoveComponentsOpts{RefreshProfile: true,
					ConflictOptions: snapstate.ConflictOptions{FromChange: s.mockContext.ChangeID()}})
			var ts state.TaskSet
			ts.AddTask(task)
			return []*state.TaskSet{&ts}, nil
		})
	}

	_, _, err := ctlcmd.Run(s.mockContext, []string{cmd, "test-snap+comp1", "+comp2"}, 0, nil)
	c.Assert(err, IsNil)
	s.st.Lock()
	// one is the task added in SetUpTest, the other is the queued task
	c.Assert(len(s.chg.Tasks()), Equals, 2)
	s.st.Unlock()
}

func (s *installSuite) TestInstallCommand(c *C) {
	s.testMngmtCommand(c, "install")
}

func (s *installSuite) TestRemoveCommand(c *C) {
	s.testMngmtCommand(c, "remove")
}

func (s *installSuite) TestInstallCommandUsesPendingValidationSets(c *C) {
	signer := assertstest.NewStoreStack("can0nical", nil)
	headers := map[string]any{
		"authority-id": "can0nical",
		"account-id":   "can0nical",
		"name":         "set-one",
		"series":       "16",
		"sequence":     "1",
		"timestamp":    time.Now().UTC().Format(time.RFC3339),
		"snaps": []any{
			map[string]any{
				"id":       snaptest.AssertedSnapID("test-snap"),
				"name":     "test-snap",
				"presence": "required",
			},
		},
	}

	assertion, err := signer.Sign(asserts.ValidationSetType, headers, nil, "")
	c.Assert(err, IsNil)
	vs := assertion.(*asserts.ValidationSet)

	encoded := map[string][]byte{
		string(snapasserts.NewValidationSetKey(vs)): asserts.Encode(assertion),
	}

	s.st.Lock()
	enforce := s.st.NewTask("enforce-validation-sets", "enforce validation sets")
	enforce.Set("validation-sets", encoded)
	s.chg.AddTask(enforce)
	task := s.st.NewTask("queued", "queued task")
	s.st.Unlock()

	var called bool
	restore := ctlcmd.MockSnapstateInstallComponentsFunc(func(ctx context.Context, st *state.State, names []string, info *snap.Info, vsets *snapasserts.ValidationSets, opts snapstate.Options) ([]*state.TaskSet, error) {
		called = true
		c.Assert(vsets, NotNil)
		c.Assert(vsets.Keys(), DeepEquals, []snapasserts.ValidationSetKey{snapasserts.NewValidationSetKey(vs)})
		c.Check(names, DeepEquals, []string{"comp1", "comp2"})
		c.Check(opts, DeepEquals, snapstate.Options{
			ExpectOneSnap:   true,
			ConflictOptions: snapstate.ConflictOptions{FromChange: s.mockContext.ChangeID()},
		})

		return []*state.TaskSet{state.NewTaskSet(task)}, nil
	})
	defer restore()

	_, _, err = ctlcmd.Run(s.mockContext, []string{"install", "test-snap+comp1", "+comp2"}, 0, nil)
	c.Assert(err, IsNil)
	c.Check(called, Equals, true)
}

func (s *installSuite) testEphemeralMngmtCommand(c *C, cmd string) {
	s.st.Lock()
	task := s.st.NewTask("test", "test task")
	setup := &hookstate.HookSetup{Snap: "test-snap", Revision: snap.R(1)}
	s.st.Unlock()

	switch cmd {
	case "install":
		ctlcmd.MockSnapstateInstallComponentsFunc(func(ctx context.Context, st *state.State, names []string, info *snap.Info, vsets *snapasserts.ValidationSets, opts snapstate.Options) ([]*state.TaskSet, error) {
			c.Check(names, DeepEquals,
				[]string{"comp1", "comp2", "comp3", "comp4", "comp5", "comp6"})
			c.Check(opts, DeepEquals, snapstate.Options{ExpectOneSnap: true})
			var ts state.TaskSet
			ts.AddTask(task)
			return []*state.TaskSet{&ts}, nil
		})
	case "remove":
		ctlcmd.MockSnapstateRemoveComponentsFunc(func(st *state.State, snapName string, compNames []string, opts snapstate.RemoveComponentsOpts) ([]*state.TaskSet, error) {
			c.Check(compNames, DeepEquals,
				[]string{"comp1", "comp2", "comp3", "comp4", "comp5", "comp6"})
			c.Check(opts, DeepEquals,
				snapstate.RemoveComponentsOpts{RefreshProfile: true})
			var ts state.TaskSet
			ts.AddTask(task)
			return []*state.TaskSet{&ts}, nil
		})
	}

	var err error
	s.mockContext, err = hookstate.NewContext(nil, s.st, setup, s.mockHandler, "")
	c.Assert(err, IsNil)

	chann := make(chan error)
	go func() {
		_, _, err := ctlcmd.Run(s.mockContext,
			[]string{cmd, "test-snap+comp1", "test-snap+comp2+comp3",
				"+comp4", "+comp5+comp6"}, 0, nil)
		chann <- err
	}()

	// Wait for the change to be created and assigned to task
	var chg *state.Change
	for i := 0; i < 50; i++ {
		s.st.Lock()
		chg = task.Change()
		s.st.Unlock()
		if chg != nil {
			break
		}
		time.Sleep(100 * time.Millisecond)
	}
	if chg == nil {
		c.Error("change not assigned to ephemeral snapctl command")
	}

	s.st.Lock()
	chg.SetStatus(state.DoneStatus)
	s.st.Unlock()

	err = <-chann
	c.Assert(err, IsNil)
}

func (s *installSuite) TestEphemeralInstallCommand(c *C) {
	s.testEphemeralMngmtCommand(c, "install")
}

func (s *installSuite) TestEphemeralRemoveCommand(c *C) {
	s.testEphemeralMngmtCommand(c, "remove")
}

func (s *installSuite) testMgmntCommandOtherSnap(c *C, cmd string) {
	_, _, err := ctlcmd.Run(s.mockContext, []string{cmd, "+comp1", "other-snap+comp2"}, 0, nil)
	c.Assert(err, ErrorMatches, "cannot install snaps using snapctl")
}

func (s *installSuite) TestInstallCommandOtherSnap(c *C) {
	s.testMgmntCommandOtherSnap(c, "install")
}

func (s *installSuite) TestRemoveCommandOtherSnap(c *C) {
	s.testMgmntCommandOtherSnap(c, "remove")
}

func (s *installSuite) testMgmntCommandBadCompName(c *C, cmd string) {
	_, _, err := ctlcmd.Run(s.mockContext, []string{cmd, "+comp_1"}, 0, nil)
	c.Assert(err, ErrorMatches, `invalid snap name: "comp_1"`)
}

func (s *installSuite) TestInstallCommandBadCompName(c *C) {
	s.testMgmntCommandBadCompName(c, "install")
}

func (s *installSuite) TestRemoveCommandBadCompName(c *C) {
	s.testMgmntCommandBadCompName(c, "remove")
}

func (s *installSuite) TestNoWaitNonEphemeralReturnsError(c *C) {
	c.Skip("Content removed for release 2.76")

	for _, cmd := range []string{"install", "remove"} {
		s.st.Lock()
		task := s.st.NewTask("test", "test task")
		s.st.Unlock()

		switch cmd {
		case "install":
			restore := ctlcmd.MockSnapstateInstallComponentsFunc(func(ctx context.Context, st *state.State, names []string, info *snap.Info, vsets *snapasserts.ValidationSets, opts snapstate.Options) ([]*state.TaskSet, error) {
				return []*state.TaskSet{state.NewTaskSet(task)}, nil
			})
			defer restore()
		case "remove":
			restore := ctlcmd.MockSnapstateRemoveComponentsFunc(func(st *state.State, snapName string, compNames []string, opts snapstate.RemoveComponentsOpts) ([]*state.TaskSet, error) {
				return []*state.TaskSet{state.NewTaskSet(task)}, nil
			})
			defer restore()
		}

		_, _, err := ctlcmd.Run(s.mockContext, []string{cmd, "+comp1", "--no-wait"}, 0, nil)
		c.Assert(err, ErrorMatches, "internal error: cannot run snap management command asynchronously from a non-ephemeral context", Commentf("cmd: %s", cmd))
	}
}

func (s *installSuite) TestNoWaitInstallAndRemoveCommands(c *C) {
	c.Skip("Content removed for release 2.76")

	for _, cmd := range []string{"install", "remove"} {
		s.st.Lock()
		task := s.st.NewTask("test", "test task")
		setup := &hookstate.HookSetup{Snap: "test-snap", Revision: snap.R(1)}
		s.st.Unlock()

		switch cmd {
		case "install":
			restore := ctlcmd.MockSnapstateInstallComponentsFunc(func(ctx context.Context, st *state.State, names []string, info *snap.Info, vsets *snapasserts.ValidationSets, opts snapstate.Options) ([]*state.TaskSet, error) {
				c.Check(names, DeepEquals, []string{"comp1"})
				c.Check(opts, DeepEquals, snapstate.Options{ExpectOneSnap: true})
				return []*state.TaskSet{state.NewTaskSet(task)}, nil
			})
			defer restore()
		case "remove":
			restore := ctlcmd.MockSnapstateRemoveComponentsFunc(func(st *state.State, snapName string, compNames []string, opts snapstate.RemoveComponentsOpts) ([]*state.TaskSet, error) {
				c.Check(compNames, DeepEquals, []string{"comp1"})
				c.Check(opts, DeepEquals, snapstate.RemoveComponentsOpts{RefreshProfile: true})
				return []*state.TaskSet{state.NewTaskSet(task)}, nil
			})
			defer restore()
		}

		var err error
		s.mockContext, err = hookstate.NewContext(nil, s.st, setup, s.mockHandler, "")
		c.Assert(err, IsNil)

		stdout, _, err := ctlcmd.Run(s.mockContext, []string{cmd, "+comp1", "--no-wait"}, 0, nil)
		c.Assert(err, IsNil)
		changeID := string(stdout)
		c.Assert(changeID, Not(Equals), "")

		s.st.Lock()
		c.Check(task.Change().ID(), Equals, changeID)
		s.st.Unlock()
	}
}

func (s *installSuite) TestInstallWithParallelInstalledSnap(c *C) {
	s.st.Lock()
	s.chg = s.st.NewChange("install change", "install change")
	task := s.st.NewTask("test-task", "my test task")
	s.chg.AddTask(task)
	setup := &hookstate.HookSetup{Snap: "test-snap_foo", Revision: snap.R(1), Hook: "test-hook"}

	// create a context for the parallel installed snap
	var err error
	s.mockContext, err = hookstate.NewContext(task, task.State(), setup, s.mockHandler, "")
	c.Assert(err, IsNil)

	installTask := s.st.NewTask("queued", "queued task")
	s.st.Unlock()

	restore := ctlcmd.MockSnapstateInstallComponentsFunc(func(ctx context.Context, st *state.State, names []string, info *snap.Info, vsets *snapasserts.ValidationSets, opts snapstate.Options) ([]*state.TaskSet, error) {
		c.Check(names, DeepEquals, []string{"two"})
		c.Check(opts, DeepEquals, snapstate.Options{ExpectOneSnap: true,
			ConflictOptions: snapstate.ConflictOptions{FromChange: s.mockContext.ChangeID()}})
		var ts state.TaskSet
		ts.AddTask(installTask)
		return []*state.TaskSet{&ts}, nil
	})
	defer restore()

	rev := snap.R(1)
	si := &snap.SideInfo{
		RealName: "test-snap",
		Revision: rev,
		SnapID:   "test-snap-id",
	}

	seq := snapstatetest.NewSequenceFromRevisionSideInfos([]*sequence.RevisionSideState{
		sequence.NewRevisionSideState(si, nil),
	})

	// component +one is already installed
	seq.AddComponentForRevision(snap.R(1), sequence.NewComponentState(&snap.ComponentSideInfo{
		Component: naming.NewComponentRef("test-snap", "one"),
		Revision:  snap.R(1),
	}, snap.StandardComponent))

	s.st.Lock()
	snapstate.Set(s.st, "test-snap_foo", &snapstate.SnapState{
		Active:   true,
		Sequence: seq,
		Current:  rev,
	})
	s.st.Unlock()

	stdout, stderr, err := ctlcmd.Run(s.mockContext, []string{"install", "+one", "+two"}, 0, nil)
	c.Check(err, IsNil)
	c.Check(stdout, HasLen, 0)
	c.Check(string(stderr), Matches, `(?sm).*snapctl: component "one" is already installed`)
}

func (s *installSuite) TestInstallAllAlreadyInstalled(c *C) {
	rev := snap.R(1)
	si := &snap.SideInfo{
		RealName: "test-snap",
		Revision: rev,
		SnapID:   "test-snap-id",
	}

	seq := snapstatetest.NewSequenceFromRevisionSideInfos([]*sequence.RevisionSideState{
		sequence.NewRevisionSideState(si, nil),
	})

	seq.AddComponentForRevision(snap.R(1), sequence.NewComponentState(&snap.ComponentSideInfo{
		Component: naming.NewComponentRef("test-snap", "one"),
		Revision:  snap.R(1),
	}, snap.StandardComponent))

	seq.AddComponentForRevision(snap.R(1), sequence.NewComponentState(&snap.ComponentSideInfo{
		Component: naming.NewComponentRef("test-snap", "two"),
		Revision:  snap.R(1),
	}, snap.StandardComponent))

	s.st.Lock()
	snapstate.Set(s.st, "test-snap", &snapstate.SnapState{
		Active:   true,
		Sequence: seq,
		Current:  rev,
	})
	s.st.Unlock()

	stdout, stderr, err := ctlcmd.Run(s.mockContext, []string{"install", "+one", "+two"}, 0, nil)
	c.Check(err, IsNil)
	c.Check(stdout, HasLen, 0)
	c.Check(string(stderr), Matches, `(?sm).*snapctl: component "one" is already installed`)
	c.Check(string(stderr), Matches, `(?sm).*snapctl: component "two" is already installed`)
}

func (s *installSuite) TestInstallSomeAlreadyInstalled(c *C) {
	s.st.Lock()
	task := s.st.NewTask("queued", "queued task")
	s.st.Unlock()

	restore := ctlcmd.MockSnapstateInstallComponentsFunc(func(ctx context.Context, st *state.State, names []string, info *snap.Info, vsets *snapasserts.ValidationSets, opts snapstate.Options) ([]*state.TaskSet, error) {
		c.Check(names, DeepEquals, []string{"two"})
		c.Check(opts, DeepEquals, snapstate.Options{ExpectOneSnap: true,
			ConflictOptions: snapstate.ConflictOptions{FromChange: s.mockContext.ChangeID()}})
		var ts state.TaskSet
		ts.AddTask(task)
		return []*state.TaskSet{&ts}, nil
	})
	defer restore()

	rev := snap.R(1)
	si := &snap.SideInfo{
		RealName: "test-snap",
		Revision: rev,
		SnapID:   "test-snap-id",
	}

	seq := snapstatetest.NewSequenceFromRevisionSideInfos([]*sequence.RevisionSideState{
		sequence.NewRevisionSideState(si, nil),
	})

	seq.AddComponentForRevision(snap.R(1), sequence.NewComponentState(&snap.ComponentSideInfo{
		Component: naming.NewComponentRef("test-snap", "one"),
		Revision:  snap.R(1),
	}, snap.StandardComponent))

	s.st.Lock()
	snapstate.Set(s.st, "test-snap", &snapstate.SnapState{
		Active:   true,
		Sequence: seq,
		Current:  rev,
	})

	s.st.Unlock()
	stdout, stderr, err := ctlcmd.Run(s.mockContext, []string{"install", "+one", "+two"}, 0, nil)
	c.Check(err, IsNil)
	c.Check(stdout, HasLen, 0)
	c.Check(string(stderr), Matches, `snapctl: component "one" is already installed\n`)
}
