notarius/internal/modules/dnd/chunk/scenes/chunker_test.go

package scenes

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/shared"
)

func TestNewModuleSpecAndRegister(t *testing.T) {
	client := &fakeScenesLLMClient{}
	chunker := newChunker(t, client)
	if chunker.Key() != Key {
		t.Fatalf("Key() = %q, want %q", chunker.Key(), Key)
	}

	want := pipeline.ModuleSpec{
		Key:                Key,
		Stage:              pipeline.StageChunk,
		ExecutionClass:     contracts.ExecutionClassLLMBacked,
		CorrectionProtocol: contracts.CorrectionProtocolSingleResponseV1,
		Requires:           []string{"source.transcript"},
		Provides:           []string{"chunks"},
		ReferenceSlots:     wantReferenceSlots(),
	}
	if got := ModuleSpec(); !reflect.DeepEqual(got, want) {
		t.Fatalf("ModuleSpec() = %#v, want %#v", got, want)
	}
	got := ModuleSpec()
	got.Requires[0] = "changed"
	got.Provides[0] = "changed"
	got.ReferenceSlots[0].AcceptedMediaTypes[0] = "changed"
	if again := ModuleSpec(); !reflect.DeepEqual(again, want) {
		t.Fatalf("ModuleSpec() after caller mutation = %#v, want %#v", again, want)
	}

	registry := pipeline.NewChunkerRegistry()
	if err := Register(registry); err != nil {
		t.Fatalf("Register() error = %v, want nil", err)
	}
	registered, ok := registry.Spec(Key)
	if !ok {
		t.Fatalf("Spec(%q) ok = false, want true", Key)
	}
	if !reflect.DeepEqual(registered, want) {
		t.Fatalf("registered spec = %#v, want %#v", registered, want)
	}
	built, err := registry.BuildWithRequest(Key, pipeline.BuildRequest{Dependencies: pipeline.ModuleDependencies{LLM: client}})
	if err != nil {
		t.Fatalf("Build(%q) error = %v, want nil", Key, err)
	}
	if built.Key() != Key {
		t.Fatalf("built Key() = %q, want %q", built.Key(), Key)
	}
	if slots := built.ReferenceSlots(); !reflect.DeepEqual(slots, want.ReferenceSlots) {
		t.Fatalf("ReferenceSlots() = %#v, want %#v", slots, want.ReferenceSlots)
	}
	if err := registry.ValidateOptions(Key, map[string]any{"unexpected": true}); err == nil || !strings.Contains(err.Error(), "unknown option") {
		t.Fatalf("ValidateOptions() error = %v, want unknown option error", err)
	}
}

func TestConstructionRejectsMissingDependencyAndUnknownOptions(t *testing.T) {
	if _, err := New(nil, Options{}); err == nil || !strings.Contains(err.Error(), "LLM client") {
		t.Fatalf("New(nil) error = %v, want LLM client error", err)
	}
	if _, err := DecodeOptions(map[string]any{"unexpected": true}); err == nil || !strings.Contains(err.Error(), "unknown option") {
		t.Fatalf("DecodeOptions() error = %v, want unknown option error", err)
	}
}

func TestRegisterNilRegistryReturnsError(t *testing.T) {
	err := Register(nil)
	if err == nil {
		t.Fatal("Register(nil) error = nil, want error")
	}
	if !strings.Contains(err.Error(), "chunker registry") {
		t.Fatalf("Register(nil) error = %q, want registry context", err.Error())
	}
}

func wantReferenceSlots() []contracts.ReferenceSlot {
	accepted := []string{"application/json", "application/x-yaml", "application/yaml", "text/markdown", "text/plain"}
	return []contracts.ReferenceSlot{
		{
			Name:               "glossary",
			Description:        "Optional campaign glossary reference material used only for scene disambiguation.",
			AcceptedMediaTypes: append([]string(nil), accepted...),
		},
		{
			Name:               "party",
			Description:        "Optional party roster reference material used only for scene disambiguation.",
			AcceptedMediaTypes: append([]string(nil), accepted...),
		},
		{
			Name:               "players",
			Description:        "Optional player list reference material used only for scene disambiguation.",
			AcceptedMediaTypes: append([]string(nil), accepted...),
		},
		{
			Name:               "roster",
			Description:        "Deprecated alias for party roster reference material used only for scene disambiguation.",
			AcceptedMediaTypes: append([]string(nil), accepted...),
		},
	}
}

