package semanticreconcile

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"math"
	"reflect"
	"strings"
	"testing"

	"gitea.maximumdirect.net/eric/notarius/internal/core/source"
)

func TestPrepareBuildsContiguousCandidatesAndOwnedSourceContext(t *testing.T) {
	document := &source.SourceDocument{ID: "private-source-id", Units: []source.SourceUnit{
		{ID: 40, Kind: "narration", Text: "zero"},
		{ID: 10, Kind: "speech", Text: "one", Metadata: map[string]any{"speaker": map[string]any{"name": "Mira"}}},
		{ID: 70, Kind: "speech", Text: "two"},
		{ID: 20, Kind: "narration", Text: "three"},
		{ID: 90, Kind: "speech", Text: "four"},
	}}
	references := []source.SourceRef{
		{SourceID: document.ID, StartUnitID: 90, EndUnitID: 90},
		{SourceID: document.ID, StartUnitID: 10, EndUnitID: 20},
		{SourceID: document.ID, StartUnitID: 10, EndUnitID: 20},
	}
	candidates := []Candidate{
		{Label: "The Tavern", SourceRefs: append([]source.SourceRef(nil), references...)},
		{Label: "The Tavern", SourceRefs: append([]source.SourceRef(nil), references...)},
		{Label: "Broken", SourceRefs: []source.SourceRef{{SourceID: document.ID, StartUnitID: 20, EndUnitID: 10}}},
	}
	before := cloneCandidates(candidates)

	preparation, err := Prepare(document, candidates, Limits{
		ContextRadius:        1,
		MaximumCandidates:    len(candidates),
		MaximumMaterialBytes: 10000,
	})
	if err != nil {
		t.Fatal(err)
	}
	if preparation.Disposition() != Ready {
		t.Fatalf("Disposition() = %v, want Ready", preparation.Disposition())
	}
	if !reflect.DeepEqual(candidates, before) {
		t.Fatalf("Prepare() mutated candidates: %#v", candidates)
	}
	if got, want := preparation.CandidateMappings(), []CandidateMapping{
		{CandidateID: 1, CandidatePosition: 0},
		{CandidateID: 2, CandidatePosition: 1},
	}; !reflect.DeepEqual(got, want) {
		t.Fatalf("CandidateMappings() = %#v, want %#v", got, want)
	}

	materials := preparation.Materials()
	if len(materials) != 2 {
		t.Fatalf("materials = %#v, want candidates and transcript", materials)
	}
	if _, ok := materials[candidateInputName]; !ok {
		t.Fatal("candidate material is missing")
	}
	if _, ok := materials[transcriptInputName]; !ok {
		t.Fatal("transcript material is missing")
	}
	for name, material := range materials {
		if material.Name != name || material.MediaType != "application/json" || material.OriginURI != "" || material.SizeBytes != int64(len(material.Content)) {
			t.Fatalf("material %q metadata = %#v", name, material)
		}
		digest := sha256.Sum256(material.Content)
		if want := "sha256:" + hex.EncodeToString(digest[:]); material.Digest != want {
			t.Fatalf("material %q digest = %q, want %q", name, material.Digest, want)
		}
	}

	candidateContent := append([]byte(nil), materials[candidateInputName].Content...)
	var candidatePayload candidateInput
	if err := json.Unmarshal(candidateContent, &candidatePayload); err != nil {
		t.Fatal(err)
	}
	wantCandidates := []visibleCandidate{
		{CandidateID: 1, Label: "The Tavern", SourceRefs: []sourceRange{{StartUnitID: 10, EndUnitID: 20}, {StartUnitID: 90, EndUnitID: 90}}},
		{CandidateID: 2, Label: "The Tavern", SourceRefs: []sourceRange{{StartUnitID: 10, EndUnitID: 20}, {StartUnitID: 90, EndUnitID: 90}}},
	}
	if !reflect.DeepEqual(candidatePayload.Candidates, wantCandidates) {
		t.Fatalf("candidate payload = %#v, want %#v", candidatePayload.Candidates, wantCandidates)
	}
	var candidateObjects struct {
		Candidates []map[string]json.RawMessage `json:"candidates"`
	}
	if err := json.Unmarshal(candidateContent, &candidateObjects); err != nil {
		t.Fatal(err)
	}
	for _, candidate := range candidateObjects.Candidates {
		if len(candidate) != 3 || candidate["candidate_id"] == nil || candidate["label"] == nil || candidate["source_refs"] == nil {
			t.Fatalf("model-facing candidate fields = %#v", candidate)
		}
	}
	combined := string(materials[candidateInputName].Content) + string(materials[transcriptInputName].Content)
	for _, forbidden := range []string{document.ID, "application_entity_id", "private-entity-id"} {
		if strings.Contains(combined, forbidden) {
			t.Fatalf("model material leaked %q: %s", forbidden, combined)
		}
	}

	var transcript transcriptInput
	if err := json.Unmarshal(materials[transcriptInputName].Content, &transcript); err != nil {
		t.Fatal(err)
	}
	if len(transcript.Windows) != 1 || len(transcript.Windows[0].Units) != len(document.Units) {
		t.Fatalf("windows = %#v, want one coalesced source window", transcript.Windows)
	}
	for index, wantID := range []int{40, 10, 70, 20, 90} {
		if transcript.Windows[0].Units[index].ID != wantID {
			t.Fatalf("unit %d id = %d, want %d", index, transcript.Windows[0].Units[index].ID, wantID)
		}
	}
	if transcript.Windows[0].Units[0].Cited {
		t.Fatal("radius-only unit marked cited")
	}
	for index := 1; index < len(transcript.Windows[0].Units); index++ {
		if !transcript.Windows[0].Units[index].Cited {
			t.Fatalf("evidence unit %d was not marked cited", index)
		}
	}

	candidates[0].SourceRefs[0].StartUnitID = 40
	document.Units[1].Metadata["speaker"].(map[string]any)["name"] = "changed"
	if got := preparation.Materials()[candidateInputName].Content; !reflect.DeepEqual(got, candidateContent) {
		t.Fatalf("candidate material changed through caller input: %s", got)
	}
	var retained transcriptInput
	if err := json.Unmarshal(preparation.Materials()[transcriptInputName].Content, &retained); err != nil {
		t.Fatal(err)
	}
	if got := retained.Windows[0].Units[1].Metadata["speaker"].(map[string]any)["name"]; got != "Mira" {
		t.Fatalf("retained metadata = %v, want Mira", got)
	}

	returnedMappings := preparation.CandidateMappings()
	returnedMappings[0].CandidatePosition = 99
	returnedMaterials := preparation.Materials()
	candidateMaterial := returnedMaterials[candidateInputName]
	candidateMaterial.Content[0] = '['
	returnedMaterials[candidateInputName] = candidateMaterial
	delete(returnedMaterials, transcriptInputName)
	if preparation.CandidateMappings()[0].CandidatePosition != 0 || !json.Valid(preparation.Materials()[candidateInputName].Content) || len(preparation.Materials()) != 2 {
		t.Fatal("preparation accessors exposed retained data")
	}
}

