package provision

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/carlosframework/platform/internal/release"
	"github.com/carlosframework/platform/internal/runs"
	"github.com/carlosframework/platform/internal/seeds"
)

// keyed is caller A's admission that carries key k and digest d.
func keyed(id, k, d string) RunAdmission { return keyedBy("A", id, k, d) }

// keyedBy is keyed for caller.
func keyedBy(caller, id, k, d string) RunAdmission {
	return RunAdmission{ID: id, By: "svc", Deadline: time.Now().Add(time.Minute), Key: k, Digest: d, Caller: caller}
}

func wantReplay(t *testing.T, what string, err error, id string) {
	t.Helper()
	var replay *RunKeyReplay
	if !errors.As(err, &replay) || replay.RunID != id {
		t.Fatalf("%s: %v, want a RunKeyReplay of %s", what, err, id)
	}
}

func TestRunAdmitRecordsAndReplaysAKey(t *testing.T) {
	s := release.NewDirStore(t.TempDir())
	rec := runTestRecord(t, s)
	got, err := setRun(t, s, rec, RunAdmit(keyed("run_a", "K", "D")))
	if err != nil || got.ActiveRun.ID != "run_a" {
		t.Fatalf("keyed admit: %+v, %v", got, err)
	}
	if len(got.RunKeys) != 1 || got.RunKeys[0].Key != "K" || got.RunKeys[0].RunID != "run_a" || got.RunKeys[0].Digest != "D" || got.RunKeys[0].At.IsZero() {
		t.Fatalf("the admission must record its key: %+v", got.RunKeys)
	}

	// While run_a is still open, the replay is answered before the open-run
	// refusal.
	_, err = setRun(t, s, rec, RunAdmit(keyed("run_b", "K", "D")))
	wantReplay(t, "the same key and digest while open", err, "run_a")
	if _, err := setRun(t, s, rec, RunAdmit(keyed("run_b", "K", "D2"))); !errors.Is(err, ErrRunKeyReused) {
		t.Fatalf("the same key with another digest: %v, want ErrRunKeyReused", err)
	}
	// Another key still meets the open run.
	if _, err := setRun(t, s, rec, RunAdmit(keyed("run_b", "K2", "D"))); !errors.Is(err, ErrRunOpen) {
		t.Fatalf("another key while open: %v, want ErrRunOpen", err)
	}

	// Picked, finished and cleared (the ack): the key still replays.
	for _, step := range []func(*Record) error{
		RunPick("run_a", "box-1", "boot-1", time.Now().Add(time.Hour)),
		RunFinish("run_a"),
		RunClear("run_a"),
	} {
		if _, err := setRun(t, s, rec, step); err != nil {
			t.Fatal(err)
		}
		_, err = setRun(t, s, rec, RunAdmit(keyed("run_b", "K", "D")))
		wantReplay(t, "after a run transition", err, "run_a")
	}

	// A later, unkeyed run replaces it in the record; the key still replays.
	if _, err := setRun(t, s, rec, RunAdmit(RunAdmission{ID: "run_c", Deadline: time.Now().Add(time.Minute)})); err != nil {
		t.Fatal(err)
	}
	if _, err := setRun(t, s, rec, RunExpire("run_c")); err != nil {
		t.Fatal(err)
	}
	_, err = setRun(t, s, rec, RunAdmit(keyed("run_d", "K", "D")))
	wantReplay(t, "after a later run replaced it", err, "run_a")

	// A refused replay wrote nothing.
	cur, err := GetRecord(context.Background(), s, rec.Account, rec.App, rec.Host)
	if err != nil || cur.ActiveRun.ID != "run_c" || len(cur.RunKeys) != 1 {
		t.Fatalf("record after the replays: run %+v keys %+v, %v", cur.ActiveRun, cur.RunKeys, err)
	}
}

// putRunKeys installs keys on the record directly, as an earlier console's
// admissions would have left them.
func putRunKeys(t *testing.T, s release.Store, rec Record, keys []RunKey) {
	t.Helper()
	if _, err := updateRecord(context.Background(), s, rec.Account, rec.App, rec.Host, func(r *Record) error {
		r.RunKeys = keys
		return nil
	}); err != nil {
		t.Fatal(err)
	}
}

