notarius/internal/cli/dnd_enemy_events_contract_test.go

package cli

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"reflect"
	"strings"
	"sync"
	"testing"
	"time"

	"gitea.maximumdirect.net/eric/notarius/internal/core/artifacts"
	"gitea.maximumdirect.net/eric/notarius/internal/core/config"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/evidencecontext"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/chunk/scenes"
	locationoccurrencecodec "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/codec/locationoccurrences"
	locationcodec "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/codec/locations"
	combat "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/extract/combatturns"
	enemyevents "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/extract/enemyevents"
	itemevents "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/extract/itemevents"
	locationoccurrences "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/extract/locationoccurrences"
	locations "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/extract/locations"
	npcinteractions "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/extract/npcinteractions"
	npcs "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/extract/npcs"
	scenedescriptions "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/extract/scenedescriptions"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/extract/spells"
	enemyeventnormalize "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/normalize/enemyevents"
	locationnormalize "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/normalize/locations"
	npcnormalize "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/normalize/npcs"
)

func TestProductionEnemyEventConfigurationResolvesGeneratedHandoffs(t *testing.T) {
	components := productionTestComponents(t)
	cfg := loadMaintainedExample(t, repositoryPath("examples", "dnd-complete.config.yml"))
	effective, err := cfg.Resolve(resolveInputForMaintainedExample(components, "dnd-session"))
	if err != nil {
		t.Fatalf("Resolve() error = %v, want nil", err)
	}
	materialized, _, err := pipeline.MaterializeReferences(effective.ResolvedPipeline, catalogFromRegistries(components.registries), pipeline.ReferenceMaterializationOptions{
		ConfigPath: repositoryPath("examples", "dnd-complete.config.yml"),
		WorkingDir: repositoryPath("examples"),
	})
	if err != nil {
		t.Fatalf("MaterializeReferences() error = %v, want nil", err)
	}
	lane := referenceContractLane(t, materialized, "enemy-events")
	if lane.ArtifactKind != dnd.EnemyEventListKind || lane.Extract.Module != enemyevents.Key || lane.Extract.Retries != 2 || lane.Merge.Module != pipeline.DefaultMergeModule || lane.Normalize.Module != enemyeventnormalize.Key {
		t.Fatalf("enemy event lane = %#v, want typed production composition", lane)
	}
	for slot, want := range map[string]struct{ step, lane string }{
		"npcs":               {step: "describe-session", lane: "npcs"},
		"scene_descriptions": {step: "describe-session", lane: "scene-descriptions"},
		"combat_turns":       {step: "extract-events", lane: "combat-turns"},
		"npc_interactions":   {step: "extract-events", lane: "npc-interactions"},
	} {
		binding, found := generatedReferenceBinding(lane.ExtractReferences.Bindings, slot)
		if !found || binding.Artifact.Step != want.step || binding.Artifact.Lane != want.lane {
			t.Fatalf("enemy event %s reference = %#v, want generated %s/%s artifact", slot, binding, want.step, want.lane)
		}
	}
	if binding, found := generatedReferenceBinding(lane.NormalizeReferences.Bindings, "npcs"); !found || binding.Artifact.Step != "describe-session" || binding.Artifact.Lane != "npcs" {
		t.Fatalf("enemy event normalizer NPC reference = %#v, want generated NPC artifact", binding)
	}

	catalog := catalogFromRegistries(components.registries)
	extractSpec, ok := catalog.Extractors.Spec(enemyevents.Key)
	if !ok || !reflect.DeepEqual(extractSpec.Requires, []string{"chunks", "source.transcript"}) || !reflect.DeepEqual(extractSpec.Provides, []string{"dnd.enemy_events"}) {
		t.Fatalf("enemy event extractor spec = %#v, want source and artifact capabilities", extractSpec)
	}
	normalizeSpec, ok := catalog.Normalizers.SpecForArtifact(enemyeventnormalize.Key, dnd.EnemyEventListKind)
	if !ok || !reflect.DeepEqual(normalizeSpec.Requires, []string{"merged"}) || !reflect.DeepEqual(normalizeSpec.Provides, []string{"normalized"}) {
		t.Fatalf("enemy event normalizer spec = %#v, want merged/normalized capabilities", normalizeSpec)
	}
	for _, slot := range []string{"npcs", "scene_descriptions", "combat_turns", "npc_interactions"} {
		if !hasReferenceSlot(extractSpec.ReferenceSlots, slot) {
			t.Fatalf("enemy event extractor slots = %#v, want %q", extractSpec.ReferenceSlots, slot)
		}
	}
	if !hasReferenceSlot(normalizeSpec.ReferenceSlots, "npcs") {
		t.Fatalf("enemy event normalizer slots = %#v, want NPC registry", normalizeSpec.ReferenceSlots)
	}

	profile := cfg.Pipelines["dnd-session"]
	profile.Steps[2].References["npcs"] = pipeline.GeneratedReference("track-enemies", "enemy-events")
	cfg.Pipelines["dnd-session"] = profile
	if _, err := cfg.Resolve(resolveInputForMaintainedExample(components, "dnd-session")); err == nil || !strings.Contains(err.Error(), "earlier step") {
		t.Fatalf("Resolve() error = %v, want future generated-reference rejection", err)
	}
}