func TestPrepareRedactsInvalidSourceMetadata(t *testing.T) {
	const sensitiveKey = "sensitive-metadata-key"
	document := &source.SourceDocument{ID: "private-source", Units: []source.SourceUnit{
		{ID: 1, Text: "private transcript", Metadata: map[string]any{sensitiveKey: math.NaN()}},
		{ID: 2, Text: "other private transcript"},
	}}
	candidates := []Candidate{
		{Label: "Private One", SourceRefs: []source.SourceRef{{SourceID: document.ID, StartUnitID: 1, EndUnitID: 1}}},
		{Label: "Private Two", SourceRefs: []source.SourceRef{{SourceID: document.ID, StartUnitID: 2, EndUnitID: 2}}},
	}
	_, err := Prepare(document, candidates, DefaultLimits())
	if err == nil || !strings.Contains(err.Error(), "invalid source metadata") {
		t.Fatalf("Prepare() error = %v, want redacted metadata failure", err)
	}
	for _, forbidden := range []string{sensitiveKey, document.ID, document.Units[0].Text, candidates[0].Label, "non-finite", "float64"} {
		if strings.Contains(err.Error(), forbidden) {
			t.Fatalf("Prepare() error leaked %q: %v", forbidden, err)
		}
	}
}

func TestPrepareFiltersUnsafeCandidatesAndCoalescesAdjacentWindows(t *testing.T) {
	document := &source.SourceDocument{ID: "session", Units: []source.SourceUnit{
		{ID: 9}, {ID: 3}, {ID: 8}, {ID: 1}, {ID: 7},
	}}
	candidates := []Candidate{
		{Label: "One", SourceRefs: []source.SourceRef{{SourceID: document.ID, StartUnitID: 3, EndUnitID: 3}}},
		{Label: "Two", SourceRefs: []source.SourceRef{{SourceID: document.ID, StartUnitID: 8, EndUnitID: 8}}},
		{Label: "Missing", SourceRefs: []source.SourceRef{{SourceID: document.ID, StartUnitID: 99, EndUnitID: 99}}},
		{Label: "Foreign", SourceRefs: []source.SourceRef{{SourceID: "other", StartUnitID: 1, EndUnitID: 1}}},
		{Label: "No references"},
		{Label: "Partly invalid", SourceRefs: []source.SourceRef{
			{SourceID: document.ID, StartUnitID: 1, EndUnitID: 1},
			{SourceID: document.ID, StartUnitID: 100, EndUnitID: 100},
		}},
	}
	limits := Limits{ContextRadius: 0, MaximumCandidates: len(candidates), MaximumMaterialBytes: 10000}

	preparation, err := Prepare(document, candidates, limits)
	if err != nil || preparation.Disposition() != Ready {
		t.Fatalf("Prepare() disposition = %v, error = %v", preparation.Disposition(), err)
	}
	if got, want := preparation.CandidateMappings(), []CandidateMapping{
		{CandidateID: 1, CandidatePosition: 0},
		{CandidateID: 2, CandidatePosition: 1},
	}; !reflect.DeepEqual(got, want) {
		t.Fatalf("CandidateMappings() = %#v, want %#v", got, want)
	}
	var transcript transcriptInput
	if err := json.Unmarshal(preparation.Materials()[transcriptInputName].Content, &transcript); err != nil {
		t.Fatal(err)
	}
	if len(transcript.Windows) != 1 || len(transcript.Windows[0].Units) != 2 || transcript.Windows[0].Units[0].ID != 3 || transcript.Windows[0].Units[1].ID != 8 {
		t.Fatalf("windows = %#v, want adjacent source-order units coalesced", transcript.Windows)
	}
	oneCandidate, err := Prepare(document, candidates[:1], limits)
	if err != nil {
		t.Fatal(err)
	}
	if got, want := oneCandidate.CandidateMappings(), []CandidateMapping{{CandidateID: 1, CandidatePosition: 0}}; oneCandidate.Disposition() != InsufficientCandidates || !reflect.DeepEqual(got, want) || len(oneCandidate.Materials()) != 0 {
		t.Fatalf("Prepare(one candidate) = disposition %v, mappings %#v, materials %#v", oneCandidate.Disposition(), got, oneCandidate.Materials())
	}

	nilPreparation, err := Prepare(nil, candidates, limits)
	if err != nil {
		t.Fatal(err)
	}
	if nilPreparation.Disposition() != InsufficientCandidates || len(nilPreparation.CandidateMappings()) != 0 || len(nilPreparation.Materials()) != 0 {
		t.Fatalf("Prepare(nil) = disposition %v, mappings %#v, materials %#v", nilPreparation.Disposition(), nilPreparation.CandidateMappings(), nilPreparation.Materials())
	}
}

