xinmi/internal/application/service/wiki_source_identity.go

package service

import (
	"context"
	"encoding/json"
	"sort"
	"strings"
	"time"
	"unicode"

	"github.com/Tencent/WeKnora/internal/logger"
	"github.com/Tencent/WeKnora/internal/types"
)

const wikiSourceIdentityVersion = "wiki-source-identity-v2"

// Identity failures are not output-format failures: retrying or splitting an
// unsupported identity must not silently discard it and report complete reading.
type wikiSourceIdentityError struct {
	Code       string `json:"code"`
	Kind       string `json:"kind"`
	Index      int    `json:"index"`
	SourceHash string `json:"source_hash,omitempty"`
	Contract   string `json:"contract"`
}

func (e *wikiSourceIdentityError) Error() string {
	b, _ := json.Marshal(e)
	return string(b)
}

func wikiIdentityFailure(code, kind string, index int, sources []string) error {
	e := &wikiSourceIdentityError{Code: code, Kind: kind, Index: index, Contract: wikiSourceIdentityVersion}
	if sources != nil {
		e.SourceHash = wikiCheckpointHash(sources)
	}
	return e
}

func wikiSourceIdentityKey(kind, title string) string {
	return kind + "\x00" + strings.TrimSpace(title)
}

// Only one layout newline between Han characters may interrupt a name. Match
// against each original source separately; never normalize the evidence itself.
func wikiSourceContainsEntity(source, name string) bool {
	if name == "" {
		return false
	}
	if strings.Contains(source, name) {
		return true
	}
	runes := []rune(name)
	for i := 1; i < len(runes); i++ {
		if !unicode.Is(unicode.Han, runes[i-1]) || !unicode.Is(unicode.Han, runes[i]) {
			continue
		}
		left, right := string(runes[:i]), string(runes[i:])
		if strings.Contains(source, left+"\n"+right) || strings.Contains(source, left+"\r\n"+right) {
			return true
		}
	}
	return false
}

func validateWikiSourceItemIdentity(item extractedItem, kind string, index int, sources []string) error {
	if len(item.Aliases) != 0 {
		return wikiIdentityFailure("wiki_identity_alias_unsupported", kind, index, sources)
	}
	if kind != types.WikiPageTypeEntity {
		return nil
	}
	name := strings.TrimSpace(item.Name)
	for _, source := range sources {
		if wikiSourceContainsEntity(source, name) {
			return nil
		}
	}
	return wikiIdentityFailure("wiki_identity_name_not_in_source", kind, index, sources)
}

func validateWikiSourceIdentities(result combinedExtraction, sources ...string) error {
	if err := validateWikiFreshSourceIdentities(result, sources...); err != nil {
		return err
	}
	return validateWikiIdentityCollisions(result)
}

// Fresh extraction proposals have no references yet. Their source identities
// are checked per leaf before complete-set routing identifiers are assigned.
func validateWikiFreshSourceIdentities(result combinedExtraction, sources ...string) error {
	if err := validateWikiCandidates(result); err != nil {
		return err
	}
	for _, group := range []struct {
		kind  string
		items []extractedItem
	}{{types.WikiPageTypeEntity, result.Entities}, {types.WikiPageTypeConcept, result.Concepts}} {
		for i, item := range group.items {
			if err := validateWikiSourceItemIdentity(item, group.kind, i, sources); err != nil {
				return err
			}
		}
	}
	return nil
}

// pruneIdentityCollisions keeps the first identity that claims a slug and drops
// later candidates that claim it for a different identity, plus candidates that
// carry unsupported aliases. Drops are counted, never silent.
func pruneIdentityCollisions(result combinedExtraction) (combinedExtraction, int) {
	owners := make(map[string]string)
	dropped := 0
	prune := func(kind string, items []extractedItem) []extractedItem {
		kept := make([]extractedItem, 0, len(items))
		for _, item := range items {
			if len(item.Aliases) != 0 {
				dropped++
				continue
			}
			identity := wikiSourceIdentityKey(kind, item.Name)
			if owner, exists := owners[item.Slug]; exists && owner != identity {
				dropped++
				continue
			}
			owners[item.Slug] = identity
			kept = append(kept, item)
		}
		return kept
	}
	pruned := combinedExtraction{
		Entities:         prune(types.WikiPageTypeEntity, result.Entities),
		Concepts:         prune(types.WikiPageTypeConcept, result.Concepts),
		CandidatePruning: result.CandidatePruning,
	}
	return pruned, dropped
}

