Switch from regex parsing to dom-parsing for response term mapping

Change-Id: I02becb3cfa1f0f9563af56edabea08058ce14c56
diff --git a/mapper/benchmark_test.go b/mapper/benchmark_test.go
index 6338661..fcf28a6 100644
--- a/mapper/benchmark_test.go
+++ b/mapper/benchmark_test.go
@@ -141,6 +141,41 @@
 	}
 }
 
+// BenchmarkApplyResponseMappings benchmarks the SAX-based response annotation
+func BenchmarkApplyResponseMappings(b *testing.B) {
+	mappingList := config.MappingList{
+		ID:       "bench-response",
+		FoundryA: "marmot",
+		LayerA:   "p",
+		FoundryB: "opennlp",
+		LayerB:   "p",
+		Mappings: []config.MappingRule{
+			"[DET] <> [DT]",
+			"[NN] <> [NOUN]",
+			"[ADJA] <> [ADJ]",
+		},
+	}
+
+	mapper, err := NewMapper([]config.MappingList{mappingList})
+	if err != nil {
+		b.Fatalf("Failed to create mapper: %v", err)
+	}
+
+	responseData := map[string]any{
+		"snippet": `<span title="marmot/p:DET">Der</span> <span title="marmot/p:ADJA">alte</span> <span title="marmot/p:NN">Mann</span>`,
+	}
+
+	opts := MappingOptions{Direction: AtoB}
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		_, err := mapper.ApplyResponseMappings("bench-response", opts, responseData)
+		if err != nil {
+			b.Fatalf("ApplyResponseMappings failed: %v", err)
+		}
+	}
+}
+
 // BenchmarkApplyQueryMappingsWorstCase benchmarks the worst case scenario
 // where we have many rules but no matches (tests optimization effectiveness)
 func BenchmarkApplyQueryMappingsWorstCase(b *testing.B) {
diff --git a/mapper/response.go b/mapper/response.go
index ced97f3..14fd978 100644
--- a/mapper/response.go
+++ b/mapper/response.go
@@ -7,6 +7,7 @@
 	"github.com/KorAP/Koral-Mapper/ast"
 	"github.com/KorAP/Koral-Mapper/matcher"
 	"github.com/KorAP/Koral-Mapper/parser"
+	"github.com/orisano/gosax"
 	"github.com/rs/zerolog/log"
 )
 
@@ -195,55 +196,83 @@
 }
 
 // addAnnotationsToSnippet adds new annotations to matching tokens in the snippet
