package boxchoice

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"testing"

	"github.com/carlosframework/platform/internal/boxid"
)

const testKeyFP = "sha256:test-key"

// haBox is an eligible box's edges file in placement "uk", as JSON, so the
// tests also hold Box's tags to the edges file's.
func haBox(id string, mut func(m map[string]any)) string {
	m := map[string]any{
		"v": 1, "box": id, "box_id": id, "eip": "1.2.3.4", "region": "eu-west-2", "placement": "uk",
		"hosts": map[string]any{}, "hibernation_configured": true, "serves_shared_pool": true,
		"accepts_inbound": true, "secrets_key_fp": testKeyFP, "capabilities": []string{"region-ha-1"},
		"headroom": map[string]any{"mem_free_mib": 1000, "disk_free_mib": 1000, "awake": 1},
	}
	if mut != nil {
		mut(m)
	}
	b, _ := json.Marshal(m)
	return string(b)
}

// views is a complete snapshot of fresh boxes.
func views(t *testing.T, bodies ...string) Snapshot {
	t.Helper()
	snap := Snapshot{Complete: true}
	for _, body := range bodies {
		var b Box
		if err := json.Unmarshal([]byte(body), &b); err != nil {
			t.Fatal(err)
		}
		b.Fresh = true
		snap.Boxes = append(snap.Boxes, b)
	}
	return snap
}

func markStale(snap Snapshot, box string) Snapshot {
	for i := range snap.Boxes {
		if snap.Boxes[i].Box == box {
			snap.Boxes[i].Fresh = false
		}
	}
	return snap
}

func ids(cs []Candidate) []string {
	out := make([]string, len(cs))
	for i, c := range cs {
		out[i] = c.BoxID
	}
	return out
}

// rosterFor enrolls each box in placement "uk" as an instance in eu-west-2.
func rosterFor(boxes ...string) map[string]RosterBox {
	out := map[string]RosterBox{}
	for _, b := range boxes {
		out[b] = RosterBox{Placement: "uk", Cloud: boxid.Cloud{Provider: "aws", Account: "1", Region: "eu-west-2", InstanceID: "i-" + b}}
	}
	return out
}

func TestEligibleAppliesThePredicateOneBoxAtATime(t *testing.T) {
	del := func(k string) func(map[string]any) { return func(m map[string]any) { delete(m, k) } }
	set := func(k string, v any) func(map[string]any) { return func(m map[string]any) { m[k] = v } }
	cases := []struct {
		name string
		mut  func(map[string]any)
		want bool
	}{
		{"all present", nil, true},
		{"no dial address", set("eip", ""), false},
		{"no box id", set("box_id", ""), false},
		{"hibernation unconfigured", set("hibernation_configured", false), false},
		{"hibernation unknown", del("hibernation_configured"), false},
		{"refuses the shared pool", set("serves_shared_pool", false), false},
		{"shared pool unknown", del("serves_shared_pool"), false},
		{"refuses inbound", set("accepts_inbound", false), false},
		{"inbound unknown", del("accepts_inbound"), false},
		{"wrong secrets key", set("secrets_key_fp", "sha256:other"), false},
		{"no capability", del("capabilities"), false},
		{"draining", set("draining", true), false},
		{"another placement", set("placement", "us"), false},
		{"box_id names no roster entry", set("box_id", "i-0zzz"), false},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			snap := views(t, haBox("i-0aaa", tc.mut))
			got := Eligible(snap, rosterFor("i-0aaa"), "uk", "eu-west-2", "ie", testKeyFP, nil)
			if (len(got) == 1) != tc.want {
				t.Fatalf("eligible = %v, want %v", ids(got), tc.want)
			}
		})
	}
}

func TestEligibleCarriesTheHeadroomAndTheRostersRegion(t *testing.T) {
	got := Eligible(views(t, haBox("i-0aaa", nil)), rosterFor("i-0aaa"), "uk", "eu-west-2", "ie", testKeyFP, nil)
	want := Candidate{BoxID: "i-0aaa", Region: "eu-west-2", MemFreeMiB: 1000, DiskFreeMiB: 1000, Awake: 1}
	if len(got) != 1 || got[0] != want {
		t.Fatalf("got %+v, want %+v", got, want)
	}
}

