package main

import (
	"encoding/json"
	"testing"
)

// recs builds a backendResult from name + literal JSON element strings.
func recs(t *testing.T, name string, elems ...string) backendResult {
	t.Helper()
	body := "[" + join(elems) + "]"
	r, err := decodeRecords([]byte(body))
	if err != nil {
		t.Fatalf("decodeRecords(%s): %v", body, err)
	}
	return backendResult{name: name, records: r}
}

func join(elems []string) string {
	out := ""
	for i, e := range elems {
		if i > 0 {
			out += ","
		}
		out += e
	}
	return out
}

// certnames extracts the certname field from a merged result set.
func certnames(t *testing.T, raws []json.RawMessage) []string {
	t.Helper()
	var out []string
	for _, r := range raws {
		var m recordMeta
		if err := json.Unmarshal(r, &m); err != nil {
			t.Fatalf("unmarshal %s: %v", r, err)
		}
		out = append(out, m.Certname)
	}
	return out
}

// factValues extracts "certname:name=value" for /facts records to assert which
// backend's facts survived.
func factValues(t *testing.T, raws []json.RawMessage) []string {
	t.Helper()
	var out []string
	for _, r := range raws {
		var m struct {
			Certname string `json:"certname"`
			Name     string `json:"name"`
			Value    string `json:"value"`
		}
		if err := json.Unmarshal(r, &m); err != nil {
			t.Fatalf("unmarshal %s: %v", r, err)
		}
		out = append(out, m.Certname+":"+m.Name+"="+m.Value)
	}
	return out
}

func node(cn, ts string) string {
	return `{"certname":"` + cn + `","report_timestamp":"` + ts + `","latest_report_status":"changed"}`
}

func fact(cn, name, val, ts string) string {
	// facts records don't carry report_timestamp in real PuppetDB, but including
	// it is harmless and lets a couple of tests reuse the same helper. Merge
	// attribution for facts comes from the owner func, not the record.
	if ts == "" {
		return `{"certname":"` + cn + `","name":"` + name + `","value":"` + val + `"}`
	}
	return `{"certname":"` + cn + `","name":"` + name + `","value":"` + val + `","report_timestamp":"` + ts + `"}`
}

// report builds a /reports record with the fields the merge and ordering paths
// care about.
func report(cn, hash, receive string) string {
	return `{"certname":"` + cn + `","hash":"` + hash + `","receive_time":"` + receive +
		`","end_time":"` + receive + `","status":"changed","environment":"production"}`
}

// event builds an /events record, which carries its report's hash but no id of
// its own.
func event(cn, reportHash, resource string) string {
	return `{"certname":"` + cn + `","report":"` + reportHash + `","resource_title":"` + resource +
		`","status":"success","timestamp":"2026-07-01T00:00:00Z"}`
}

func TestMergeNodes_NewerWins(t *testing.T) {
	a := recs(t, "a", node("h1", "2026-07-01T00:00:00Z"), node("h2", "2026-07-10T00:00:00Z"))
	b := recs(t, "b", node("h1", "2026-07-20T00:00:00Z"), node("h3", "2026-07-05T00:00:00Z"))

	merged := mergeNodes([]backendResult{a, b}, nil)
	got := map[string]string{}
	for _, r := range merged {
		var m recordMeta
		_ = json.Unmarshal(r, &m)
		got[m.Certname] = m.ReportTimestamp
	}
	if got["h1"] != "2026-07-20T00:00:00Z" {
		t.Errorf("h1: newer (b) should win, got %s", got["h1"])
	}
	if got["h2"] != "2026-07-10T00:00:00Z" {
		t.Errorf("h2: only in a, got %s", got["h2"])
	}
	if got["h3"] != "2026-07-05T00:00:00Z" {
		t.Errorf("h3: only in b, got %s", got["h3"])
	}
	if len(merged) != 3 {
		t.Errorf("expected 3 deduped nodes, got %d", len(merged))
	}
}

func TestMergeNodes_OneBackendOnly(t *testing.T) {
	a := recs(t, "a", node("h1", "2026-07-01T00:00:00Z"))
	// b returned nothing (e.g. empty result).
	b := backendResult{name: "b"}
	merged := mergeNodes([]backendResult{a, b}, nil)
	if len(merged) != 1 || certnames(t, merged)[0] != "h1" {
		t.Fatalf("expected only h1, got %v", certnames(t, merged))
	}
}

func TestMergeNodes_TieKeepsEarlierBackend(t *testing.T) {
	// Equal timestamps: the backend listed first wins, as a tie-break.
	first := recs(t, "b", node("h1", "2026-07-01T00:00:00Z"))
	second := recs(t, "a", node("h1", "2026-07-01T00:00:00Z"))
	merged := mergeNodes([]backendResult{first, second}, nil)
	if len(merged) != 1 {
		t.Fatalf("expected 1 record, got %d", len(merged))
	}
	// Ensure the kept record is the first backend's (identical here, but assert count/dedupe).
	if certnames(t, merged)[0] != "h1" {
		t.Fatalf("expected h1")
	}
}