func TestRunKeysPrunedByAgeAndCount(t *testing.T) {
	s := release.NewDirStore(t.TempDir())
	rec := runTestRecord(t, s)
	now := time.Now().UTC()
	var keys []RunKey
	keys = append(keys, RunKey{Key: "old", RunID: "run_old", Digest: "D", Caller: "A", At: now.Add(-runs.KeyTTL - time.Minute)})
	for i := 0; i < runs.MaxRunKeysPerCaller+2; i++ {
		keys = append(keys, RunKey{Key: fmt.Sprintf("k%03d", i), RunID: fmt.Sprintf("run_%03d", i), Digest: "D", Caller: "A", At: now.Add(-time.Hour + time.Duration(i)*time.Second)})
	}
	putRunKeys(t, s, rec, keys)

	// An expired key is not a replay: it admits a fresh run.
	got, err := setRun(t, s, rec, RunAdmit(keyed("run_new", "old", "D")))
	if err != nil || got.ActiveRun.ID != "run_new" {
		t.Fatalf("an expired key must admit afresh: %+v, %v", got, err)
	}
	if len(got.RunKeys) != runs.MaxRunKeysPerCaller {
		t.Fatalf("%d keys kept, want %d", len(got.RunKeys), runs.MaxRunKeysPerCaller)
	}
	kept := map[string]bool{}
	for _, k := range got.RunKeys {
		kept[k.Key+"/"+k.RunID] = true
	}
	if kept["old/run_old"] {
		t.Error("the entry older than KeyTTL survived")
	}
	if !kept["old/run_new"] {
		t.Error("the newest entry was not kept")
	}
	// 34 fresh + 1 new, the caller's newest 32 kept: k000, k001 and k002 go.
	for i := 0; i < runs.MaxRunKeysPerCaller+2; i++ {
		name := fmt.Sprintf("k%03d/run_%03d", i, i)
		if want := i >= 3; kept[name] != want {
			t.Errorf("%s kept %v, want %v", name, kept[name], want)
		}
	}

	// Under the count, age alone prunes.
	s2 := release.NewDirStore(t.TempDir())
	rec2 := runTestRecord(t, s2)
	putRunKeys(t, s2, rec2, []RunKey{
		{Key: "old", RunID: "run_old", Digest: "D", Caller: "A", At: now.Add(-runs.KeyTTL - time.Minute)},
		{Key: "a", RunID: "run_a", Digest: "D", Caller: "A", At: now.Add(-time.Hour)},
	})
	got, err = setRun(t, s2, rec2, RunAdmit(keyed("run_b", "b", "D")))
	if err != nil || len(got.RunKeys) != 2 || got.RunKeys[0].Key != "a" || got.RunKeys[1].Key != "b" {
		t.Fatalf("an expired entry under the count: %+v, %v; want a then b", got.RunKeys, err)
	}

	// The record's lookup ignores an expired entry too.
	r := Record{RunKeys: []RunKey{{Key: "x", RunID: "run_x", Digest: "D", At: now.Add(-runs.KeyTTL - time.Second)}}}
	if _, ok := r.RunKeyFor("x", now); ok {
		t.Error("RunKeyFor answered an expired key")
	}
	r.RunKeys[0].At = now.Add(-time.Minute)
	if k, ok := r.RunKeyFor("x", now); !ok || k.RunID != "run_x" {
		t.Errorf("RunKeyFor = %+v, %v", k, ok)
	}
}

// TestRunKeysCappedPerCaller: the count cap is each caller's own, so caller
// B's admissions never evict caller A's key, and A's own 33rd admission
// evicts A's oldest.
func TestRunKeysCappedPerCaller(t *testing.T) {
	s := release.NewDirStore(t.TempDir())
	rec := runTestRecord(t, s)
	admit := func(a RunAdmission) {
		t.Helper()
		if _, err := setRun(t, s, rec, RunAdmit(a)); err != nil {
			t.Fatalf("admit %s: %v", a.ID, err)
		}
		if _, err := setRun(t, s, rec, RunExpire(a.ID)); err != nil {
			t.Fatal(err)
		}
	}
	admit(keyedBy("A", "run_a0", "KA", "D"))
	for i := range 40 {
		admit(keyedBy("B", fmt.Sprintf("run_b%02d", i), fmt.Sprintf("KB%02d", i), "D"))
	}
	_, err := setRun(t, s, rec, RunAdmit(keyedBy("A", "run_x", "KA", "D")))
	wantReplay(t, "A's key after B's 40 admissions", err, "run_a0")
	cur, err := GetRecord(context.Background(), s, rec.Account, rec.App, rec.Host)
	if err != nil {
		t.Fatal(err)
	}
	per := map[string]int{}
	for _, k := range cur.RunKeys {
		per[k.Caller]++
	}
	if per["A"] != 1 || per["B"] != runs.MaxRunKeysPerCaller {
		t.Fatalf("entries per caller %v, want A 1 and B %d", per, runs.MaxRunKeysPerCaller)
	}
	for i := 1; i < runs.MaxRunKeysPerCaller; i++ {
		admit(keyedBy("A", fmt.Sprintf("run_a%02d", i), fmt.Sprintf("KA%02d", i), "D"))
	}
	_, err = setRun(t, s, rec, RunAdmit(keyedBy("A", "run_y", "KA", "D")))
	wantReplay(t, "A's key with 32 of A's own", err, "run_a0")
	admit(keyedBy("A", "run_a32", "KA32", "D")) // A's 33rd
	got, err := setRun(t, s, rec, RunAdmit(keyedBy("A", "run_z", "KA", "D")))
	if err != nil || got.ActiveRun.ID != "run_z" {
		t.Fatalf("A's oldest key after A's 33rd admission: %+v, %v; want it evicted and a fresh run", got, err)
	}
}