func TestMaintainedCompleteExampleProducesEnemyEventsThroughGeneratedHandoffs(t *testing.T) {
	t.Chdir(repositoryPath())
	outputRoot := filepath.Join(t.TempDir(), "output")
	configPath := completeExampleConfigWithTemporaryCache(t)
	client := &enemyEventLLMClient{}
	options := productionCLIOptions(t)
	options.Now = func() time.Time { return time.Unix(1700000000, 0).UTC() }
	options.RunIDGenerator = func(time.Time) (string, error) { return productionRunID, nil }
	options.UserCacheDir = func() (string, error) { return "", errors.New("user cache must not be used") }
	options.LLMClientFactory = func(context.Context, config.Config, string, LLMRuntimeOverrides) (contracts.StructuredLLMClient, []artifacts.LLMProfileManifest, error) {
		return client, nil, nil
	}

	var stdout, stderr strings.Builder
	code := RunWithOptions([]string{
		"run", "dnd-session",
		"--config", configPath,
		"--input", repositoryPath("examples", "dnd-complete-transcript.json"),
		"--chunk_cache", "bypass", "--output-dir", outputRoot, "--session-id", "enemy-event-session",
	}, &stdout, &stderr, options)
	if code != 0 {
		t.Fatalf("code=%d stdout=%q stderr=%q", code, stdout.String(), stderr.String())
	}

	runRoot := filepath.Join(outputRoot, productionRunID)
	index := readProductionJSON[exampleOutputIndex](t, filepath.Join(runRoot, "index.json"))
	var enemyOutput exampleOutputIndexEntry
	var locationOutput, occurrenceOutput exampleOutputIndexEntry
	for _, entry := range index.OutputFiles {
		switch entry.LaneID {
		case "enemy-events":
			enemyOutput = entry
		case "locations":
			locationOutput = entry
		case "location-occurrences":
			occurrenceOutput = entry
		}
	}
	if enemyOutput.File != "lanes/enemy-events.json" || enemyOutput.SchemaID != "notarius.dnd.enemy_events" || enemyOutput.SchemaVersion != "v1" {
		t.Fatalf("enemy event output = %#v, want typed enemy-event JSON", enemyOutput)
	}
	value := readProductionJSON[dnd.EnemyEventList](t, filepath.Join(runRoot, enemyOutput.File))
	if len(value.Events) != 1 || value.Events[0].Name != "Kesh" || value.Events[0].Kind != dnd.EnemyEventKindFled || len(value.Events[0].SourceRefs) != 1 || value.Events[0].SourceRefs[0].SourceID != "session-ravenfall" || value.Events[0].SourceRefs[0].StartUnitID != 10 {
		t.Fatalf("enemy event artifact = %#v, want source-linked Kesh fleeing event", value)
	}
	if locationOutput.File != "lanes/locations.json" || locationOutput.SchemaID != locationcodec.SchemaID || locationOutput.SchemaVersion != locationcodec.SchemaVersion {
		t.Fatalf("location output = %#v, want typed location registry JSON", locationOutput)
	}
	locationsValue := readProductionJSON[dnd.LocationList](t, filepath.Join(runRoot, locationOutput.File))
	if len(locationsValue.Locations) != 2 || locationsValue.Locations[0].Name != "Moon Gate" || locationsValue.Locations[1].Name != "Moon Gate" || locationsValue.Locations[0].ID == locationsValue.Locations[1].ID {
		t.Fatalf("location registry = %#v, want distinct source-grounded identities for same-name locations", locationsValue)
	}
	if occurrenceOutput.File != "lanes/location-occurrences.json" || occurrenceOutput.SchemaID != locationoccurrencecodec.SchemaID || occurrenceOutput.SchemaVersion != locationoccurrencecodec.SchemaVersion {
		t.Fatalf("location occurrence output = %#v, want typed occurrence JSON", occurrenceOutput)
	}
	occurrencesValue := readProductionJSON[dnd.LocationOccurrenceList](t, filepath.Join(runRoot, occurrenceOutput.File))
	if len(occurrencesValue.Occurrences) != 2 || occurrencesValue.Occurrences[0].LocationID == occurrencesValue.Occurrences[1].LocationID || occurrencesValue.Occurrences[0].Name != "Moon Gate" || occurrencesValue.Occurrences[1].Name != "Moon Gate" {
		t.Fatalf("location occurrences = %#v, want source-grounded references to distinct registry identities", occurrencesValue)
	}

	evidence := readProductionJSON[evidencecontext.Document](t, filepath.Join(runRoot, "evidence-context.json"))
	for _, laneID := range []string{"enemy-events", "locations", "location-occurrences"} {
		if !containsString(evidence.SelectedLanes, laneID) || !evidenceHasLane(evidence, laneID) {
			t.Fatalf("evidence context = %#v, want direct %s evidence", evidence, laneID)
		}
	}

	requests := client.requestsFor(enemyevents.PromptID)
	if len(requests) != 1 {
		t.Fatalf("enemy event requests = %#v, want only the combat scene request", requests)
	}
	request := requests[0]
	if request.SessionID != "enemy-event-session" {
		t.Fatalf("enemy event session = %q, want shared session", request.SessionID)
	}
	for slot, required := range map[string]string{
		"npcs":             "Kesh",
		"combat_turns":     "Kesh",
		"npc_interactions": "Kesh",
	} {
		input, ok := request.Inputs[slot]
		if !ok || !strings.Contains(string(input.Content), required) || strings.Contains(string(input.Content), "source_refs") || strings.Contains(string(input.Content), "start_unit_id") {
			t.Fatalf("enemy event %s prompt input = %q, want compact source-free grounding", slot, input.Content)
		}
	}
	locationRequests := client.requestsFor(locationoccurrences.PromptID)
	if len(locationRequests) != 2 {
		t.Fatalf("location occurrence requests = %#v, want one request per scene", locationRequests)
	}
	for _, request := range locationRequests {
		registryInput := request.Inputs["locations"]
		if !strings.Contains(string(registryInput.Content), "Moon Gate") || !strings.Contains(string(registryInput.Content), `"id"`) || strings.Contains(string(registryInput.Content), "source_refs") {
			t.Fatalf("location occurrence registry input = %q, want source-free ID grounding", registryInput.Content)
		}
	}
}