func TestPrepareValidatesLimitsBeforeBuildingMaterials(t *testing.T) {
	if err := DefaultLimits().Validate(); err != nil {
		t.Fatalf("DefaultLimits().Validate() error = %v", err)
	}
	cycle := map[string]any{}
	cycle["self"] = cycle
	document := &source.SourceDocument{ID: "session", Units: []source.SourceUnit{
		{ID: 1, Metadata: cycle}, {ID: 2},
	}}
	candidates := candidatesForEveryUnit(document)
	tests := []struct {
		name   string
		limits Limits
		want   string
	}{
		{name: "negative radius", limits: Limits{ContextRadius: -1, MaximumCandidates: 2, MaximumMaterialBytes: 100}, want: "radius"},
		{name: "zero candidates", limits: Limits{ContextRadius: 0, MaximumCandidates: 0, MaximumMaterialBytes: 100}, want: "candidates"},
		{name: "negative candidates", limits: Limits{ContextRadius: 0, MaximumCandidates: -1, MaximumMaterialBytes: 100}, want: "candidates"},
		{name: "zero bytes", limits: Limits{ContextRadius: 0, MaximumCandidates: 2, MaximumMaterialBytes: 0}, want: "bytes"},
		{name: "negative bytes", limits: Limits{ContextRadius: 0, MaximumCandidates: 2, MaximumMaterialBytes: -1}, want: "bytes"},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			if _, err := Prepare(document, candidates, test.limits); err == nil || !strings.Contains(err.Error(), test.want) {
				t.Fatalf("Prepare() error = %v, want %q validation", err, test.want)
			}
		})
	}
}