+// using SAX-based parsing for structural identification of text nodes.
 func (m *Mapper) addAnnotationsToSnippet(snippet string, matchingTokens []matcher.TokenSpan, annotationStrings []string) (string, error) {
 	if len(matchingTokens) == 0 || len(annotationStrings) == 0 {
 		return snippet, nil
 	}
 
-	result := snippet
+	tokenByStartPos := make(map[int]matcher.TokenSpan)
+	for _, tok := range matchingTokens {
+		tokenByStartPos[tok.StartPos] = tok
+	}
 
-	// Process each matching token
-	for _, token := range matchingTokens {
-		// For nested span structure, we need to find the innermost text and wrap it
-		// Look for the actual token text within span tags
-		tokenText := token.Text
+	reader := strings.NewReader(snippet)
+	r := gosax.NewReader(reader)
 
-		// Find all occurrences of the token text in the current snippet
-		// We need to be careful about which occurrence to replace
-		startPos := 0
-		for {
-			tokenStart := strings.Index(result[startPos:], tokenText)
-			if tokenStart == -1 {
-				break // No more occurrences
+	var result strings.Builder
+	result.Grow(len(snippet) + len(matchingTokens)*100)
+
+	var textPos int
+
+	for {
+		e, err := r.Event()
+		if err != nil {
+			return "", fmt.Errorf("failed to parse snippet for annotation: %w", err)
+		}
+		if e.Type() == gosax.EventEOF {
+			break
+		}
+
+		switch e.Type() {
+		case gosax.EventStart:
+			result.Write(e.Bytes)
+
+		case gosax.EventEnd:
+			result.Write(e.Bytes)
+
+		case gosax.EventText:
+			charData, err := gosax.CharData(e.Bytes)
+			if err != nil {
+				result.Write(e.Bytes)
+				break
 			}
-			tokenStart += startPos
-			tokenEnd := tokenStart + len(tokenText)
 
-			// Check if this token text is within the expected context
-			// Look backwards and forwards to see if we're in the right span context
-			beforeContext := result[:tokenStart]
-			afterContext := result[tokenEnd:]
+			text := string(charData)
+			trimmed := strings.TrimSpace(text)
 
-			// Simple heuristic: if we're immediately preceded by a > and followed by a <
-			// then we're likely at the innermost text node
-			if strings.HasSuffix(beforeContext, ">") && (strings.HasPrefix(afterContext, "<") || len(afterContext) == 0 || afterContext[0] == ' ') {
-				// Build the replacement with nested spans for each annotation
-				replacement := tokenText
+			if token, ok := tokenByStartPos[textPos]; ok && trimmed != "" && trimmed == token.Text {
+				trimStart := strings.Index(text, trimmed)
+				leadingWS := text[:trimStart]
+				trailingWS := text[trimStart+len(trimmed):]
+
+				result.WriteString(leadingWS)
+
+				annotated := escapeXMLText(trimmed)
 				for i := len(annotationStrings) - 1; i >= 0; i-- {
-					replacement = fmt.Sprintf(`<span title="%s" class="notinindex">%s</span>`, annotationStrings[i], replacement)
+					annotated = fmt.Sprintf(`<span title="%s" class="notinindex">%s</span>`, annotationStrings[i], annotated)
 				}
-
-				// Replace this occurrence
-				result = result[:tokenStart] + replacement + result[tokenEnd:]
-				break // Only replace the first appropriate occurrence for this token
+				result.WriteString(annotated)
+				result.WriteString(trailingWS)
+			} else {
+				result.Write(e.Bytes)
 			}
 
-			// Move past this occurrence
-			startPos = tokenEnd
+			textPos += len(text)
+
+		default:
+			result.Write(e.Bytes)
 		}
 	}
 
-	return result, nil
+	return result.String(), nil
+}
+
+func escapeXMLText(s string) string {
+	s = strings.ReplaceAll(s, "&", "&amp;")
+	s = strings.ReplaceAll(s, "<", "&lt;")
+	s = strings.ReplaceAll(s, ">", "&gt;")
+	return s
 }
 
 // applyReplacementWithLayerPrecedence applies RestrictToObligatory with proper layer precedence
diff --git a/mapper/response_test.go b/mapper/response_test.go
index 88273c5..9a0c668 100644
--- a/mapper/response_test.go
+++ b/mapper/response_test.go
@@ -493,6 +493,444 @@
 	assert.Equal(t, "John Doe", author)
 }
 
+// TestResponseAnnotationDuplicateTokenText tests that when the same token text
+// appears multiple times, only the correct occurrence is annotated based on its
+// annotation context (not string position).
+func TestResponseAnnotationDuplicateTokenText(t *testing.T) {
+	// "Der" appears twice: first as NN (no match), then as DET (match).
+	// The old string-heuristic would annotate the first "Der" because it
+	// finds the first occurrence preceded by ">".
+	responseSnippet := `{
+		"snippet": "<span title=\"marmot/p:NN\">Der</span> <span title=\"marmot/p:DET\">Der</span>"
+	}`
+
+	mappingList := config.MappingList{
+		ID:       "test-dup-mapper",
+		FoundryA: "marmot",
+		LayerA:   "p",
+		FoundryB: "opennlp",
+		LayerB:   "p",
+		Mappings: []config.MappingRule{
+			"[DET] <> [DT]",
+		},
+	}
+
+	m, err := NewMapper([]config.MappingList{mappingList})
+	require.NoError(t, err)
+
+	var inputData any
+	err = json.Unmarshal([]byte(responseSnippet), &inputData)
+	require.NoError(t, err)
+
+	result, err := m.ApplyResponseMappings("test-dup-mapper", MappingOptions{Direction: AtoB}, inputData)
+	require.NoError(t, err)
+
+	resultMap := result.(map[string]any)
+	snippet := resultMap["snippet"].(string)
+
+	// Only the second "Der" (DET) should be annotated
+	expected := `<span title="marmot/p:NN">Der</span> <span title="marmot/p:DET"><span title="opennlp/p:DT" class="notinindex">Der</span></span>`
+	assert.Equal(t, expected, snippet)
+}
+
+// TestResponseAnnotationTextInTitle verifies that the SAX rewriter only wraps
+// text nodes, not content inside title attributes, even when the token text
+// matches part of an attribute value.
+func TestResponseAnnotationTextInTitle(t *testing.T) {
+	responseSnippet := `{
+		"snippet": "<span title=\"marmot/p:NN\">NN</span>"
+	}`
+
+	mappingList := config.MappingList{
+		ID:       "test-title-mapper",
+		FoundryA: "marmot",
+		LayerA:   "p",
+		FoundryB: "opennlp",
+		LayerB:   "p",
+		Mappings: []config.MappingRule{
+			"[NN] <> [NOUN]",
+		},
+	}
+
+	m, err := NewMapper([]config.MappingList{mappingList})
+	require.NoError(t, err)
+
+	var inputData any
+	err = json.Unmarshal([]byte(responseSnippet), &inputData)
+	require.NoError(t, err)
+
+	result, err := m.ApplyResponseMappings("test-title-mapper", MappingOptions{Direction: AtoB}, inputData)
+	require.NoError(t, err)
+
+	resultMap := result.(map[string]any)
+	snippet := resultMap["snippet"].(string)
+
+	expected := `<span title="marmot/p:NN"><span title="opennlp/p:NOUN" class="notinindex">NN</span></span>`
+	assert.Equal(t, expected, snippet)
+}
+
+// TestResponseAnnotationWhitespaceAroundText tests that annotations are applied
+// even when there is whitespace between the enclosing tag and the text content.
+// The old string-heuristic fails because it requires ">" immediately before the text.
+func TestResponseAnnotationWhitespaceAroundText(t *testing.T) {
+	responseSnippet := `{
+		"snippet": "<span title=\"marmot/p:DET\"> Der </span>"
+	}`
+
+	mappingList := config.MappingList{
+		ID:       "test-ws-mapper",
+		FoundryA: "marmot",
+		LayerA:   "p",
+		FoundryB: "opennlp",
+		LayerB:   "p",
+		Mappings: []config.MappingRule{
+			"[DET] <> [DT]",
+		},
+	}
+
+	m, err := NewMapper([]config.MappingList{mappingList})
+	require.NoError(t, err)
+
+	var inputData any
+	err = json.Unmarshal([]byte(responseSnippet), &inputData)
+	require.NoError(t, err)
+
+	result, err := m.ApplyResponseMappings("test-ws-mapper", MappingOptions{Direction: AtoB}, inputData)
+	require.NoError(t, err)
+
+	resultMap := result.(map[string]any)
+	snippet := resultMap["snippet"].(string)
+
+	// Whitespace should be preserved, annotation wraps only the token text
+	expected := `<span title="marmot/p:DET"> <span title="opennlp/p:DT" class="notinindex">Der</span> </span>`
+	assert.Equal(t, expected, snippet)
+}
+
+// TestResponseAnnotationCrossElementText tests annotation of individual tokens
+// whose text spans across sibling/child elements.
+func TestResponseAnnotationCrossElementText(t *testing.T) {
+	responseSnippet := `{
+		"snippet": "<span title=\"marmot/p:DET\">Die</span> <span title=\"base/s:s\"><span title=\"marmot/p:NN\">Sonne</span></span>"
+	}`
+
+	mappingList := config.MappingList{
+		ID:       "test-cross-mapper",
+		FoundryA: "marmot",
+		LayerA:   "p",
+		FoundryB: "opennlp",
+		LayerB:   "p",
+		Mappings: []config.MappingRule{
+			"[DET] <> [DT]",
+			"[NN] <> [NOUN]",
+		},
+	}
+
+	m, err := NewMapper([]config.MappingList{mappingList})
+	require.NoError(t, err)
+
+	var inputData any
+	err = json.Unmarshal([]byte(responseSnippet), &inputData)
+	require.NoError(t, err)
+
+	result, err := m.ApplyResponseMappings("test-cross-mapper", MappingOptions{Direction: AtoB}, inputData)
+	require.NoError(t, err)
+
+	resultMap := result.(map[string]any)
+	snippet := resultMap["snippet"].(string)
+
+	assert.Contains(t, snippet, `<span title="opennlp/p:DT" class="notinindex">Die</span>`)
+	assert.Contains(t, snippet, `<span title="opennlp/p:NOUN" class="notinindex">Sonne</span>`)
+	assert.Contains(t, snippet, `title="base/s:s"`)
+}
+
+// TestResponseAnnotationSubstringToken tests that a short token ("er") is
+// annotated only in its own text node and not when it appears as a prefix of
+// another word ("er Mann") in an earlier text node.
+func TestResponseAnnotationSubstringToken(t *testing.T) {
+	// "er" appears at the start of "er Mann" (NN span) and as standalone (PPER span).
+	// The old heuristic matches the first occurrence because "er" is preceded by ">"
+	// and followed by " ".
+	responseSnippet := `{
+		"snippet": "<span title=\"marmot/p:NN\">er Mann</span> <span title=\"marmot/p:PPER\">er</span>"
+	}`
+
+	mappingList := config.MappingList{
+		ID:       "test-sub-mapper",
+		FoundryA: "marmot",
+		LayerA:   "p",
+		FoundryB: "opennlp",
+		LayerB:   "p",
+		Mappings: []config.MappingRule{
+			"[PPER] <> [PRP]",
+		},
+	}
+
+	m, err := NewMapper([]config.MappingList{mappingList})
+	require.NoError(t, err)
+
+	var inputData any
+	err = json.Unmarshal([]byte(responseSnippet), &inputData)
+	require.NoError(t, err)
+
+	result, err := m.ApplyResponseMappings("test-sub-mapper", MappingOptions{Direction: AtoB}, inputData)
+	require.NoError(t, err)
+
+	resultMap := result.(map[string]any)
+	snippet := resultMap["snippet"].(string)
+
+	// The NN "er Mann" must remain unchanged; only the PPER "er" gets annotated
+	expected := `<span title="marmot/p:NN">er Mann</span> <span title="marmot/p:PPER"><span title="opennlp/p:PRP" class="notinindex">er</span></span>`
+	assert.Equal(t, expected, snippet)
+}
+
+// TestResponseAnnotationSelfClosingTags verifies that self-closing tags like
+// <br/> are preserved and do not interfere with annotation insertion.
+func TestResponseAnnotationSelfClosingTags(t *testing.T) {
+	responseSnippet := `{
+		"snippet": "<span title=\"marmot/p:DET\">Der</span><br/><span title=\"marmot/p:NN\">Mann</span>"
+	}`
+
+	mappingList := config.MappingList{
+		ID:       "test-br-mapper",
+		FoundryA: "marmot",
+		LayerA:   "p",
+		FoundryB: "opennlp",
+		LayerB:   "p",
+		Mappings: []config.MappingRule{
+			"[DET] <> [DT]",
+			"[NN]  <> [NOUN]",
+		},
+	}
+
+	m, err := NewMapper([]config.MappingList{mappingList})
+	require.NoError(t, err)
+
+	var inputData any
+	err = json.Unmarshal([]byte(responseSnippet), &inputData)
+	require.NoError(t, err)
+
+	result, err := m.ApplyResponseMappings("test-br-mapper", MappingOptions{Direction: AtoB}, inputData)
+	require.NoError(t, err)
+
+	snippet := result.(map[string]any)["snippet"].(string)
+
+	assert.Contains(t, snippet, "<br/>")
+	assert.Contains(t, snippet, `<span title="opennlp/p:DT" class="notinindex">Der</span>`)
+	assert.Contains(t, snippet, `<span title="opennlp/p:NOUN" class="notinindex">Mann</span>`)
+}
+
+// TestResponseAnnotationEntityReferences verifies that entity references
+// (&amp;, &lt;, etc.) are correctly preserved in output.
+func TestResponseAnnotationEntityReferences(t *testing.T) {
+	responseSnippet := `{
+		"snippet": "<span title=\"marmot/p:NN\">Haus &amp; Hof</span>"
+	}`
+
+	mappingList := config.MappingList{
+		ID:       "test-entity-mapper",
+		FoundryA: "marmot",
+		LayerA:   "p",
+		FoundryB: "opennlp",
+		LayerB:   "p",
+		Mappings: []config.MappingRule{
+			"[NN] <> [NOUN]",
+		},
+	}
+
+	m, err := NewMapper([]config.MappingList{mappingList})
+	require.NoError(t, err)
+
+	var inputData any
+	err = json.Unmarshal([]byte(responseSnippet), &inputData)
+	require.NoError(t, err)
+
+	result, err := m.ApplyResponseMappings("test-entity-mapper", MappingOptions{Direction: AtoB}, inputData)
+	require.NoError(t, err)
+
+	snippet := result.(map[string]any)["snippet"].(string)
+
+	// Entity reference must be preserved (re-encoded) in the annotated output
+	expected := `<span title="marmot/p:NN"><span title="opennlp/p:NOUN" class="notinindex">Haus &amp; Hof</span></span>`
+	assert.Equal(t, expected, snippet)
+}
+
+// TestResponseAnnotationEntityLtGt verifies &lt; and &gt; are re-encoded.
+func TestResponseAnnotationEntityLtGt(t *testing.T) {
+	responseSnippet := `{
+		"snippet": "<span title=\"marmot/p:SYM\">&lt;tag&gt;</span>"
+	}`
+
+	mappingList := config.MappingList{
+		ID:       "test-ltgt-mapper",
+		FoundryA: "marmot",
+		LayerA:   "p",
+		FoundryB: "opennlp",
+		LayerB:   "p",
+		Mappings: []config.MappingRule{
+			"[SYM] <> [PUNCT]",
+		},
+	}
+
+	m, err := NewMapper([]config.MappingList{mappingList})
+	require.NoError(t, err)
+
+	var inputData any
+	err = json.Unmarshal([]byte(responseSnippet), &inputData)
+	require.NoError(t, err)
+
+	result, err := m.ApplyResponseMappings("test-ltgt-mapper", MappingOptions{Direction: AtoB}, inputData)
+	require.NoError(t, err)
+
+	snippet := result.(map[string]any)["snippet"].(string)
+
+	expected := `<span title="marmot/p:SYM"><span title="opennlp/p:PUNCT" class="notinindex">&lt;tag&gt;</span></span>`
+	assert.Equal(t, expected, snippet)
+}
+
+// TestResponseAnnotationCDATAGraceful verifies that a CDATA section in the
+// snippet does not cause errors and is passed through unchanged.
+func TestResponseAnnotationCDATAGraceful(t *testing.T) {
+	responseSnippet := `{
+		"snippet": "<span title=\"marmot/p:DET\">Der</span><![CDATA[ raw ]]><span title=\"marmot/p:NN\">Mann</span>"
+	}`
+
+	mappingList := config.MappingList{
+		ID:       "test-cdata-mapper",
+		FoundryA: "marmot",
+		LayerA:   "p",
+		FoundryB: "opennlp",
+		LayerB:   "p",
+		Mappings: []config.MappingRule{
+			"[NN] <> [NOUN]",
+		},
+	}
+
+	m, err := NewMapper([]config.MappingList{mappingList})
+	require.NoError(t, err)
+
+	var inputData any
+	err = json.Unmarshal([]byte(responseSnippet), &inputData)
+	require.NoError(t, err)
+
+	result, err := m.ApplyResponseMappings("test-cdata-mapper", MappingOptions{Direction: AtoB}, inputData)
+	require.NoError(t, err)
+
+	snippet := result.(map[string]any)["snippet"].(string)
+
+	assert.Contains(t, snippet, "<![CDATA[ raw ]]>")
+	assert.Contains(t, snippet, `<span title="opennlp/p:NOUN" class="notinindex">Mann</span>`)
+}
+
+// TestResponseAnnotationOverlappingSpans verifies that when two independent
+// rules match the same token, both annotations are applied.
+func TestResponseAnnotationOverlappingSpans(t *testing.T) {
+	responseSnippet := `{
+		"snippet": "<span title=\"marmot/p:DET\"><span title=\"marmot/m:case:nom\">Der</span></span>"
+	}`
+
+	mappingList := config.MappingList{
+		ID:       "test-overlap-mapper",
+		FoundryA: "marmot",
+		LayerA:   "p",
+		FoundryB: "opennlp",
+		LayerB:   "p",
+		Mappings: []config.MappingRule{
+			"[DET] <> [DT]",
+		},
+	}
+
+	m, err := NewMapper([]config.MappingList{mappingList})
+	require.NoError(t, err)
+
+	var inputData any
+	err = json.Unmarshal([]byte(responseSnippet), &inputData)
+	require.NoError(t, err)
+
+	result, err := m.ApplyResponseMappings("test-overlap-mapper", MappingOptions{Direction: AtoB}, inputData)
+	require.NoError(t, err)
+
+	snippet := result.(map[string]any)["snippet"].(string)
+
+	// The existing nested structure must be preserved, with new annotation added
+	assert.Contains(t, snippet, `title="marmot/p:DET"`)
+	assert.Contains(t, snippet, `title="marmot/m:case:nom"`)
+	assert.Contains(t, snippet, `title="opennlp/p:DT" class="notinindex"`)
+	assert.Contains(t, snippet, "Der")
+}
+
+// TestResponseAnnotationEmptyTextNodes verifies that empty or whitespace-only
+// text nodes are passed through without errors and without spurious annotations.
+func TestResponseAnnotationEmptyTextNodes(t *testing.T) {
+	responseSnippet := `{
+		"snippet": "<span title=\"marmot/p:DET\"></span> <span title=\"marmot/p:NN\">Mann</span>"
+	}`
+
+	mappingList := config.MappingList{
+		ID:       "test-empty-mapper",
+		FoundryA: "marmot",
+		LayerA:   "p",
+		FoundryB: "opennlp",
+		LayerB:   "p",
+		Mappings: []config.MappingRule{
+			"[DET] <> [DT]",
+			"[NN]  <> [NOUN]",
+		},
+	}
+
+	m, err := NewMapper([]config.MappingList{mappingList})
+	require.NoError(t, err)
+
+	var inputData any
+	err = json.Unmarshal([]byte(responseSnippet), &inputData)
+	require.NoError(t, err)
+
+	result, err := m.ApplyResponseMappings("test-empty-mapper", MappingOptions{Direction: AtoB}, inputData)
+	require.NoError(t, err)
+
+	snippet := result.(map[string]any)["snippet"].(string)
+
+	// The empty DET span should not get an annotation
+	// The NN token "Mann" should be annotated
+	assert.Contains(t, snippet, `<span title="marmot/p:DET"></span>`)
+	assert.Contains(t, snippet, `<span title="opennlp/p:NOUN" class="notinindex">Mann</span>`)
+}
+
+// TestResponseAnnotationWhitespaceOnlyNodes verifies that whitespace-only text
+// nodes are preserved without annotations.
+func TestResponseAnnotationWhitespaceOnlyNodes(t *testing.T) {
+	responseSnippet := `{
+		"snippet": "<span title=\"marmot/p:DET\">   </span><span title=\"marmot/p:NN\">Mann</span>"
+	}`
+
+	mappingList := config.MappingList{
+		ID:       "test-wsonly-mapper",
+		FoundryA: "marmot",
+		LayerA:   "p",
+		FoundryB: "opennlp",
+		LayerB:   "p",
+		Mappings: []config.MappingRule{
+			"[DET] <> [DT]",
+			"[NN]  <> [NOUN]",
+		},
+	}
+
+	m, err := NewMapper([]config.MappingList{mappingList})
+	require.NoError(t, err)
+
+	var inputData any
+	err = json.Unmarshal([]byte(responseSnippet), &inputData)
+	require.NoError(t, err)
+
+	result, err := m.ApplyResponseMappings("test-wsonly-mapper", MappingOptions{Direction: AtoB}, inputData)
+	require.NoError(t, err)
+
+	snippet := result.(map[string]any)["snippet"].(string)
+
+	// Whitespace-only text should not be annotated
+	assert.Contains(t, snippet, `<span title="marmot/p:DET">   </span>`)
+	assert.Contains(t, snippet, `<span title="opennlp/p:NOUN" class="notinindex">Mann</span>`)
+}
+
 // TestResponseMappingWithLayerOverride tests layer precedence rules
 func TestResponseMappingWithLayerOverride(t *testing.T) {
 	// Test 1: Explicit layer in mapping rule should take precedence over MappingOptions