func completeExampleConfigWithTemporaryCache(t *testing.T) string {
	t.Helper()
	content, err := os.ReadFile(repositoryPath("examples", "dnd-complete.config.yml"))
	if err != nil {
		t.Fatal(err)
	}
	cacheRoot := t.TempDir()
	updated := strings.Replace(string(content), "directory: ./notarius-cache/chunk-plans", fmt.Sprintf("directory: %q", filepath.Join(cacheRoot, "chunk-plans")), 1)
	updated = strings.Replace(updated, "directory: ./notarius-cache/checkpoints", fmt.Sprintf("directory: %q", filepath.Join(cacheRoot, "checkpoints")), 1)
	for relative, absolute := range map[string]string{
		"./dnd-party.txt":          repositoryPath("examples", "dnd-party.txt"),
		"./dnd-glossary.txt":       repositoryPath("examples", "dnd-glossary.txt"),
		"./dnd-spell-catalog.json": repositoryPath("examples", "dnd-spell-catalog.json"),
	} {
		updated = strings.ReplaceAll(updated, relative, fmt.Sprintf("%q", absolute))
	}
	path := filepath.Join(t.TempDir(), "dnd-complete.config.yml")
	if err := os.WriteFile(path, []byte(updated), 0o600); err != nil {
		t.Fatal(err)
	}
	return path
}

type enemyEventLLMClient struct {
	mu       sync.Mutex
	requests []contracts.StructuredCompletionRequest
}