func TestPrepareEnforcesCandidateLimitBeforeRenderingContext(t *testing.T) {
	cycle := map[string]any{}
	cycle["self"] = cycle
	document := &source.SourceDocument{ID: "session", Units: []source.SourceUnit{
		{ID: 1, Text: "one", Metadata: cycle},
		{ID: 2, Text: "two"},
		{ID: 3, Text: "three"},
	}}
	candidates := candidatesForEveryUnit(document)
	limits := Limits{ContextRadius: 0, MaximumCandidates: 2, MaximumMaterialBytes: 10000}

	exceeded, err := Prepare(document, candidates, limits)
	if err != nil {
		t.Fatalf("Prepare(over limit) error = %v; context should not be rendered", err)
	}
	if exceeded.Disposition() != LimitExceeded || len(exceeded.CandidateMappings()) != 3 || len(exceeded.Materials()) != 0 {
		t.Fatalf("Prepare(over limit) = disposition %v, mappings %#v, materials %#v", exceeded.Disposition(), exceeded.CandidateMappings(), exceeded.Materials())
	}

	document.Units[0].Metadata = nil
	exact, err := Prepare(document, candidates[:2], limits)
	if err != nil || exact.Disposition() != Ready {
		t.Fatalf("Prepare(at limit) = disposition %v, error %v", exact.Disposition(), err)
	}
}

func TestPrepareStopsRenderingContextWhenMaterialLimitIsExceeded(t *testing.T) {
	cycle := map[string]any{}
	cycle["self"] = cycle
	document := &source.SourceDocument{ID: "session", Units: []source.SourceUnit{
		{ID: 1, Text: strings.Repeat("oversized", 100)},
		{ID: 2, Text: "must not be inspected"},
	}}
	candidates := candidatesForEveryUnit(document)
	base, err := Prepare(document, candidates, Limits{
		ContextRadius:        0,
		MaximumCandidates:    len(candidates),
		MaximumMaterialBytes: 10000,
	})
	if err != nil || base.Disposition() != Ready {
		t.Fatalf("Prepare(base) = disposition %v, error %v", base.Disposition(), err)
	}
	candidateBytes := len(base.Materials()[candidateInputName].Content)
	document.Units[1].Metadata = cycle

	limited, err := Prepare(document, candidates, Limits{
		ContextRadius:        0,
		MaximumCandidates:    len(candidates),
		MaximumMaterialBytes: candidateBytes + 64,
	})
	if err != nil {
		t.Fatalf("Prepare(limited) error = %v; rendering should stop at the material bound", err)
	}
	if limited.Disposition() != LimitExceeded || len(limited.Materials()) != 0 {
		t.Fatalf("Prepare(limited) = disposition %v, materials %#v, want bounded skip", limited.Disposition(), limited.Materials())
	}
}