func TestRunAdmitRefusesAMalformedKeyedAdmission(t *testing.T) {
	s := release.NewDirStore(t.TempDir())
	rec := runTestRecord(t, s)
	if _, err := setRun(t, s, rec, RunAdmit(keyed("run_a", "K", ""))); !errors.Is(err, ErrRunConflict) {
		t.Fatalf("a key without a digest: %v, want ErrRunConflict", err)
	}
	if _, err := setRun(t, s, rec, RunAdmit(keyedBy("", "run_a", "K", "D"))); !errors.Is(err, ErrRunConflict) {
		t.Fatalf("a key without a caller: %v, want ErrRunConflict", err)
	}
}

// keysJSON is the record's run_keys as stored, byte for byte.
func keysJSON(t *testing.T, s release.Store, rec Record) string {
	t.Helper()
	raw, _, err := s.Get(context.Background(), RecordKey(rec.Account, rec.App, rec.Host))
	if err != nil {
		t.Fatal(err)
	}
	var m map[string]json.RawMessage
	if err := json.Unmarshal(raw, &m); err != nil {
		t.Fatal(err)
	}
	return string(m["run_keys"])
}

// TestRunKeysSurviveEveryRecordMutation: every run transition, the
// ordinary writers, the move transitions and the import transitions carry
// run_keys through unchanged. Only a keyed admission adds to them.
func TestRunKeysSurviveEveryRecordMutation(t *testing.T) {
	ctx := context.Background()
	at := time.Now().UTC().Truncate(time.Second)
	seeded := []RunKey{{Key: strings.Repeat("a", 64), RunID: "run_k", Digest: strings.Repeat("d", 64), At: at}}
	fresh := func(t *testing.T, r Record) (release.Store, Record, string) {
		t.Helper()
		s := release.NewDirStore(t.TempDir())
		if err := CreateRecord(ctx, s, r); err != nil {
			t.Fatal(err)
		}
		putRunKeys(t, s, r, seeded)
		want := keysJSON(t, s, r)
		if !strings.Contains(want, "run_k") {
			t.Fatalf("run_keys = %s", want)
		}
		return s, r, want
	}
	base := func() Record {
		return Record{Schema: 1, Host: "a.go.example.com", Slug: "s", App: "app", Account: "bab", Region: "ie", Channel: "stable", CreatedAt: time.Now().UTC()}
	}
	type step struct {
		name  string
		apply func(release.Store, Record) error
	}
	run := func(mut func(*Record) error) func(release.Store, Record) error {
		return func(s release.Store, r Record) error {
			_, err := SetRecordRun(ctx, s, r.Account, r.App, r.Host, mut)
			return err
		}
	}
	move := func(mut func(*Record) error) func(release.Store, Record) error {
		return func(s release.Store, r Record) error {
			_, err := SetRecordMove(ctx, s, r.Account, r.App, r.Host, mut)
			return err
		}
	}
	seed := func(mut func(*Record) error) func(release.Store, Record) error {
		return func(s release.Store, r Record) error {
			_, err := SetRecordSeed(ctx, s, r.Account, r.App, r.Host, mut)
			return err
		}
	}
	deadline := at.Add(time.Hour)
	moveID := "mv_" + strings.Repeat("a", 26)
	seedID, err := seeds.NewID()
	if err != nil {
		t.Fatal(err)
	}
	importRecord := base()
	importRecord.Seed = &SeedSource{ID: seedID, Kind: seeds.KindImport, State: seeds.StateAwaitingUpload}
	abortRecord := base()
	abortRecord.Seed = &SeedSource{ID: seedID, Kind: seeds.KindImport, State: seeds.StateAwaitingUpload}

	for _, path := range []struct {
		name  string
		rec   Record
		steps []step
	}{
		{"the six run transitions", base(), []step{
			{"RunAdmit, unkeyed", run(RunAdmit(RunAdmission{ID: "run_a", Deadline: deadline}))},
			{"RunPick", run(RunPick("run_a", "box-1", "boot-1", deadline))},
			{"RunCancelTransition, picked", run(RunCancelTransition("run_a"))},
			{"RunFinish", run(RunFinish("run_a"))},
			{"RunClear", run(RunClear("run_a"))},
			{"RunAdmit, unkeyed, again", run(RunAdmit(RunAdmission{ID: "run_b", Deadline: deadline}))},
			{"RunExpire", run(RunExpire("run_b"))},
			{"RunAdmit replacing a terminal run", run(RunAdmit(RunAdmission{ID: "run_c", Deadline: deadline}))},
			{"RunCancelTransition, queued", run(RunCancelTransition("run_c"))},
			{"RunClear, cancelled", run(RunClear("run_c"))},
		}},
		{"the ordinary writers", base(), []step{
			{"SetRecordRegion", func(s release.Store, r Record) error { return SetRecordRegion(ctx, s, r.Account, r.App, r.Host, "uk") }},
			{"UpdateRecord", func(s release.Store, r Record) error {
				cur, etag, err := GetRecordETag(ctx, s, r.Account, r.App, r.Host)
				if err != nil {
					return err
				}
				cur.Channel = "beta"
				return UpdateRecord(ctx, s, *cur, etag)
			}},
		}},
		{"a legacy move to its end", base(), []step{
			{"MoveRequest", move(MoveRequest("uk", at))},
			{"MovePromoteReady", move(MovePromoteReady(at))},
			{"MoveMarkCutting", move(MoveMarkCutting())},
			{"MoveComplete", move(MoveComplete())},
		}},
		{"a legacy move aborted", base(), []step{
			{"MoveRequest", move(MoveRequest("uk", at))},
			{"MoveAbort", move(MoveAbort())},
			{"MoveAbortComplete", move(MoveAbortComplete())},
		}},
		{"a phase-1 move start", base(), []step{
			{"MoveStartTransition", move(MoveStartTransition(MoveStart{ID: moveID, Incarnation: strings.Repeat("a", 64), From: "ie", To: "uk", Cut: "auto", Now: at}))},
		}},
		{"an offline move aborted", base(), []step{
			{"OfflineStart", move(OfflineStart(moveID, "ie", "uk", "paul", at))},
			{"OfflineAbort", move(OfflineAbort(moveID))},
			{"OfflineFinish", move(OfflineFinish(moveID, at))},
		}},
		{"an import to its commit", importRecord, []step{
			{"SeedUploaded", seed(SeedUploaded(seedID))},
			{"SeedStopped", seed(SeedStopped(seedID))},
			{"SeedSwapped", seed(SeedSwapped(seedID))},
			{"SeedPlace", seed(SeedPlace(seedID))},
			{"SeedCommit", seed(SeedCommit(seedID))},
		}},
		{"an import aborted", abortRecord, []step{
			{"SeedAbort", seed(SeedAbort(seedID))},
		}},
	} {
		t.Run(path.name, func(t *testing.T) {
			s, r, want := fresh(t, path.rec)
			for _, st := range path.steps {
				if err := st.apply(s, r); err != nil {
					t.Fatalf("%s: %v", st.name, err)
				}
				if got := keysJSON(t, s, r); got != want {
					t.Fatalf("%s: run_keys %s, want %s", st.name, got, want)
				}
			}
		})
	}
}