// pruneIncomingWindowCandidates keeps the first identity that owns a slug and
// drops later candidates that claim it for a different identity, plus candidates
// that cannot be published because they carry aliases. Drops are counted so the
// decision is auditable; a window whose candidates are all dropped fails closed
// instead of silently contributing nothing.
func pruneIncomingWindowCandidates(accumulated, incoming combinedExtraction) (combinedExtraction, int, error) {
	combined := combinedExtraction{
		Entities: append(append([]extractedItem(nil), accumulated.Entities...), incoming.Entities...),
		Concepts: append(append([]extractedItem(nil), accumulated.Concepts...), incoming.Concepts...),
	}
	pruned, dropped := pruneIdentityCollisions(combined)
	// Earlier windows own their slugs, so pruning must never shorten the prefix.
	if len(pruned.Entities) < len(accumulated.Entities) || len(pruned.Concepts) < len(accumulated.Concepts) {
		return combinedExtraction{}, dropped, wikiIdentityFailure("wiki_identity_slug_conflict", types.WikiPageTypeEntity, -1, nil)
	}
	kept := combinedExtraction{
		Entities: pruned.Entities[len(accumulated.Entities):],
		Concepts: pruned.Concepts[len(accumulated.Concepts):],
	}
	if len(kept.Entities)+len(kept.Concepts) == 0 && len(incoming.Entities)+len(incoming.Concepts) > 0 {
		return combinedExtraction{}, dropped, wikiIdentityFailure("wiki_identity_slug_conflict", types.WikiPageTypeEntity, -1, nil)
	}
	if err := validateWikiIdentityCollisions(pruned); err != nil {
		return combinedExtraction{}, dropped, err
	}
	return kept, dropped, nil
}

// Check before any merge can hide the second title by converting it to an alias.
func validateWikiIdentityCollisions(result combinedExtraction) error {
	seen := make(map[string]string)
	for _, group := range []struct {
		kind  string
		items []extractedItem
	}{{types.WikiPageTypeEntity, result.Entities}, {types.WikiPageTypeConcept, result.Concepts}} {
		for i, item := range group.items {
			if len(item.Aliases) != 0 {
				return wikiIdentityFailure("wiki_identity_alias_unsupported", group.kind, i, nil)
			}
			identity := wikiSourceIdentityKey(group.kind, item.Name)
			if previous, exists := seen[item.Slug]; exists && previous != identity {
				return wikiIdentityFailure("wiki_identity_slug_conflict", group.kind, i, nil)
			}
			seen[item.Slug] = identity
		}
	}
	return nil
}

func validateWikiCitationIdentityMerge(entities, concepts []extractedItem, declarations []newSlugFromCitation) error {
	all := combinedExtraction{Entities: append([]extractedItem(nil), entities...), Concepts: append([]extractedItem(nil), concepts...)}
	for _, declaration := range declarations {
		item := extractedItem{Name: declaration.Name, Slug: declaration.Slug, Aliases: declaration.Aliases}
		if declaration.Type == types.WikiPageTypeEntity {
			all.Entities = append(all.Entities, item)
		} else {
			all.Concepts = append(all.Concepts, item)
		}
	}
	return validateWikiIdentityCollisions(all)
}

func mergeWikiSourceCitationItems(entities, concepts []extractedItem, citations map[string][]string, declarations []newSlugFromCitation) ([]extractedItem, []extractedItem, int, error) {
	if err := validateWikiCitationIdentityMerge(entities, concepts, declarations); err != nil {
		return nil, nil, 0, err
	}
	trimNames := func(items []extractedItem) []extractedItem {
		out := append([]extractedItem(nil), items...)
		for i := range out {
			out[i].Name = strings.TrimSpace(out[i].Name)
		}
		return out
	}
	declarations = append([]newSlugFromCitation(nil), declarations...)
	for i := range declarations {
		declarations[i].Name = strings.TrimSpace(declarations[i].Name)
	}
	entities, concepts, uncited := mergeCitationsIntoItems(trimNames(entities), trimNames(concepts), citations, declarations)
	if err := validateWikiIdentityCollisions(combinedExtraction{Entities: entities, Concepts: concepts}); err != nil {
		return nil, nil, 0, err
	}
	return entities, concepts, uncited, nil
}

