package locationoccurrences

import (
	"context"
	"encoding/json"
	"errors"
	"reflect"
	"strings"
	"testing"

	"gitea.maximumdirect.net/eric/notarius/internal/core/source"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd"
	locationcodec "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/codec/locations"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/locations/identity"
)

func TestExtractMapsKindsOrdersOccurrencesAndPreservesIndependentFacts(t *testing.T) {
	locations := locationRegistry(t, "The Tavern", "The Tavern")
	first, second := locations.Locations[0], locations.Locations[1]
	client := &fakeOccurrencesLLMClient{response: extractionResponse{Occurrences: []occurrenceResponse{
		{LocationID: second.ID, Name: second.Name, Kind: "mentioned", SourceRefs: occurrenceRefs(30, 30)},
		{LocationID: first.ID, Name: first.Name, Kind: "mentioned", SourceRefs: occurrenceRefs(10, 10)},
		{LocationID: first.ID, Name: first.Name, Kind: "recalled", SourceRefs: occurrenceRefs(10, 10)},
		{LocationID: first.ID, Name: first.Name, Kind: "planned", SourceRefs: occurrenceRefs(10, 10)},
		{LocationID: first.ID, Name: first.Name, Kind: "visited", SourceRefs: append(occurrenceRefs(10, 10), occurrenceRefs(10, 10)...)},
		{LocationID: first.ID, Name: first.Name, Kind: "visited", SourceRefs: occurrenceRefs(20, 20)},
		{LocationID: first.ID, Name: first.Name, Kind: "visited", SourceRefs: occurrenceRefs(10, 10)},
	}}}
	references := registryReferences(t, locations)
	req := extractionRequest()
	req.References = references
	result, err := newExtractor(t, client, references).Extract(context.Background(), req)
	if err != nil {
		t.Fatal(err)
	}
	if len(result.Value.Occurrences) != 6 {
		t.Fatalf("occurrences = %#v, want exact duplicate removed", result.Value.Occurrences)
	}
	got := result.Value.Occurrences
	if kinds := []dnd.LocationOccurrenceKind{got[0].Kind, got[1].Kind, got[2].Kind, got[3].Kind}; !reflect.DeepEqual(kinds, []dnd.LocationOccurrenceKind{dnd.LocationOccurrenceKindVisited, dnd.LocationOccurrenceKindPlanned, dnd.LocationOccurrenceKindRecalled, dnd.LocationOccurrenceKindMentioned}) {
		t.Fatalf("same-evidence kind order = %#v", kinds)
	}
	if got[4].Kind != dnd.LocationOccurrenceKindVisited || got[4].SourceRefs[0].StartUnitID != 20 || got[5].LocationID != second.ID || got[5].SourceRefs[0].StartUnitID != 30 {
		t.Fatalf("occurrence order = %#v", got)
	}
	if !reflect.DeepEqual(got[0].SourceRefs, []source.SourceRef{{SourceID: req.Source.ID, StartUnitID: 10, EndUnitID: 10}}) {
		t.Fatalf("canonical evidence = %#v", got[0].SourceRefs)
	}
}

func TestExtractUsesIDsNamesAndCurrentTranscriptEvidenceOnly(t *testing.T) {
	locations := locationRegistry(t, "The Tavern", "The Tavern")
	first, second := locations.Locations[0], locations.Locations[1]
	client := &fakeOccurrencesLLMClient{response: extractionResponse{Occurrences: []occurrenceResponse{{
		LocationID: second.ID, Name: second.Name, Kind: "visited", SourceRefs: occurrenceRefs(10, 10),
	}}}}
	references := registryReferences(t, locations)
	req := extractionRequest()
	req.References = references
	result, err := newExtractor(t, client, references).Extract(context.Background(), req)
	if err != nil {
		t.Fatal(err)
	}
	if occurrence := result.Value.Occurrences[0]; occurrence.LocationID != second.ID || occurrence.Name != second.Name || occurrence.SourceRefs[0].SourceID != req.Source.ID {
		t.Fatalf("occurrence = %#v", occurrence)
	}
	input := client.requests[0].Inputs[LocationRegistryReferenceSlot]
	if input.Name != LocationRegistryReferenceSlot || !strings.Contains(string(input.Content), first.ID) || !strings.Contains(string(input.Content), second.ID) {
		t.Fatalf("location prompt input = %#v", input)
	}
	for _, forbidden := range []string{"source_refs", "source_id", "other-session"} {
		if strings.Contains(string(input.Content), forbidden) {
			t.Fatalf("location prompt leaked %q: %s", forbidden, input.Content)
		}
	}
	if strings.Contains(string(client.requests[0].Inputs["transcript"].Content), "other-session") {
		t.Fatal("transcript input contains registry evidence")
	}
	metadata, err := json.Marshal(newExtractor(t, &fakeOccurrencesLLMClient{}, references).ManifestMetadata())
	if err != nil || strings.Contains(string(metadata), "other-session") || strings.Contains(string(metadata), first.ID) {
		t.Fatalf("manifest metadata = %s, %v", metadata, err)
	}
}