func TestPlanReturnsAnnotationFreeSceneRangesFromStructuredOutput(t *testing.T) {
	client := &fakeScenesLLMClient{
		response: chunkResponse{
			Scenes: []sceneResponse{
				scene(1, 2),
				scene(3, 4),
			},
		},
	}

	result, err := newChunker(t, client).Plan(context.Background(), chunkRequest())
	if err != nil {
		t.Fatalf("Plan() error = %v, want nil", err)
	}

	if len(client.requests) != 1 {
		t.Fatalf("LLM calls = %d, want 1", len(client.requests))
	}
	req := client.requests[0]
	if req.StageName != Key {
		t.Fatalf("StageName = %q, want %q", req.StageName, Key)
	}
	if req.PromptID != PromptID || req.PromptVersion != ResponseSchemaVersion {
		t.Fatalf("prompt = %q/%q, want %q/%q", req.PromptID, req.PromptVersion, PromptID, ResponseSchemaVersion)
	}
	if req.SessionID != "session-123" || req.ProfileID != "profile-scenes" {
		t.Fatalf("session/profile = %q/%q, want session-123/profile-scenes", req.SessionID, req.ProfileID)
	}
	transcript, ok := req.Inputs["transcript"]
	if !ok {
		t.Fatalf("transcript input missing from %#v", req.Inputs)
	}
	if transcript.Name != "transcript" || transcript.MediaType != "application/json" || transcript.Digest != "sha256:transcript" || transcript.OriginURI != "file:///session-alpha.json" {
		t.Fatalf("transcript metadata = %#v", transcript)
	}
	if got := string(transcript.Content); got != sceneTranscriptJSON {
		t.Fatalf("transcript content = %q, want original source input", got)
	}
	if got := string(req.Inputs["players"].Content); got != " " {
		t.Fatalf("players input = %q, want empty reference placeholder", got)
	}
	if got := string(req.Inputs["party"].Content); got != " " {
		t.Fatalf("party input = %q, want empty reference placeholder", got)
	}
	if got := string(req.Inputs["glossary"].Content); got != " " {
		t.Fatalf("glossary input = %q, want empty reference placeholder", got)
	}

	if result.Plan.SourceDigest != "sha256:source" {
		t.Fatalf("SourceDigest = %q, want source digest", result.Plan.SourceDigest)
	}
	wantRanges := []source.ChunkRange{{StartUnitID: 1, EndUnitID: 2}, {StartUnitID: 3, EndUnitID: 4}}
	if len(result.Plan.Ranges) != len(wantRanges) {
		t.Fatalf("ranges = %#v, want two", result.Plan.Ranges)
	}
	for i, want := range wantRanges {
		got := result.Plan.Ranges[i]
		if got.StartUnitID != want.StartUnitID || got.EndUnitID != want.EndUnitID {
			t.Fatalf("range[%d] = %#v, want %#v", i, result.Plan.Ranges[i], want)
		}
		if len(got.Annotations) != 0 {
			t.Fatalf("range[%d] annotations = %#v, want absent", i, got.Annotations)
		}
	}
	if len(result.Plan.Annotations) != 0 {
		t.Fatalf("plan annotations = %#v, want absent", result.Plan.Annotations)
	}
	wantCandidate, err := json.Marshal(client.response)
	if err != nil {
		t.Fatalf("marshal expected candidate: %v", err)
	}
	if result.ModelCandidate == nil || result.ModelCandidate.Protocol != contracts.CorrectionProtocolSingleResponseV1 || !reflect.DeepEqual(result.ModelCandidate.Response, wantCandidate) {
		t.Fatalf("model candidate = %#v, want exact validated response", result.ModelCandidate)
	}
}