func validateWikiPageIdentityUpdates(page *types.WikiPage, updates []SlugUpdate) error {
	var candidates combinedExtraction
	for i, update := range updates {
		if update.Type != types.WikiPageTypeEntity && update.Type != types.WikiPageTypeConcept {
			continue
		}
		if page != nil && (page.PageType != update.Type || strings.TrimSpace(page.Title) != strings.TrimSpace(update.Item.Name)) {
			return wikiIdentityFailure("wiki_identity_slug_conflict", update.Type, i, nil)
		}
		if update.Item.Slug != update.Slug {
			return wikiIdentityFailure("wiki_identity_slug_conflict", update.Type, i, nil)
		}
		if update.Type == types.WikiPageTypeEntity {
			candidates.Entities = append(candidates.Entities, update.Item)
		} else {
			candidates.Concepts = append(candidates.Concepts, update.Item)
		}
	}
	if err := validateWikiCandidates(candidates); err != nil {
		return err
	}
	return validateWikiIdentityCollisions(candidates)
}

func (s *wikiIngestService) claimWikiSourceIdentity(ctx context.Context, kbID, kind, title, proposed string, authoritative bool, batch *WikiBatchContext) (string, error) {
	key := wikiSourceIdentityVersion + ":" + wikiSourceIdentityKey(kind, title)
	claim := proposed
	if s.redisClient != nil {
		auth := "0"
		if authoritative {
			auth = "1"
		}
		value, err := s.redisClient.Eval(ctx, wikiIdentityClaimScript,
			[]string{wikiIdentityClaimPrefix + wikiSourceIdentityVersion + ":" + kbID + ":" + wikiCheckpointHash(key)},
			proposed, int64(wikiIdentityClaimTTL/time.Second), auth, kind+"/").Result()
		if err != nil {
			return "", wikiIdentityFailure("wiki_identity_claim_failed", kind, 0, nil)
		}
		claim = identityClaimString(value)
		if !strings.HasPrefix(claim, kind+"/") || strings.TrimPrefix(claim, kind+"/") == "" {
			return "", wikiIdentityFailure("wiki_identity_claim_invalid", kind, 0, nil)
		}
	}
	if batch != nil {
		if authoritative || s.redisClient != nil {
			batch.identityClaims.Store(key, claim)
		} else {
			actual, _ := batch.identityClaims.LoadOrStore(key, claim)
			claim, _ = actual.(string)
		}
	}
	return claim, nil
}