func TestEligibleNeedsARosterEntryFreshnessAndNotLost(t *testing.T) {
	snap := views(t, haBox("i-0aaa", nil), haBox("i-0bbb", nil), haBox("i-0ccc", nil))
	snap = markStale(snap, "i-0bbb")
	got := Eligible(snap, rosterFor("i-0aaa", "i-0bbb"), "uk", "eu-west-2", "ie", testKeyFP, nil)
	if len(got) != 1 || got[0].BoxID != "i-0aaa" {
		t.Fatalf("got %v, want only i-0aaa (i-0bbb is stale, i-0ccc has no roster entry)", ids(got))
	}
	got = Eligible(snap, rosterFor("i-0aaa"), "uk", "eu-west-2", "ie", testKeyFP, map[string]bool{"i-0aaa": true})
	if len(got) != 0 {
		t.Fatalf("a lost box is eligible: %v", ids(got))
	}
}

// Nearest-region's rule for boxes: one box holding the key and another
// accepting inbound do not add up to one.
func TestEligibleNeverCombinesTwoHalfBoxes(t *testing.T) {
	snap := views(t,
		haBox("i-0aaa", func(m map[string]any) { m["secrets_key_fp"] = "" }),
		haBox("i-0bbb", func(m map[string]any) { m["accepts_inbound"] = false }))
	if got := Eligible(snap, rosterFor("i-0aaa", "i-0bbb"), "uk", "eu-west-2", "ie", testKeyFP, nil); len(got) != 0 {
		t.Fatalf("got %v, want none", ids(got))
	}
}

func TestEligibleFromAnIncompleteSnapshotOrWithoutAKeyOrRegionIsNothing(t *testing.T) {
	snap := views(t, haBox("i-0aaa", nil))
	snap.Complete = false
	if got := Eligible(snap, rosterFor("i-0aaa"), "uk", "eu-west-2", "ie", testKeyFP, nil); len(got) != 0 {
		t.Fatalf("got %v from an incomplete snapshot", ids(got))
	}
	if got := Eligible(views(t, haBox("i-0aaa", nil)), rosterFor("i-0aaa"), "uk", "eu-west-2", "ie", "", nil); len(got) != 0 {
		t.Fatalf("got %v with no expected key", ids(got))
	}
	if got := Eligible(views(t, haBox("i-0aaa", nil)), rosterFor("i-0aaa"), "uk", "", "ie", testKeyFP, nil); len(got) != 0 {
		t.Fatalf("got %v with no expected region", ids(got))
	}
}

func TestEligibleHomePlacementMatchesTheEmptyCode(t *testing.T) {
	home := haBox("i-0aaa", func(m map[string]any) { m["placement"] = "" })
	named := haBox("i-0bbb", func(m map[string]any) { m["placement"] = "ie" })
	snap := views(t, home, named)
	roster := rosterFor("i-0aaa", "i-0bbb")
	for _, b := range []string{"i-0aaa", "i-0bbb"} {
		rb := roster[b]
		rb.Placement = "ie"
		roster[b] = rb
	}
	if got := Eligible(snap, roster, "", "eu-west-2", "ie", testKeyFP, nil); len(got) != 2 {
		t.Fatalf("got %v, want both home boxes", ids(got))
	}
	if got := Eligible(snap, roster, "ie", "eu-west-2", "ie", testKeyFP, nil); len(got) != 2 {
		t.Fatalf("got %v, want both home boxes by the home's code", ids(got))
	}
}

// A box enrolled in another AWS region than the placement's is not in that
// placement, even when its own file says it is.
func TestEligibleRefusesABoxEnrolledInAnotherRegion(t *testing.T) {
	snap := views(t, haBox("i-0aaa", nil), haBox("i-0bbb", nil))
	roster := rosterFor("i-0aaa", "i-0bbb")
	roster["i-0bbb"] = RosterBox{Placement: "uk", Cloud: boxid.Cloud{Provider: "aws", Account: "1", Region: "us-east-1", InstanceID: "i-0bbb"}}
	got := Eligible(snap, roster, "uk", "eu-west-2", "ie", testKeyFP, nil)
	if len(got) != 1 || got[0].BoxID != "i-0aaa" {
		t.Fatalf("got %v, want only i-0aaa", ids(got))
	}
	// And a file that advertises another AWS region than the one expected.
	snap = views(t, haBox("i-0aaa", func(m map[string]any) { m["region"] = "us-east-1" }))
	if got := Eligible(snap, rosterFor("i-0aaa"), "uk", "eu-west-2", "ie", testKeyFP, nil); len(got) != 0 {
		t.Fatalf("got %v, want none", ids(got))
	}
}

