package supervisor_test

import (
	"git.cin.su/m/Assistant/runtime/supervisor"
	"path/filepath"
	"testing"
)

func TestLifeboatRollsBackCandidateToLastKnownGood(t *testing.T) {
	l := supervisor.Lifeboat{}
	l.Start(supervisor.Release{ID: "good", Version: "1"})
	l.Record(supervisor.Observation{ReleaseID: "good", Ready: true})
	l.Start(supervisor.Release{ID: "candidate", Version: "2"})
	l.Record(supervisor.Observation{ReleaseID: "candidate", Ready: false})
	if !l.Rollback() || l.Current.ID != "good" || l.Current.State != supervisor.ReleaseRolledBack {
		t.Fatalf("%#v", l)
	}
}

func TestPersistentLifeboatRestoresLastKnownGoodAfterCandidateFailure(t *testing.T) {
	store, err := supervisor.NewDurableStore(filepath.Join(t.TempDir(), "control", "lifeboat.json"))
	if err != nil {
		t.Fatal(err)
	}
	first := supervisor.NewPersistentLifeboat(store)
	if ok, err := first.Start(supervisor.Release{ID: "good", Version: "1"}); err != nil || !ok {
		t.Fatalf("start good: %v %v", ok, err)
	}
	if ok, err := first.Record(supervisor.Observation{ReleaseID: "good", Ready: true, ObservedAtUnixMillis: 1}); err != nil || !ok {
		t.Fatalf("record good: %v %v", ok, err)
	}
	if ok, err := first.Start(supervisor.Release{ID: "candidate", Version: "2"}); err != nil || !ok {
		t.Fatalf("start candidate: %v %v", ok, err)
	}
	if ok, err := first.Record(supervisor.Observation{ReleaseID: "candidate", Ready: false, Reason: "candidate crashed", ObservedAtUnixMillis: 2}); err != nil || !ok {
		t.Fatalf("record failure: %v %v", ok, err)
	}
	second := supervisor.NewPersistentLifeboat(store)
	if ok, err := second.Rollback(); err != nil || !ok {
		t.Fatalf("rollback after restart: %v %v", ok, err)
	}
	state, err := second.Snapshot()
	if err != nil {
		t.Fatal(err)
	}
	if state.Current.ID != "good" || state.Current.State != supervisor.ReleaseRolledBack || state.LastKnownGood.ID != "good" {
		t.Fatalf("state=%+v", state)
	}
}

func TestPersistentSupervisorRestoresTrafficModeAfterRestart(t *testing.T) {
	store, err := supervisor.NewDurableStore(filepath.Join(t.TempDir(), "control", "release-state.json"))
	if err != nil {
		t.Fatal(err)
	}
	first := supervisor.NewPersistentLifeboat(store)
	if ok, err := first.Start(supervisor.Release{ID: "candidate", Version: "2"}); err != nil || !ok {
		t.Fatalf("start: %v %v", ok, err)
	}
	if ok, err := first.Record(supervisor.Observation{ReleaseID: "candidate", Ready: true, ObservedAtUnixMillis: 1}); err != nil || !ok {
		t.Fatalf("observe: %v %v", ok, err)
	}
	if ok, err := first.SetTraffic(supervisor.TrafficCanaryHeader, "http://stable:8084", "http://candidate:8084"); err != nil || !ok {
		t.Fatalf("canary: %v %v", ok, err)
	}
	second := supervisor.NewPersistentLifeboat(store)
	state, err := second.Snapshot()
	if err != nil {
		t.Fatal(err)
	}
	if state.TrafficMode != supervisor.TrafficCanaryHeader || state.CandidateUpstream != "http://candidate:8084" || state.LastKnownGood.ID != "candidate" {
		t.Fatalf("state=%+v", state)
	}
	if ok, err := second.RollbackTraffic(); err != nil || !ok {
		t.Fatalf("rollback traffic: %v %v", ok, err)
	}
	state, err = second.Snapshot()
	if err != nil {
		t.Fatal(err)
	}
	if state.TrafficMode != supervisor.TrafficStable || state.Current.State != supervisor.ReleaseRolledBack {
		t.Fatalf("rollback state=%+v", state)
	}
}