func TestPrepareAcceptsExactCombinedByteLimitAndSkipsOneOver(t *testing.T) {
	document := &source.SourceDocument{ID: "session", Units: []source.SourceUnit{
		{ID: 1, Text: "one"},
		{ID: 2, Text: "two"},
	}}
	candidates := candidatesForEveryUnit(document)
	baseLimits := Limits{ContextRadius: 0, MaximumCandidates: len(candidates), MaximumMaterialBytes: 10000}
	base, err := Prepare(document, candidates, baseLimits)
	if err != nil || base.Disposition() != Ready {
		t.Fatalf("Prepare(base) = disposition %v, error %v", base.Disposition(), err)
	}
	materials := base.Materials()
	totalBytes := len(materials[candidateInputName].Content) + len(materials[transcriptInputName].Content)

	exactLimits := baseLimits
	exactLimits.MaximumMaterialBytes = totalBytes
	exact, err := Prepare(document, candidates, exactLimits)
	if err != nil || exact.Disposition() != Ready {
		t.Fatalf("Prepare(exact bytes) = disposition %v, error %v", exact.Disposition(), err)
	}
	oneOverLimits := exactLimits
	oneOverLimits.MaximumMaterialBytes--
	oneOver, err := Prepare(document, candidates, oneOverLimits)
	if err != nil || oneOver.Disposition() != LimitExceeded || len(oneOver.Materials()) != 0 {
		t.Fatalf("Prepare(one over) = disposition %v, materials %#v, error %v", oneOver.Disposition(), oneOver.Materials(), err)
	}
}

func TestPrepareSerializationIsDeterministic(t *testing.T) {
	document := &source.SourceDocument{ID: "session", Units: []source.SourceUnit{
		{ID: 5, Text: "five", Metadata: map[string]any{"z": 1, "a": []any{"first", "second"}}},
		{ID: 2, Text: "two", Metadata: map[string]any{"nested": map[string]any{"b": true, "a": false}}},
	}}
	candidates := candidatesForEveryUnit(document)
	limits := Limits{ContextRadius: 0, MaximumCandidates: len(candidates), MaximumMaterialBytes: 10000}

	first, err := Prepare(document, candidates, limits)
	if err != nil {
		t.Fatal(err)
	}
	second, err := Prepare(document, cloneCandidates(candidates), limits)
	if err != nil {
		t.Fatal(err)
	}
	for _, name := range []string{candidateInputName, transcriptInputName} {
		firstMaterial := first.Materials()[name]
		secondMaterial := second.Materials()[name]
		if !reflect.DeepEqual(firstMaterial, secondMaterial) {
			t.Fatalf("material %q is not deterministic:\n%#v\n%#v", name, firstMaterial, secondMaterial)
		}
	}
}

func candidatesForEveryUnit(document *source.SourceDocument) []Candidate {
	candidates := make([]Candidate, len(document.Units))
	for index, unit := range document.Units {
		candidates[index] = Candidate{
			Label: "candidate",
			SourceRefs: []source.SourceRef{{
				SourceID:    document.ID,
				StartUnitID: unit.ID,
				EndUnitID:   unit.ID,
			}},
		}
	}
	return candidates
}

func cloneCandidates(candidates []Candidate) []Candidate {
	cloned := append([]Candidate(nil), candidates...)
	for index := range cloned {
		cloned[index].SourceRefs = append([]source.SourceRef(nil), candidates[index].SourceRefs...)
	}
	return cloned
}