func TestPlanUsesDocumentOrderForNonconsecutiveUnitIDs(t *testing.T) {
	doc := &source.SourceDocument{
		ID:     "session-nonnumeric",
		Kind:   "transcript",
		Format: "application/vnd.seriatim.minimal+json",
		Digest: "sha256:nonnumeric",
		Units: []source.SourceUnit{
			{ID: 10, Kind: "transcript_segment", Text: "The party arrives.", Ref: source.SourceRef{SourceID: "session-nonnumeric", StartUnitID: 10, EndUnitID: 10}},
			{ID: 3, Kind: "transcript_segment", Text: "The party explores.", Ref: source.SourceRef{SourceID: "session-nonnumeric", StartUnitID: 3, EndUnitID: 3}},
			{ID: 20, Kind: "transcript_segment", Text: "The party rests.", Ref: source.SourceRef{SourceID: "session-nonnumeric", StartUnitID: 20, EndUnitID: 20}},
		},
	}
	req := chunkRequest()
	req.Source = doc

	t.Run("accepts document-ordered ranges", func(t *testing.T) {
		client := &fakeScenesLLMClient{response: chunkResponse{Scenes: []sceneResponse{
			scene(10, 3),
			scene(20, 20),
		}}}
		result, err := newChunker(t, client).Plan(context.Background(), req)
		if err != nil {
			t.Fatalf("Plan() error = %v, want nil", err)
		}
		if got := result.Plan.Ranges; !reflect.DeepEqual(got, []source.ChunkRange{{StartUnitID: 10, EndUnitID: 3}, {StartUnitID: 20, EndUnitID: 20}}) {
			t.Fatalf("ranges = %#v, want source-document order", got)
		}
	})

	t.Run("rejects reversed document positions", func(t *testing.T) {
		client := &fakeScenesLLMClient{response: chunkResponse{Scenes: []sceneResponse{
			scene(3, 10),
			scene(20, 20),
		}}}
		_, err := newChunker(t, client).Plan(context.Background(), req)
		if err == nil || !strings.Contains(err.Error(), "appears after") {
			t.Fatalf("Plan() error = %v, want document-position reversal", err)
		}
	})
}

func TestPlanPassesReferencesAsPromptInputs(t *testing.T) {
	client := &fakeScenesLLMClient{response: chunkResponse{
		Scenes: []sceneResponse{
			scene(1, 4),
		},
	}}
	req := chunkRequest()
	req.References = contracts.ReferenceSet{
		Slots: map[string]contracts.ResolvedReferenceSlot{
			"players": {
				Slot: contracts.ReferenceSlot{Name: "players"},
				Items: []contracts.ReferenceItem{
					{SlotName: "players", Content: []byte("Alice: Aria")},
				},
			},
			"party": {
				Slot: contracts.ReferenceSlot{Name: "party"},
				Items: []contracts.ReferenceItem{
					{SlotName: "party", Content: []byte("Aria: cleric")},
				},
			},
			"glossary": {
				Slot: contracts.ReferenceSlot{Name: "glossary"},
				Items: []contracts.ReferenceItem{
					{SlotName: "glossary", Content: []byte("Brightmantle: local temple")},
				},
			},
		},
	}

	if _, err := newChunker(t, client).Plan(context.Background(), req); err != nil {
		t.Fatalf("Plan() error = %v, want nil", err)
	}
	request := client.requests[0]
	if got := string(request.Inputs["players"].Content); got != "Alice: Aria" {
		t.Fatalf("players input = %q, want reference content", got)
	}
	if got := string(request.Inputs["party"].Content); got != "Aria: cleric" {
		t.Fatalf("party input = %q, want reference content", got)
	}
	if got := string(request.Inputs["glossary"].Content); got != "Brightmantle: local temple" {
		t.Fatalf("glossary input = %q, want reference content", got)
	}
	if strings.Contains(string(request.Inputs["transcript"].Content), "Aria: cleric") {
		t.Fatalf("transcript input contains reference content")
	}
}

func TestPromptInputsMapLegacyRosterReferenceToParty(t *testing.T) {
	inputs := shared.PromptInputs(sceneSourceInput(), contracts.ReferenceSet{
		Slots: map[string]contracts.ResolvedReferenceSlot{
			"roster": {
				Slot: contracts.ReferenceSlot{Name: "roster"},
				Items: []contracts.ReferenceItem{
					{SlotName: "roster", Content: []byte("Legacy roster text")},
				},
			},
		},
	})

	if got := string(inputs["party"].Content); got != "Legacy roster text" {
		t.Fatalf("party input = %q, want legacy roster content", got)
	}
	if _, ok := inputs["roster"]; ok {
		t.Fatalf("roster prompt input was present; want only party input")
	}
}

func TestChunkerManifestMetadataIncludesPromptAndSchemaProvenance(t *testing.T) {
	metadata := newChunker(t, &fakeScenesLLMClient{}).ManifestMetadata()

	tests := map[string]string{
		"prompt_id":               PromptID,
		"prompt_version":          ResponseSchemaVersion,
		"response_schema_key":     string(ResponseSchemaKey),
		"response_schema_id":      ResponseSchemaID,
		"response_schema_name":    ResponseSchemaName,
		"response_schema_version": ResponseSchemaVersion,
	}
	for key, want := range tests {
		if metadata[key] != want {
			t.Fatalf("metadata[%q] = %#v, want %q", key, metadata[key], want)
		}
	}
	for _, key := range []string{"prompt_sha256", "response_schema_sha256"} {
		value, ok := metadata[key].(string)
		if !ok || !strings.HasPrefix(value, "sha256:") {
			t.Fatalf("metadata[%q] = %#v, want sha256 value", key, metadata[key])
		}
	}
	for _, forbidden := range []string{"prompt", "schema", "source", "text"} {
		if _, ok := metadata[forbidden]; ok {
			t.Fatalf("metadata includes raw %q field: %#v", forbidden, metadata)
		}
	}
}