// TestOnlyRunAdmitWritesRunKeys reads the package's source: RunAdmit is the
// one function that assigns RunKeys, and no function overwrites a whole
// record through a pointer, which would drop run_keys with everything else
// it did not copy. Every other writer mutates the record it read in place,
// which is what TestRunKeysSurviveEveryRecordMutation relies on.
func TestOnlyRunAdmitWritesRunKeys(t *testing.T) {
	files, err := filepath.Glob("*.go")
	if err != nil {
		t.Fatal(err)
	}
	fset := token.NewFileSet()
	for _, file := range files {
		if strings.HasSuffix(file, "_test.go") {
			continue
		}
		src, err := os.ReadFile(file)
		if err != nil {
			t.Fatal(err)
		}
		f, err := parser.ParseFile(fset, file, src, 0)
		if err != nil {
			t.Fatal(err)
		}
		for _, decl := range f.Decls {
			fn, ok := decl.(*ast.FuncDecl)
			if !ok || fn.Body == nil {
				continue
			}
			ast.Inspect(fn.Body, func(n ast.Node) bool {
				as, ok := n.(*ast.AssignStmt)
				if !ok {
					return true
				}
				for _, lhs := range as.Lhs {
					switch l := lhs.(type) {
					case *ast.SelectorExpr:
						if l.Sel.Name == "RunKeys" && !(file == "run.go" && fn.Name.Name == "RunAdmit") {
							t.Errorf("%s: %s assigns RunKeys; only RunAdmit may", fset.Position(l.Pos()), fn.Name.Name)
						}
					case *ast.StarExpr:
						if id, ok := l.X.(*ast.Ident); ok && (id.Name == "rec" || id.Name == "r") {
							t.Errorf("%s: %s overwrites a whole record through *%s", fset.Position(l.Pos()), fn.Name.Name, id.Name)
						}
					}
				}
				return true
			})
		}
	}
}