func TestExtractPreservesUnknownOrMismatchedGroundingForValidators(t *testing.T) {
	locations := locationRegistry(t, "The Mill")
	known := locations.Locations[0]
	client := &fakeOccurrencesLLMClient{response: extractionResponse{Occurrences: []occurrenceResponse{
		{LocationID: "location:sha256:unknown", Name: "The Mill", Kind: "mentioned", SourceRefs: occurrenceRefs(10, 10)},
		{LocationID: known.ID, Name: "A Different Mill", Kind: "mentioned", SourceRefs: occurrenceRefs(20, 20)},
	}}}
	references := registryReferences(t, locations)
	req := extractionRequest()
	req.References = references
	result, err := newExtractor(t, client, references).Extract(context.Background(), req)
	if err != nil {
		t.Fatal(err)
	}
	if result.Value.Occurrences[0].LocationID != "location:sha256:unknown" || result.Value.Occurrences[1].Name != "A Different Mill" {
		t.Fatalf("extractor repaired validator-owned grounding errors: %#v", result.Value.Occurrences)
	}
}

func TestExtractRequiresRegistryAndAcceptsEmptyRegistryWithNoOccurrences(t *testing.T) {
	client := &fakeOccurrencesLLMClient{response: extractionResponse{Occurrences: []occurrenceResponse{}}}
	if _, err := newExtractor(t, client).Extract(context.Background(), extractionRequest()); err == nil || !strings.Contains(err.Error(), "location registry reference is required") {
		t.Fatalf("Extract() error = %v", err)
	}
	if len(client.requests) != 0 {
		t.Fatalf("LLM calls = %d", len(client.requests))
	}
	empty := dnd.LocationList{Locations: []dnd.Location{}}
	references := registryReferences(t, empty)
	req := extractionRequest()
	req.References = references
	result, err := newExtractor(t, client, references).Extract(context.Background(), req)
	if err != nil || result.Value.Occurrences == nil || len(result.Value.Occurrences) != 0 {
		t.Fatalf("empty registry result = %#v, %v", result, err)
	}
}

func TestExtractResolvesGeneratedRegistryAtOperationTimeAndDoesNotMutateResponse(t *testing.T) {
	locations := locationRegistry(t, "The Mill")
	location := locations.Locations[0]
	client := &fakeOccurrencesLLMClient{response: extractionResponse{Occurrences: []occurrenceResponse{{
		LocationID: location.ID, Name: location.Name, Kind: "mentioned", SourceRefs: occurrenceRefs(30, 30),
	}}}}
	references := registryReferences(t, locations)
	req := extractionRequest()
	req.References = references
	extractor := newExtractor(t, client)
	result, err := extractor.Extract(context.Background(), req)
	if err != nil || result.Value.Occurrences[0].Name != "The Mill" {
		t.Fatalf("Extract() = %#v, %v", result, err)
	}
	if input := client.requests[0].Inputs[LocationRegistryReferenceSlot]; !strings.Contains(string(input.Content), location.ID) || input.OriginURI != "" {
		t.Fatalf("generated registry prompt input = %#v", input)
	}
	if _, ok := extractor.ManifestMetadata()["location_registry_digest"]; ok {
		t.Fatalf("operation registry leaked into static metadata: %#v", extractor.ManifestMetadata())
	}
	if client.response.Occurrences[0].SourceRefs[0].StartUnitID != 30 {
		t.Fatalf("model response mutated: %#v", client.response)
	}
}