func TestMergeNodes_PreservesUnknownFields(t *testing.T) {
	a := recs(t, "a", `{"certname":"h1","report_timestamp":"2026-07-01T00:00:00Z","extra":{"deep":42}}`)
	merged := mergeNodes([]backendResult{a}, nil)
	if len(merged) != 1 {
		t.Fatalf("expected 1 record")
	}
	var m map[string]json.RawMessage
	_ = json.Unmarshal(merged[0], &m)
	if _, ok := m["extra"]; !ok {
		t.Fatalf("unknown field 'extra' was dropped: %s", merged[0])
	}
}

func TestMergeFacts_NilOwnerUsesConfiguredOrder(t *testing.T) {
	// Static merge passes no owner: h1 is in both, so the first backend in the
	// slice supplies its facts.
	a := recs(t, "a", fact("h1", "role", "web-a", ""), fact("h2", "role", "db-a", ""))
	b := recs(t, "b", fact("h1", "role", "web-b", ""))

	merged := mergeFacts([]backendResult{b, a}, nil, nil)
	got := factValues(t, merged)
	assertContains(t, got, "h1:role=web-b")
	assertNotContains(t, got, "h1:role=web-a")
	// h2 only in a -> still served from a.
	assertContains(t, got, "h2:role=db-a")
}

func TestMergeFacts_Freshness_NewerBackendWins(t *testing.T) {
	// owner map says h1 belongs to a and h2 to b. Multiple facts per node must
	// all come from the winner.
	a := recs(t, "a",
		fact("h1", "role", "web-a", ""), fact("h1", "ip", "10.0.0.1", ""),
		fact("h2", "role", "db-a", ""))
	b := recs(t, "b",
		fact("h1", "role", "web-b", ""), fact("h1", "ip", "10.9.9.9", ""),
		fact("h2", "role", "db-b", ""), fact("h2", "ip", "10.0.0.2", ""))

	owner := func(cn string) string {
		if cn == "h1" {
			return "a"
		}
		return "b"
	}
	merged := mergeFacts([]backendResult{b, a}, owner, nil)
	got := factValues(t, merged)
	// h1 -> all a facts, no b facts.
	assertContains(t, got, "h1:role=web-a")
	assertContains(t, got, "h1:ip=10.0.0.1")
	assertNotContains(t, got, "h1:role=web-b")
	assertNotContains(t, got, "h1:ip=10.9.9.9")
	// h2 -> all b facts.
	assertContains(t, got, "h2:role=db-b")
	assertContains(t, got, "h2:ip=10.0.0.2")
	assertNotContains(t, got, "h2:role=db-a")
}

func TestMergeFacts_OwnerMissingFallsBackToConfiguredOrder(t *testing.T) {
	// owner returns a backend with no facts for h1 -> fall back to the first
	// backend in the slice that has some.
	first := recs(t, "b", fact("h1", "role", "web-b", ""))
	second := recs(t, "a", fact("h1", "role", "web-a", ""))
	merged := mergeFacts([]backendResult{first, second}, func(string) string { return "ghost" }, nil)
	got := factValues(t, merged)
	assertContains(t, got, "h1:role=web-b") // b is first in slice
	assertNotContains(t, got, "h1:role=web-a")
}

func TestBuildFreshness(t *testing.T) {
	// a holds the newer report for h1; b holds the newer one for h2.
	a := recs(t, "a", node("h1", "2026-07-20T00:00:00Z"), node("h2", "2026-07-01T00:00:00Z"))
	b := recs(t, "b", node("h1", "2026-07-01T00:00:00Z"), node("h2", "2026-07-20T00:00:00Z"))
	f := buildFreshness([]backendResult{a, b})
	if f["h1"] != "a" {
		t.Errorf("h1 should belong to a, got %q", f["h1"])
	}
	if f["h2"] != "b" {
		t.Errorf("h2 should belong to b, got %q", f["h2"])
	}
}

func TestDecodeRecords_NotArray(t *testing.T) {
	if _, err := decodeRecords([]byte(`{"not":"array"}`)); err == nil {
		t.Fatal("expected error decoding non-array body")
	}
}

func assertContains(t *testing.T, hay []string, needle string) {
	t.Helper()
	for _, h := range hay {
		if h == needle {
			return
		}
	}
	t.Errorf("expected %q in %v", needle, hay)
}

func assertNotContains(t *testing.T, hay []string, needle string) {
	t.Helper()
	for _, h := range hay {
		if h == needle {
			t.Errorf("did not expect %q in %v", needle, hay)
		}
	}
}