// The roster's enrolled placement is the enrollment's word and must agree
// with the file's; an incomplete cloud is no identity at all.
func TestEligibleNeedsTheRosterToAgree(t *testing.T) {
	snap := views(t, haBox("i-0aaa", nil), haBox("i-0bbb", nil))
	roster := rosterFor("i-0aaa", "i-0bbb")
	roster["i-0aaa"] = RosterBox{Placement: "us", Cloud: roster["i-0aaa"].Cloud}
	roster["i-0bbb"] = RosterBox{Placement: "uk", Cloud: boxid.Cloud{Provider: "aws", Region: "eu-west-2", InstanceID: "i-0bbb"}}
	if got := Eligible(snap, roster, "uk", "eu-west-2", "ie", testKeyFP, nil); len(got) != 0 {
		t.Fatalf("got %v, want none", ids(got))
	}
}

func TestEligibleSortsByBoxID(t *testing.T) {
	snap := views(t, haBox("i-0ccc", nil), haBox("i-0aaa", nil), haBox("i-0bbb", nil))
	got := ids(Eligible(snap, rosterFor("i-0aaa", "i-0bbb", "i-0ccc"), "uk", "eu-west-2", "ie", testKeyFP, nil))
	if len(got) != 3 || got[0] != "i-0aaa" || got[1] != "i-0bbb" || got[2] != "i-0ccc" {
		t.Fatalf("got %v", got)
	}
}

func TestChoose(t *testing.T) {
	a, b := Candidate{BoxID: "i-0aaa"}, Candidate{BoxID: "i-0bbb"}
	load := func(hosts int, apps ...string) Load {
		l := Load{Hosts: hosts, Apps: map[string]int{}}
		for _, x := range apps {
			l.Apps[x]++
		}
		return l
	}
	const host = "shop.bab.oncarlos.com"
	if _, ok := Choose(host, nil, nil, "bab/jam"); ok {
		t.Fatal("chose from nothing")
	}
	got, _ := Choose(host, []Candidate{a, b}, map[string]Load{"i-0aaa": load(5), "i-0bbb": load(2)}, "bab/jam")
	if got.BoxID != "i-0bbb" {
		t.Fatalf("fewest hosts: got %s", got.BoxID)
	}
	got, _ = Choose(host, []Candidate{a, b}, map[string]Load{"i-0aaa": load(5), "i-0bbb": load(2, "bab/jam")}, "bab/jam")
	if got.BoxID != "i-0aaa" {
		t.Fatalf("spread beats load: got %s", got.BoxID)
	}
	got, ok := Choose(host, []Candidate{a, b}, map[string]Load{"i-0aaa": load(5, "bab/jam"), "i-0bbb": load(2, "bab/jam", "bab/jam")}, "bab/jam")
	if !ok || got.BoxID != "i-0bbb" {
		t.Fatalf("every box runs the app: fewest hosts overall, never a refusal; got %s ok=%v", got.BoxID, ok)
	}
}

// H12, claim P5-D2: an equal load breaks by the greatest SHA-256 of
// host + "\x00" + reference, so the choice differs host by host instead
// of always landing on the lowest name, and does not depend on the order
// the candidates arrive in.
func TestChooseBreaksATieByTheHostsHash(t *testing.T) {
	cands := []Candidate{{BoxID: "annie-welding"}, {BoxID: "bert-sawing"}, {BoxID: "cleo-mixing"}}
	want := func(host string) string {
		best, bestSum := "", ""
		for _, c := range cands {
			sum := sha256.Sum256([]byte(host + "\x00" + c.BoxID))
			if h := hex.EncodeToString(sum[:]); h > bestSum {
				best, bestSum = c.BoxID, h
			}
		}
		return best
	}
	seen := map[string]bool{}
	for _, host := range []string{"a.bab.oncarlos.com", "b.bab.oncarlos.com", "c.bab.oncarlos.com", "d.bab.oncarlos.com", "e.bab.oncarlos.com", "f.bab.oncarlos.com"} {
		forward, _ := Choose(host, cands, nil, "bab/jam")
		backward, _ := Choose(host, []Candidate{cands[2], cands[1], cands[0]}, nil, "bab/jam")
		if forward.BoxID != want(host) || backward.BoxID != want(host) {
			t.Fatalf("%s: chose %s and %s, want %s", host, forward.BoxID, backward.BoxID, want(host))
		}
		seen[forward.BoxID] = true
	}
	if len(seen) < 2 {
		t.Fatalf("six hosts all chose %v: the tie-break is not per host", seen)
	}
}