func TestPlanRejectsInvalidRequests(t *testing.T) {
	validClient := &fakeScenesLLMClient{response: validSceneResponse()}
	validReq := chunkRequest()
	canceledCtx, cancel := context.WithCancel(context.Background())
	cancel()
	invalidDoc := sceneSourceDocument()
	invalidDoc.Units[0].ID = 0
	emptyDoc := sceneSourceDocument()
	emptyDoc.Units = nil

	tests := []struct {
		name    string
		chunker *Chunker
		ctx     context.Context
		req     contracts.ChunkRequest
		want    string
	}{
		{name: "nil chunker", chunker: nil, ctx: context.Background(), req: validReq, want: "chunker"},
		{name: "nil context", chunker: newChunker(t, validClient), ctx: nil, req: validReq, want: "context"},
		{name: "canceled context", chunker: newChunker(t, validClient), ctx: canceledCtx, req: validReq, want: "context"},
		{name: "nil source", chunker: newChunker(t, validClient), ctx: context.Background(), req: contracts.ChunkRequest{}, want: "source"},
		{name: "empty source units", chunker: newChunker(t, validClient), ctx: context.Background(), req: contracts.ChunkRequest{Source: emptyDoc}, want: "units"},
		{name: "invalid source", chunker: newChunker(t, validClient), ctx: context.Background(), req: contracts.ChunkRequest{Source: invalidDoc}, want: "validate source document"},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			_, err := tt.chunker.Plan(tt.ctx, tt.req)
			if err == nil {
				t.Fatal("Plan() error = nil, want error")
			}
			if !strings.Contains(err.Error(), "dnd scenes") || !strings.Contains(err.Error(), tt.want) {
				t.Fatalf("Plan() error = %q, want module context and %q", err.Error(), tt.want)
			}
		})
	}
}

func TestPlanRejectsMalformedStructuredOutput(t *testing.T) {
	tests := []struct {
		name     string
		response chunkResponse
		want     string
	}{
		{name: "missing scenes", response: chunkResponse{}, want: "scenes"},
		{name: "empty scenes", response: chunkResponse{Scenes: []sceneResponse{}}, want: "scenes"},
		{
			name: "unknown boundary id",
			response: replaceScenes(validSceneResponse(), []sceneResponse{
				scene(1, 999),
			}),
			want: "was not found",
		},
		{
			name: "out of order boundaries",
			response: replaceScenes(validSceneResponse(), []sceneResponse{
				scene(3, 2),
			}),
			want: "appears after",
		},
		{
			name: "gap",
			response: replaceScenes(validSceneResponse(), []sceneResponse{
				scene(1, 1),
				scene(3, 4),
			}),
			want: "gap",
		},
		{
			name: "overlap",
			response: replaceScenes(validSceneResponse(), []sceneResponse{
				scene(1, 2),
				scene(2, 4),
			}),
			want: "overlap",
		},
		{
			name: "incomplete coverage",
			response: replaceScenes(validSceneResponse(), []sceneResponse{
				scene(1, 3),
			}),
			want: "final scene",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			client := &fakeScenesLLMClient{response: tt.response}
			_, err := newChunker(t, client).Plan(context.Background(), chunkRequest())
			if err == nil {
				t.Fatal("Plan() error = nil, want error")
			}
			if !strings.Contains(err.Error(), "dnd scenes") || !strings.Contains(err.Error(), tt.want) {
				t.Fatalf("Plan() error = %q, want module context and %q", err.Error(), tt.want)
			}
		})
	}
}

func TestPlanWrapsLLMClientError(t *testing.T) {
	client := &fakeScenesLLMClient{err: errors.New("provider unavailable")}

	_, err := newChunker(t, client).Plan(context.Background(), chunkRequest())
	if err == nil {
		t.Fatal("Plan() error = nil, want LLM error")
	}
	if !strings.Contains(err.Error(), "dnd scenes") || !strings.Contains(err.Error(), "provider unavailable") {
		t.Fatalf("Plan() error = %q, want wrapped LLM context", err.Error())
	}
}