func TestExtractorContractsMetadataAndFailures(t *testing.T) {
	if _, err := New(nil, Options{}); err == nil || !strings.Contains(err.Error(), "LLM client") {
		t.Fatalf("New(nil) error = %v", err)
	}
	if _, err := New(&fakeOccurrencesLLMClient{}, Options{}, contracts.ReferenceSet{}, contracts.ReferenceSet{}); err == nil || !strings.Contains(err.Error(), "at most one reference set") {
		t.Fatalf("New() error = %v", err)
	}
	malformed := contracts.ReferenceSet{Slots: map[string]contracts.ResolvedReferenceSlot{LocationRegistryReferenceSlot: {
		Items: []contracts.ReferenceItem{{SlotName: LocationRegistryReferenceSlot, MediaType: "application/json", Content: []byte(`{"secret":"registry evidence"}`)}},
	}}}
	if _, err := New(&fakeOccurrencesLLMClient{}, Options{}, malformed); err == nil || !strings.Contains(err.Error(), "prepare location registry") || strings.Contains(err.Error(), "registry evidence") {
		t.Fatalf("New() error = %v", err)
	}

	locations := locationRegistry(t, "The Mill")
	references := registryReferences(t, locations)
	req := extractionRequest()
	req.References = references
	extractor := newExtractor(t, &fakeOccurrencesLLMClient{}, references)
	var nilExtractor *Extractor
	for _, test := range []struct {
		name      string
		extractor *Extractor
		req       contracts.TypedExtractionRequest
		want      string
	}{
		{"nil extractor", nilExtractor, req, "extractor"},
		{"nil client", &Extractor{}, req, "LLM client"},
		{"invalid request", extractor, mismatchedSourceInputRequest(req), "must match chunk"},
	} {
		t.Run(test.name, func(t *testing.T) {
			if _, err := test.extractor.Extract(context.Background(), test.req); err == nil || !strings.Contains(err.Error(), test.want) {
				t.Fatalf("Extract() error = %v", err)
			}
		})
	}
	if _, err := newExtractor(t, &fakeOccurrencesLLMClient{err: errors.New("provider unavailable")}, references).Extract(context.Background(), req); err == nil || !strings.Contains(err.Error(), "provider unavailable") {
		t.Fatalf("provider error = %v", err)
	}

	spec := ModuleSpec()
	if spec.Key != Key || spec.Stage != pipeline.StageExtract || spec.ExecutionClass != contracts.ExecutionClassLLMBacked || spec.ArtifactKind != dnd.LocationOccurrenceListKind {
		t.Fatalf("ModuleSpec() = %#v", spec)
	}
	var slot contracts.ReferenceSlot
	for _, candidate := range spec.ReferenceSlots {
		if candidate.Name == LocationRegistryReferenceSlot {
			slot = candidate
		}
	}
	if !slot.Required || !reflect.DeepEqual(slot.AcceptedArtifactKinds, []contracts.ArtifactKind{dnd.LocationListKind}) || slot.MaxBytes != LocationRegistryMaxBytes {
		t.Fatalf("location registry slot = %#v", slot)
	}
	registry := pipeline.NewExtractorRegistry()
	if err := Register(registry); err != nil {
		t.Fatal(err)
	}
	if _, ok := registry.Spec(Key); !ok {
		t.Fatalf("registration missing %q", Key)
	}
	if _, err := DecodeOptions(map[string]any{"unexpected": true}); err == nil {
		t.Fatal("DecodeOptions() accepted unknown options")
	}
	metadata := newExtractor(t, &fakeOccurrencesLLMClient{}, references).ManifestMetadata()
	for _, key := range []string{"prompt_sha256", "response_schema_sha256", "location_registry_digest"} {
		if value, ok := metadata[key].(string); !ok || !strings.HasPrefix(value, "sha256:") {
			t.Fatalf("metadata[%q] = %#v", key, metadata[key])
		}
	}
	if got := newExtractor(t, &fakeOccurrencesLLMClient{}, references).CheckpointFingerprints(); len(got) != 4 || got[3].Name != "location_registry" {
		t.Fatalf("fingerprints = %#v", got)
	}
}

func locationRegistry(t *testing.T, names ...string) dnd.LocationList {
	t.Helper()
	locations := make([]dnd.Location, len(names))
	for index, name := range names {
		refs := []source.SourceRef{{SourceID: "other-session", StartUnitID: index + 1, EndUnitID: index + 1}}
		locations[index] = dnd.Location{ID: identity.DeriveID(name, refs), Name: name, SourceRefs: refs}
	}
	return dnd.LocationList{Locations: locations}
}

func registryReferences(t *testing.T, locations dnd.LocationList) contracts.ReferenceSet {
	t.Helper()
	content, err := locationcodec.New().Encode(locations)
	if err != nil {
		t.Fatal(err)
	}
	return contracts.ReferenceSet{Slots: map[string]contracts.ResolvedReferenceSlot{LocationRegistryReferenceSlot: {
		Slot:  contracts.ReferenceSlot{Name: LocationRegistryReferenceSlot},
		Items: []contracts.ReferenceItem{{SlotName: LocationRegistryReferenceSlot, MediaType: locationcodec.MediaType, Content: content, Origin: contracts.ReferenceOrigin{Type: "generated"}}},
	}}}
}