func (client *enemyEventLLMClient) CompleteStructured(ctx context.Context, request contracts.StructuredCompletionRequest, output any) (contracts.StructuredCompletionResponse, error) {
	if err := ctx.Err(); err != nil {
		return contracts.StructuredCompletionResponse{}, err
	}
	combatScene := strings.Contains(string(request.Inputs["transcript"].Content), "Roll initiative")
	var content []byte
	switch request.PromptID {
	case scenes.PromptID:
		content = []byte(`{"scenes":[{"start_unit_id":1,"end_unit_id":6},{"start_unit_id":7,"end_unit_id":11}]}`)
	case npcs.PromptID:
		if combatScene {
			content = []byte(`{"npcs":[{"name":"Kesh","source_refs":[{"start_unit_id":7,"end_unit_id":7}]}]}`)
		} else {
			content = []byte(`{"npcs":[]}`)
		}
	case npcnormalize.PromptID:
		content = []byte(`{"duplicate_groups":[]}`)
	case scenedescriptions.PromptID:
		kind, title := "narrative", "Arrival"
		if combatScene {
			kind, title = "combat", "Raiders attack"
		}
		content = []byte(fmt.Sprintf(`{"kind":%q,"title":%q,"summary":"session scene"}`, kind, title))
	case locations.PromptID:
		unitID := 1
		if combatScene {
			unitID = 7
		}
		content = []byte(fmt.Sprintf(`{"locations":[{"name":"Moon Gate","source_refs":[{"start_unit_id":%d,"end_unit_id":%d}]}]}`, unitID, unitID))
	case locationnormalize.PromptID:
		content = []byte(`{"duplicate_groups":[]}`)
	case spells.PromptID:
		content = []byte(`{"spell_casts":[]}`)
	case itemevents.PromptID:
		content = []byte(`{"events":[]}`)
	case combat.PromptID:
		content = []byte(`{"combat_turns":[{"actor":"Kesh","turn_kind":"turn","source_refs":[{"start_unit_id":8,"end_unit_id":8}]}]}`)
	case npcinteractions.PromptID:
		if combatScene {
			content = []byte(`{"interactions":[{"name":"Kesh","kind":"combat_opponent","source_refs":[{"start_unit_id":7,"end_unit_id":7}]}]}`)
		} else {
			content = []byte(`{"interactions":[]}`)
		}
	case locationoccurrences.PromptID:
		var registry struct {
			Locations []struct {
				ID string `json:"id"`
			} `json:"locations"`
		}
		if err := json.Unmarshal(request.Inputs["locations"].Content, &registry); err != nil {
			return contracts.StructuredCompletionResponse{}, fmt.Errorf("decode generated location registry: %w", err)
		}
		if len(registry.Locations) == 0 {
			return contracts.StructuredCompletionResponse{}, fmt.Errorf("generated location registry has no locations")
		}
		unitID := 1
		locationID := registry.Locations[0].ID
		if combatScene {
			unitID = 7
			if len(registry.Locations) > 1 {
				locationID = registry.Locations[1].ID
			}
		}
		content = []byte(fmt.Sprintf(`{"occurrences":[{"location_id":%q,"name":"Moon Gate","kind":"visited","source_refs":[{"start_unit_id":%d,"end_unit_id":%d}]}]}`, locationID, unitID, unitID))
	case enemyevents.PromptID:
		content = []byte(`{"events":[{"name":"Kesh","kind":"fled","source_refs":[{"start_unit_id":10,"end_unit_id":10}]}]}`)
	default:
		return contracts.StructuredCompletionResponse{}, fmt.Errorf("unexpected prompt %q", request.PromptID)
	}
	if err := json.Unmarshal(content, output); err != nil {
		return contracts.StructuredCompletionResponse{}, fmt.Errorf("populate fake structured target: %w", err)
	}
	client.mu.Lock()
	client.requests = append(client.requests, request)
	client.mu.Unlock()
	return contracts.StructuredCompletionResponse{Content: content, Provider: "test", Model: "deterministic", ProfileID: request.ProfileID}, nil
}

func (client *enemyEventLLMClient) requestsFor(promptID string) []contracts.StructuredCompletionRequest {
	client.mu.Lock()
	defer client.mu.Unlock()
	var requests []contracts.StructuredCompletionRequest
	for _, request := range client.requests {
		if request.PromptID == promptID {
			requests = append(requests, request)
		}
	}
	return requests
}

func containsString(values []string, want string) bool {
	for _, value := range values {
		if value == want {
			return true
		}
	}
	return false
}

func evidenceHasLane(value evidencecontext.Document, laneID string) bool {
	for _, context := range value.Contexts {
		for _, reference := range context.EvidenceRefs {
			if reference.LaneID == laneID {
				return true
			}
		}
	}
	return false
}

func generatedReferenceBinding(bindings []pipeline.ReferenceBinding, slotName string) (pipeline.ReferenceBinding, bool) {
	for _, binding := range bindings {
		if binding.SlotName == slotName && binding.Artifact != nil {
			return binding, true
		}
	}
	return pipeline.ReferenceBinding{}, false
}