func testReleaseSet(id, revision, web, core, cognitive, research string) supervisor.ReleaseSet {
	return supervisor.ReleaseSet{ID: id, Revision: revision, Components: map[string]supervisor.ComponentRelease{
		supervisor.ComponentCreatorWeb:      {Name: supervisor.ComponentCreatorWeb, Revision: revision, Upstream: web, HealthURL: web + "/healthz", ReadyURL: web + "/readyz"},
		supervisor.ComponentCoreRuntime:     {Name: supervisor.ComponentCoreRuntime, Revision: revision, Upstream: core, HealthURL: core + "/healthz", ReadyURL: core + "/readyz"},
		supervisor.ComponentCognitiveWorker: {Name: supervisor.ComponentCognitiveWorker, Revision: revision, Upstream: cognitive, HealthURL: cognitive + "/healthz", ReadyURL: cognitive + "/readyz"},
		supervisor.ComponentResearchRuntime: {Name: supervisor.ComponentResearchRuntime, Revision: revision, Upstream: research, HealthURL: research + "/healthz", ReadyURL: research + "/readyz"},
	}}
}

func TestReleaseSetRejectsMixedRevisionsAndMissingComponents(t *testing.T) {
	set := testReleaseSet("candidate", "r1", "http://web", "http://core", "http://cognitive", "http://research")
	if !set.Validate() {
		t.Fatal("complete release set should validate")
	}
	set.Components[supervisor.ComponentCoreRuntime] = supervisor.ComponentRelease{Name: supervisor.ComponentCoreRuntime, Revision: "r0", Upstream: "http://core", HealthURL: "http://core/healthz", ReadyURL: "http://core/readyz"}
	if set.Validate() {
		t.Fatal("mixed revision release set must be rejected")
	}
}

func TestPersistentReleaseSetRollbackPreservesWholeLastKnownGood(t *testing.T) {
	store, err := supervisor.NewDurableStore(filepath.Join(t.TempDir(), "control", "release-set.json"))
	if err != nil {
		t.Fatal(err)
	}
	first := supervisor.NewPersistentLifeboat(store)
	good := testReleaseSet("stable-r1", "r1", "http://web", "http://core", "http://cognitive", "http://research")
	candidate := testReleaseSet("candidate-r2", "r2", "http://web-c", "http://core-c", "http://cognitive-c", "http://research-c")
	if ok, err := first.EstablishBaselineSet(good); err != nil || !ok {
		t.Fatalf("baseline: %v %v", ok, err)
	}
	if ok, err := first.StartSet(candidate); err != nil || !ok {
		t.Fatalf("start: %v %v", ok, err)
	}
	observation := supervisor.ReleaseSetObservation{ReleaseID: candidate.ID, Ready: true, Components: map[string]supervisor.ComponentObservation{}}
	for name := range candidate.Components {
		observation.Components[name] = supervisor.ComponentObservation{Revision: candidate.Revision, Health: true, Ready: true}
	}
	if ok, err := first.RecordSet(observation); err != nil || !ok {
		t.Fatalf("observe: %v %v", ok, err)
	}
	if ok, err := first.SetTrafficForSets(supervisor.TrafficCanaryHeader); err != nil || !ok {
		t.Fatalf("canary: %v %v", ok, err)
	}
	second := supervisor.NewPersistentLifeboat(store)
	if ok, err := second.RollbackSet(); err != nil || !ok {
		t.Fatalf("rollback: %v %v", ok, err)
	}
	state, err := second.Snapshot()
	if err != nil {
		t.Fatal(err)
	}
	if state.StableReleaseSet == nil || state.StableReleaseSet.Revision != "r1" || state.TrafficMode != supervisor.TrafficStable {
		t.Fatalf("state=%+v", state)
	}
}

func TestReleaseSetPromotionRetainsPreviousStableForRollback(t *testing.T) {
	good := testReleaseSet("stable-r1", "r1", "http://web", "http://core", "http://cognitive", "http://research")
	candidate := testReleaseSet("candidate-r2", "r2", "http://web-c", "http://core-c", "http://cognitive-c", "http://research-c")
	l := supervisor.Lifeboat{}
	if !l.EstablishBaselineSet(good) || !l.StartSet(candidate) {
		t.Fatal("release set setup failed")
	}
	observation := supervisor.ReleaseSetObservation{ReleaseID: candidate.ID, Ready: true, Components: map[string]supervisor.ComponentObservation{}}
	for name := range candidate.Components {
		observation.Components[name] = supervisor.ComponentObservation{Revision: candidate.Revision, Health: true, Ready: true}
	}
	if !l.RecordSet(observation) || !l.PromoteSet() {
		t.Fatal("promotion failed")
	}
	if l.RollbackTargetReleaseSet == nil || l.RollbackTargetReleaseSet.Revision != "r1" {
		t.Fatalf("rollback target=%+v", l.RollbackTargetReleaseSet)
	}
}