// resolveWikiIdentityTargets reconciles candidates with the pages that already own
// their routing slug. A spelling-only title difference is adopted from the existing
// page so one identity keeps owning one page; a genuinely different identity on the
// same slug is dropped and counted, because a document that already read its whole
// source must not be discarded for a model-written name variant. Only a set that
// cannot be published at all still fails closed.
func (s *wikiIngestService) resolveWikiIdentityTargets(ctx context.Context, kbID string, result combinedExtraction) (combinedExtraction, int, error) {
	if err := validateWikiIdentityCollisions(result); err != nil {
		return combinedExtraction{}, 0, err
	}
	if s.wikiService == nil {
		return result, 0, nil
	}
	items := append(append([]extractedItem(nil), result.Entities...), result.Concepts...)
	slugs := make([]string, 0, len(items))
	for _, item := range items {
		slugs = append(slugs, item.Slug)
	}
	if len(slugs) == 0 {
		return result, 0, nil
	}
	pages, err := s.wikiService.ListBySlugs(ctx, kbID, slugs)
	if err != nil {
		return combinedExtraction{}, 0, wikiIdentityFailure("wiki_identity_lookup_failed", "targets", 0, nil)
	}
	dropped := 0
	adjust := func(kind string, list []extractedItem) []extractedItem {
		kept := make([]extractedItem, 0, len(list))
		for _, item := range list {
			page := pages[item.Slug]
			if page == nil {
				kept = append(kept, item)
				continue
			}
			title, name := strings.TrimSpace(page.Title), strings.TrimSpace(item.Name)
			switch {
			case page.PageType != kind:
				dropped++
			case title == name:
				kept = append(kept, item)
			case normalizeWikiIdentityTitle(title) == normalizeWikiIdentityTitle(name):
				// Same identity, different spelling: keep the page's own title.
				item.Name = title
				kept = append(kept, item)
			default:
				dropped++
			}
		}
		return kept
	}
	out := combinedExtraction{
		Entities:         adjust(types.WikiPageTypeEntity, result.Entities),
		Concepts:         adjust(types.WikiPageTypeConcept, result.Concepts),
		CandidatePruning: result.CandidatePruning,
	}
	if len(out.Entities)+len(out.Concepts) == 0 && len(result.Entities)+len(result.Concepts) > 0 {
		return combinedExtraction{}, dropped, wikiIdentityFailure("wiki_identity_slug_conflict", types.WikiPageTypeEntity, -1, nil)
	}
	if err := validateWikiIdentityCollisions(out); err != nil {
		return combinedExtraction{}, dropped, err
	}
	out.CandidatePruning.SlugConflicts += dropped
	return out, dropped, nil
}

