chatlog2/internal/wechatdb/repository/index.go

package repository

import (
	"context"
	"errors"
	"fmt"
	"sort"
	"strings"
	"time"

	"github.com/rs/zerolog/log"

	"github.com/sjzar/chatlog/internal/model"
	"github.com/sjzar/chatlog/internal/wechatdb/indexer"
	"github.com/sjzar/chatlog/pkg/util"
)

type bleveIndexable interface {
	ListTalkers(ctx context.Context) ([]string, error)
	IterateMessages(ctx context.Context, talkers []string, fn func(*model.Message) error) error
}

func (r *Repository) initIndex() error {
	if r.indexPath == "" {
		return nil
	}

	idx, err := indexer.Open(r.indexPath)
	if err != nil {
		return err
	}

	r.index = idx
	r.indexCtx, r.indexCancel = context.WithCancel(context.Background())

	go func() {
		ready, err := r.ensureIndex(r.indexCtx)
		if err != nil && !errors.Is(err, context.Canceled) {
			log.Warn().Err(err).Msg("ensure bleve index failed")
			return
		}
		if ready {
			log.Info().Msg("bleve index ready")
		}
	}()

	return nil
}

func (r *Repository) ensureIndex(ctx context.Context) (bool, error) {
	if r.index == nil {
		return false, nil
	}

	if ctx == nil {
		ctx = context.Background()
	}

	versionMatched, err := r.index.EnsureVersion()
	if err != nil {
		return false, err
	}

	if !versionMatched {
		r.indexMu.Lock()
		r.indexStatus.Ready = false
		r.indexStatus.Progress = 0
		r.indexFingerprint = ""
		r.indexMu.Unlock()
	}

	fp, err := r.ds.GetDatasetFingerprint(ctx)
	if err != nil {
		return false, err
	}
	if strings.TrimSpace(fp) == "" {
		return false, fmt.Errorf("dataset fingerprint is empty")
	}

	r.indexMu.Lock()
	if r.indexFingerprint == fp && r.indexStatus.Ready && !r.indexStatus.InProgress {
		status := r.indexStatus
		r.indexMu.Unlock()
		if status.Progress < 1 {
			r.updateIndexProgress(1)
		}
		return true, nil
	}
	if r.indexStatus.InProgress {
		r.indexMu.Unlock()
		return false, nil
	}
	r.indexStatus.InProgress = true
	r.indexStatus.Ready = false
	r.indexStatus.Progress = 0
	r.indexStatus.LastStartedAt = time.Now()
	r.indexStatus.LastError = ""
	r.indexMu.Unlock()

	err = r.rebuildIndex(ctx, fp)
	if err != nil {
		if err == context.Canceled || errors.Is(err, context.Canceled) {
			r.indexMu.Lock()
			r.indexStatus.InProgress = false
			r.indexMu.Unlock()
			return false, err
		}

		r.indexMu.Lock()
		r.indexStatus.InProgress = false
		r.indexStatus.LastError = err.Error()
		r.indexMu.Unlock()
		return false, err
	}

	r.indexMu.Lock()
	r.indexFingerprint = fp
	r.indexStatus.InProgress = false
	r.indexStatus.Ready = true
	r.indexStatus.Progress = 1
	r.indexStatus.LastCompletedAt = time.Now()
	r.indexMu.Unlock()

	return true, nil
}