func chunkRequest() contracts.ChunkRequest {
	return contracts.ChunkRequest{
		Source:      sceneSourceDocument(),
		SourceInput: sceneSourceInput(),
		SessionID:   "session-123",
		LLMProfile:  "profile-scenes",
	}
}

const sceneTranscriptJSON = `{"id":"session-alpha","segments":[{"id":1,"text":"Aria asks whether the goblin will parley."}]}`

func sceneSourceInput() contracts.LLMInputMaterial {
	return contracts.NewLLMInputMaterial("source", "application/json", []byte(sceneTranscriptJSON), "sha256:transcript", "file:///session-alpha.json")
}

func newChunker(t *testing.T, client contracts.StructuredLLMClient) *Chunker {
	t.Helper()
	chunker, err := New(client, Options{})
	if err != nil {
		t.Fatalf("New() error = %v, want nil", err)
	}
	return chunker
}

func sceneSourceDocument() *source.SourceDocument {
	return &source.SourceDocument{
		ID:     "session-alpha",
		Kind:   "transcript",
		Format: "application/vnd.seriatim.minimal+json",
		Digest: "sha256:source",
		Units: []source.SourceUnit{
			{ID: 1, Kind: "transcript_segment", Text: "Aria asks whether the goblin will parley.", Ref: source.SourceRef{SourceID: "session-alpha", StartUnitID: 1, EndUnitID: 1}, Metadata: map[string]any{"speaker": "Alice"}},
			{ID: 2, Kind: "transcript_segment", Text: "The goblin scout describes the gate guards.", Ref: source.SourceRef{SourceID: "session-alpha", StartUnitID: 2, EndUnitID: 2}},
			{ID: 3, Kind: "transcript_segment", Text: "The guards rush out with blades drawn.", Ref: source.SourceRef{SourceID: "session-alpha", StartUnitID: 3, EndUnitID: 3}},
			{ID: 4, Kind: "transcript_segment", Text: "The party defeats the ambushers.", Ref: source.SourceRef{SourceID: "session-alpha", StartUnitID: 4, EndUnitID: 4}},
		},
	}
}

func validSceneResponse() chunkResponse {
	return chunkResponse{
		Scenes: []sceneResponse{
			scene(1, 4),
		},
	}
}

func replaceScenes(response chunkResponse, scenes []sceneResponse) chunkResponse {
	response.Scenes = scenes
	return response
}

func scene(startUnitID int, endUnitID int) sceneResponse {
	return sceneResponse{
		StartUnitID: shared.UnitRefFromInt(startUnitID),
		EndUnitID:   shared.UnitRefFromInt(endUnitID),
	}
}

func chunkIDs(chunks []source.Chunk) []string {
	ids := make([]string, 0, len(chunks))
	for _, chunk := range chunks {
		ids = append(ids, chunk.ID)
	}
	return ids
}

func unitIDs(units []source.SourceUnit) []int {
	ids := make([]int, 0, len(units))
	for _, unit := range units {
		ids = append(ids, unit.ID)
	}
	return ids
}

type fakeScenesLLMClient struct {
	response chunkResponse
	err      error
	requests []contracts.StructuredCompletionRequest
}

func (client *fakeScenesLLMClient) CompleteStructured(ctx context.Context, req contracts.StructuredCompletionRequest, out any) (contracts.StructuredCompletionResponse, error) {
	client.requests = append(client.requests, cloneStructuredCompletionRequest(req))
	if client.err != nil {
		return contracts.StructuredCompletionResponse{}, client.err
	}

	target, ok := out.(*chunkResponse)
	if !ok {
		return contracts.StructuredCompletionResponse{}, errors.New("unexpected output target")
	}
	*target = client.response
	content, err := json.Marshal(client.response)
	if err != nil {
		return contracts.StructuredCompletionResponse{}, err
	}
	return contracts.StructuredCompletionResponse{Content: content}, nil
}

func cloneStructuredCompletionRequest(req contracts.StructuredCompletionRequest) contracts.StructuredCompletionRequest {
	req.Inputs = req.Inputs.Clone()
	correction, err := contracts.CloneSemanticCorrection(req.Correction)
	if err != nil {
		panic(err)
	}
	req.Correction = correction
	req.Vars = cloneVars(req.Vars)
	return req
}

func cloneVars(in map[string]any) map[string]any {
	if len(in) == 0 {
		return nil
	}
	out := make(map[string]any, len(in))
	for key, value := range in {
		out[key] = value
	}
	return out
}