// The normalized repository query is only a shortlist. Literal title equality,
// versioned claims, and target validation decide identity without an LLM merge.
func (s *wikiIngestService) resolveWikiSourceIdentities(ctx context.Context, kbID string, result combinedExtraction, batch *WikiBatchContext) (combinedExtraction, error) {
	if err := validateWikiCandidates(result); err != nil {
		return combinedExtraction{}, err
	}
	// Two model-written names can claim one routing slug. The first identity keeps
	// it; the later candidate is dropped and counted so the document continues.
	pruned, prunedCount := pruneIdentityCollisions(result)
	if prunedCount > 0 {
		if len(pruned.Entities)+len(pruned.Concepts) == 0 {
			return combinedExtraction{}, wikiIdentityFailure("wiki_identity_slug_conflict", types.WikiPageTypeEntity, -1, nil)
		}
		result.CandidatePruning.SlugConflicts += prunedCount
		result = pruned
		logger.Infof(ctx, "wiki source identity pruning: policy=%s slug_conflicts=%d kept=%d",
			wikiSourceIdentityVersion, result.CandidatePruning.SlugConflicts, len(result.Entities)+len(result.Concepts))
	}
	targeted, targetDrops, err := s.resolveWikiIdentityTargets(ctx, kbID, result)
	if err != nil {
		return combinedExtraction{}, err
	}
	if targetDrops > 0 {
		// Counted, content-free evidence that a candidate collided with the page
		// that already owns its slug instead of ending the document.
		logger.Infof(ctx, "wiki identity target pruning: policy=%s dropped=%d kept=%d",
			wikiSourceIdentityVersion, targetDrops, len(targeted.Entities)+len(targeted.Concepts))
	}
	result = targeted
	// Seed the audit counters: identity work already performed for this document
	// must stay visible in the final pruning report.
	out := combinedExtraction{Entities: []extractedItem{}, Concepts: []extractedItem{}, CandidatePruning: result.CandidatePruning}
	for _, group := range []struct {
		kind  string
		items []extractedItem
		dst   *[]extractedItem
	}{{types.WikiPageTypeEntity, result.Entities, &out.Entities}, {types.WikiPageTypeConcept, result.Concepts, &out.Concepts}} {
		exact := make(map[string][]string)
		if s.wikiService != nil && len(group.items) != 0 {
			normalized := make([]string, 0, len(group.items))
			for _, item := range group.items {
				normalized = appendUniqueString(normalized, normalizeWikiIdentityTitle(item.Name))
			}
			pages, err := s.wikiService.FindPagesByNormalizedTitles(ctx, kbID, group.kind, normalized)
			if err != nil {
				return combinedExtraction{}, wikiIdentityFailure("wiki_identity_lookup_failed", group.kind, 0, nil)
			}
			for _, page := range pages {
				if page != nil && page.PageType == group.kind {
					name := strings.TrimSpace(page.Title)
					exact[name] = append(exact[name], page.Slug)
				}
			}
		}
		for _, original := range group.items {
			item := original
			item.Name = strings.TrimSpace(item.Name)
			targets := exact[item.Name]
			if len(targets) > 0 {
				sort.Strings(targets)
				item.Slug = targets[0]
				for _, target := range targets {
					if target == original.Slug {
						item.Slug = target
					}
				}
			}
			var err error
			item.Slug, err = s.claimWikiSourceIdentity(ctx, kbID, group.kind, item.Name, item.Slug, len(targets) > 0, batch)
			if err != nil {
				return combinedExtraction{}, err
			}
			*group.dst = append(*group.dst, item)
		}
	}
	// Claiming may resolve two identities onto one slug; prune and count rather
	// than discarding a whole document that has already read its source.
	if prunedOut, prunedCount := pruneIdentityCollisions(out); prunedCount > 0 {
		if len(prunedOut.Entities)+len(prunedOut.Concepts) == 0 {
			return combinedExtraction{}, wikiIdentityFailure("wiki_identity_slug_conflict", types.WikiPageTypeEntity, -1, nil)
		}
		prunedOut.CandidatePruning.SlugConflicts += prunedCount
		out = prunedOut
		logger.Infof(ctx, "wiki source identity pruning after claim: policy=%s slug_conflicts=%d kept=%d",
			wikiSourceIdentityVersion, out.CandidatePruning.SlugConflicts, len(out.Entities)+len(out.Concepts))
	}
	targetedOut, claimedDrops, err := s.resolveWikiIdentityTargets(ctx, kbID, out)
	if err != nil {
		return combinedExtraction{}, err
	}
	if claimedDrops > 0 {
		logger.Infof(ctx, "wiki identity target pruning after claim: policy=%s dropped=%d kept=%d",
			wikiSourceIdentityVersion, claimedDrops, len(targetedOut.Entities)+len(targetedOut.Concepts))
	}
	out = targetedOut
	merged := combinedExtraction{}
	if err := mergeWikiWindowCandidates(&merged, out); err != nil {
		return combinedExtraction{}, err
	}
	// Carry the counted drops forward: the merge builds a fresh accumulator, and the
	// audit trail must not lose what identity resolution already removed.
	merged.CandidatePruning.SlugConflicts += out.CandidatePruning.SlugConflicts
	merged.CandidatePruning.Unmatched += out.CandidatePruning.Unmatched
	merged.CandidatePruning.Capped += out.CandidatePruning.Capped
	merged.CandidatePruning.Duplicates += out.CandidatePruning.Duplicates
	return merged, nil
}

func (s *wikiIngestService) remapWikiSourceIdentityUpdates(ctx context.Context, kbID string, updates map[string][]SlugUpdate, batch *WikiBatchContext) (map[string][]SlugUpdate, error) {
	var candidates combinedExtraction
	for _, group := range updates {
		for _, update := range group {
			switch update.Type {
			case types.WikiPageTypeEntity:
				candidates.Entities = append(candidates.Entities, update.Item)
			case types.WikiPageTypeConcept:
				candidates.Concepts = append(candidates.Concepts, update.Item)
			}
		}
	}
	resolved, err := s.resolveWikiSourceIdentities(ctx, kbID, candidates, batch)
	if err != nil {
		return nil, err
	}
	targets := make(map[string]string)
	for _, item := range append(resolved.Entities, resolved.Concepts...) {
		kind, _, _ := strings.Cut(item.Slug, "/")
		targets[wikiSourceIdentityKey(kind, item.Name)] = item.Slug
	}
	out := make(map[string][]SlugUpdate)
	for _, group := range updates {
		for _, update := range group {
			if target := targets[wikiSourceIdentityKey(update.Type, update.Item.Name)]; target != "" {
				update.Slug, update.Item.Slug = target, target
			}
			out[update.Slug] = append(out[update.Slug], update)
		}
	}
	return out, nil
}