func (r *Repository) rebuildIndex(ctx context.Context, fp string) error {
	indexable, ok := r.ds.(bleveIndexable)
	if !ok {
		return fmt.Errorf("datasource does not support bleve indexing")
	}

	if err := r.index.Reset(); err != nil {
		return err
	}
	if _, err := r.index.EnsureVersion(); err != nil {
		return err
	}

	talkers, err := indexable.ListTalkers(ctx)
	if err != nil {
		return err
	}

	if len(talkers) == 0 {
		if err := r.index.UpdateFingerprint(fp); err != nil {
			return err
		}
		return r.index.UpdateLastBuilt(time.Now())
	}

	sort.Strings(talkers)

	batch := make([]*model.Message, 0, 512)
	flush := func() error {
		if len(batch) == 0 {
			return nil
		}
		if err := r.index.IndexMessages(batch); err != nil {
			return err
		}
		batch = batch[:0]
		return nil
	}

	total := float64(len(talkers))
	for i, talker := range talkers {
		if err := ctx.Err(); err != nil {
			return err
		}

		handler := func(msg *model.Message) error {
			if msg == nil {
				return nil
			}
			batch = append(batch, msg)
			if len(batch) >= cap(batch) {
				return flush()
			}
			return nil
		}

		if err := indexable.IterateMessages(ctx, []string{talker}, handler); err != nil {
			return err
		}

		if err := flush(); err != nil {
			return err
		}

		r.updateIndexProgress(float64(i+1) / total)
	}

	if err := r.index.UpdateFingerprint(fp); err != nil {
		return err
	}
	if err := r.index.UpdateLastBuilt(time.Now()); err != nil {
		return err
	}

	return nil
}

func (r *Repository) updateIndexProgress(progress float64) {
	if r.index == nil {
		return
	}

	if progress < 0 {
		progress = 0
	}
	if progress > 1 {
		progress = 1
	}

	r.indexMu.Lock()
	r.indexStatus.Progress = progress
	r.indexMu.Unlock()
}

func (r *Repository) indexStatusSnapshot() *model.SearchIndexStatus {
	if r.index == nil {
		return nil
	}

	r.indexMu.Lock()
	status := r.indexStatus
	r.indexMu.Unlock()

	copied := status
	return &copied
}

func (r *Repository) searchMessagesWithIndex(ctx context.Context, req *model.SearchRequest) (*model.SearchResponse, error) {
	makeEmpty := func() *model.SearchResponse {
		return &model.SearchResponse{
			Total:      0,
			Hits:       []*model.SearchHit{},
			DurationMs: 0,
			Limit:      req.Limit,
			Offset:     req.Offset,
			Query:      req.Query,
			Talker:     req.Talker,
			Sender:     req.Sender,
			Start:      req.Start,
			End:        req.End,
			Index:      r.indexStatusSnapshot(),
		}
	}

	if strings.TrimSpace(req.Query) == "" {
		return makeEmpty(), nil
	}

	if r.index == nil {
		return makeEmpty(), nil
	}

	ready, err := r.ensureIndex(ctx)
	if err != nil {
		return nil, err
	}
	if !ready {
		return makeEmpty(), nil
	}

	talkers := util.Str2List(req.Talker, ",")
	senders := util.Str2List(req.Sender, ",")

	startUnix := int64(0)
	if !req.Start.IsZero() {
		startUnix = req.Start.Unix()
	}
	endUnix := int64(0)
	if !req.End.IsZero() {
		endUnix = req.End.Unix()
	}
	if startUnix > 0 && endUnix > 0 && endUnix < startUnix {
		startUnix, endUnix = endUnix, startUnix
	}

	begin := time.Now()
	hits, total, err := r.index.Search(req, talkers, senders, startUnix, endUnix, req.Offset, req.Limit)
	if err != nil {
		return nil, err
	}

	mapped := make([]*model.SearchHit, 0, len(hits))
	for _, hit := range hits {
		if hit == nil || hit.Message == nil {
			continue
		}
		mapped = append(mapped, &model.SearchHit{
			Message: hit.Message,
			Snippet: hit.Snippet,
			Score:   hit.Score,
		})
	}

	resp := &model.SearchResponse{
		Total:      total,
		Hits:       mapped,
		DurationMs: time.Since(begin).Milliseconds(),
		Limit:      req.Limit,
		Offset:     req.Offset,
		Query:      req.Query,
		Talker:     req.Talker,
		Sender:     req.Sender,
		Start:      req.Start,
		End:        req.End,
		Index:      r.indexStatusSnapshot(),
	}

	return resp, nil
}