Agent-System/internal/materialindex/pipeline.go

package materialindex

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"path/filepath"
	"regexp"
	"strings"
	"sync"
	"time"
	"unicode/utf8"

	"chat-agent-system/internal/config"
	"chat-agent-system/internal/kb"
	"chat-agent-system/internal/mineru"
	"chat-agent-system/internal/squad"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

const (
	engineMinerU   = "mineru"
	statePending   = "pending"
	stateRunning   = "running"
	stateDone      = "done"
	stateFailed    = "failed"
	maxAttempts    = 3
	defaultWait    = 90 * time.Second
	chunkRuneLimit = 1200
)

var pageAnchorRe = regexp.MustCompile(`(?m)^<!--\s*page\s+(\d+)\s*-->\s*$`)

// ParseRecord is a row from asset_parses.
type ParseRecord struct {
	ID           int64
	BlobSHA256   string
	FileName     string
	Engine       string
	State        string
	TaskID       string
	Markdown     string
	ContentList  json.RawMessage
	TotalPages   int
	ErrorMessage string
	Attempts     int
}

// ChunkHit is a retrieved material chunk.
type ChunkHit struct {
	FileName  string
	Title     string
	Content   string
	PageStart int
	PageEnd   int
	Score     float64
}

// Pipeline owns parse cache + chunk index.
type Pipeline struct {
	pool      *pgxpool.Pool
	client    *mineru.Client
	embedder  *kb.Embedder
	reranker  *kb.Reranker
	waitLimit time.Duration
	fullMax   int
	ragTopK   int

	inflightMu sync.Mutex
	inflight   map[string]chan struct{}
}

// NewPipeline builds a material parse/index pipeline. client may be nil/disabled.
func NewPipeline(cfg *config.Config, pool *pgxpool.Pool) *Pipeline {
	var client *mineru.Client
	var embedder *kb.Embedder
	var reranker *kb.Reranker
	wait := defaultWait
	fullMax := 60 * 1024
	topK := 8
	if cfg != nil {
		client = mineru.NewClient(mineru.Config{
			Endpoint:     cfg.MineruEndpointURL,
			APIKey:       cfg.MineruAPIKey,
			ModelVersion: cfg.MineruModelVersion,
			Language:     cfg.MineruLanguage,
			HTTPTimeout:  cfg.MineruTimeout,
		})
		if cfg.MineruTimeout > 0 {
			wait = cfg.MineruTimeout
		}
		if cfg.MaterialFulltextMaxBytes > 0 {
			fullMax = cfg.MaterialFulltextMaxBytes
		}
		if cfg.MaterialRAGTopK > 0 {
			topK = cfg.MaterialRAGTopK
		}
		if cfg.EmbeddingEndpointURL != "" {
			embedder = kb.NewEmbedder(cfg.EmbeddingEndpointURL, cfg.EmbeddingAPIKey, cfg.EmbeddingModel)
		}
		if cfg.RerankerEndpointURL != "" {
			reranker = kb.NewReranker(cfg.RerankerEndpointURL, cfg.RerankerAPIKey, cfg.RerankerModel)
		}
	}
	return &Pipeline{
		pool:      pool,
		client:    client,
		embedder:  embedder,
		reranker:  reranker,
		waitLimit: wait,
		fullMax:   fullMax,
		ragTopK:   topK,
		inflight:  make(map[string]chan struct{}),
	}
}

// Enabled reports whether MinerU credentials are present.
func (p *Pipeline) Enabled() bool {
	return p != nil && p.client != nil && p.client.Enabled()
}

// HasStore reports whether parse state persists (PostgreSQL available).
// Without a store, async parse would hit the live API with no state machine.
func (p *Pipeline) HasStore() bool {
	return p != nil && p.pool != nil
}

// FulltextMaxBytes returns the hybrid injection threshold.
func (p *Pipeline) FulltextMaxBytes() int {
	if p == nil || p.fullMax <= 0 {
		return 60 * 1024
	}
	return p.fullMax
}

// RAGTopK returns retrieval top-k.
func (p *Pipeline) RAGTopK() int {
	if p == nil || p.ragTopK <= 0 {
		return 8
	}
	return p.ragTopK
}

// ParseWaitLimit returns the bounded wait for pending/running MinerU parses.
func (p *Pipeline) ParseWaitLimit() time.Duration {
	if p == nil || p.waitLimit <= 0 {
		return defaultWait
	}
	return p.waitLimit
}

// GetParseState returns parse state for a blob (empty string if unknown).
func (p *Pipeline) GetParseState(ctx context.Context, blobSHA string) string {
	rec, err := p.getRecord(ctx, blobSHA)
	if err != nil || rec == nil {
		return ""
	}
	return rec.State
}

// ParseInfo is the user-facing parse status for one blob.
type ParseInfo struct {
	State        string `json:"state"`
	ErrorMessage string `json:"error_message,omitempty"`
}

// GetParseStates batch-loads parse states keyed by blob sha.
func (p *Pipeline) GetParseStates(ctx context.Context, blobSHAs []string) map[string]string {
	infos := p.GetParseInfos(ctx, blobSHAs)
	out := make(map[string]string, len(infos))
	for sha, info := range infos {
		out[sha] = info.State
	}
	return out
}

// GetParseInfos batch-loads parse state + error message keyed by blob sha.
func (p *Pipeline) GetParseInfos(ctx context.Context, blobSHAs []string) map[string]ParseInfo {
	out := make(map[string]ParseInfo, len(blobSHAs))
	if p == nil || p.pool == nil || len(blobSHAs) == 0 {
		return out
	}
	rows, err := p.pool.Query(ctx, `
		SELECT blob_sha256, state, COALESCE(error_message, '') FROM asset_parses
		WHERE engine = $1 AND blob_sha256 = ANY($2)
	`, engineMinerU, blobSHAs)
	if err != nil {
		return out
	}
	defer rows.Close()
	for rows.Next() {
		var sha, state, errMsg string
		if rows.Scan(&sha, &state, &errMsg) == nil {
			out[sha] = ParseInfo{State: state, ErrorMessage: errMsg}
		}
	}
	return out
}

// ResetForReparse rolls a failed/exhausted parse back to pending so the next
// EnqueueAsync/EnsureParsed cycle retries MinerU. Returns false when the blob
// is unknown or still in progress (pending/running/done).
func (p *Pipeline) ResetForReparse(ctx context.Context, blobSHA string) bool {
	if p == nil || p.pool == nil || blobSHA == "" {
		return false
	}
	tag, err := p.pool.Exec(ctx, `
		UPDATE asset_parses SET state = $2, attempts = 0, error_message = '', updated_at = now()
		WHERE blob_sha256 = $1 AND engine = $3 AND state NOT IN ($4, $5)
	`, blobSHA, statePending, engineMinerU, stateRunning, stateDone)
	if err != nil {
		return false
	}
	if tag.RowsAffected() > 0 {
		return true
	}
	// No row updated: either the record is done/running (must not reset) or it
	// never existed. When it never existed, allow a fresh parse to start.
	rec, err := p.getRecord(ctx, blobSHA)
	if err != nil || rec != nil {
		return false
	}
	return true
}

// EnsureParsed waits for or triggers MinerU parse. On miss/disable/fail returns
// fallback text from squad.ExtractTextWithAnchors (may be empty).
func (p *Pipeline) EnsureParsed(ctx context.Context, fileName string, blobSHA string, data []byte) (markdown string, fromMinerU bool, err error) {
	fallback := func() (string, bool, error) {
		txt, e := squad.ExtractTextWithAnchors(fileName, data)
		return strings.TrimSpace(txt), false, e
	}
	if p == nil {
		return fallback()
	}
	if blobSHA == "" && len(data) > 0 {
		sum := sha256.Sum256(data)
		blobSHA = hex.EncodeToString(sum[:])
	}
	if blobSHA == "" {
		return fallback()
	}

	rec, _ := p.getRecord(ctx, blobSHA)
	if rec != nil && rec.State == stateDone && strings.TrimSpace(rec.Markdown) != "" {
		return rec.Markdown, true, nil
	}
	if !p.Enabled() || !mineru.SupportsExt(fileName) {
		return fallback()
	}
	if rec != nil && rec.State == stateFailed && rec.Attempts >= maxAttempts {
		return fallback()
	}

	waitCtx, cancel := context.WithTimeout(ctx, p.waitLimit)
	defer cancel()

	if rec != nil && (rec.State == statePending || rec.State == stateRunning) {
		if done, md := p.waitUntilDone(waitCtx, blobSHA); done {
			return md, true, nil
		}
		return fallback()
	}

	if err := p.triggerParse(waitCtx, fileName, blobSHA, data); err != nil {
		slog.Warn("materialindex parse failed", "file", fileName, "sha", blobSHA[:min(12, len(blobSHA))], "error", err)
		return fallback()
	}
	rec, _ = p.getRecord(ctx, blobSHA)
	if rec != nil && rec.State == stateDone && strings.TrimSpace(rec.Markdown) != "" {
		return rec.Markdown, true, nil
	}
	return fallback()
}

// EnqueueAsync starts background parse without blocking the caller.
// Returns true when a parse job was actually enqueued.
func (p *Pipeline) EnqueueAsync(fileName string, blobSHA string, data []byte, projectID int64) bool {
	if p == nil || !p.Enabled() || !p.HasStore() || !mineru.SupportsExt(fileName) || blobSHA == "" || len(data) == 0 {
		return false
	}
	go func() {
		ctx, cancel := context.WithTimeout(context.Background(), p.waitLimit+2*time.Minute)
		defer cancel()
		md, ok, err := p.EnsureParsed(ctx, fileName, blobSHA, data)
		if err != nil || !ok || strings.TrimSpace(md) == "" {
			return
		}
		if idxErr := p.IndexChunks(ctx, projectID, blobSHA, fileName, md); idxErr != nil {
			slog.Warn("materialindex index chunks failed", "sha", blobSHA[:min(12, len(blobSHA))], "error", idxErr)
		}
	}()
	return true
}

func (p *Pipeline) triggerParse(ctx context.Context, fileName, blobSHA string, data []byte) error {
	key := blobSHA
	p.inflightMu.Lock()
	if ch, ok := p.inflight[key]; ok {
		p.inflightMu.Unlock()
		select {
		case <-ch:
		case <-ctx.Done():
			return ctx.Err()
		}
		return nil
	}
	done := make(chan struct{})
	p.inflight[key] = done
	p.inflightMu.Unlock()
	defer func() {
		close(done)
		p.inflightMu.Lock()
		delete(p.inflight, key)
		p.inflightMu.Unlock()
	}()

	if err := p.upsertPending(ctx, blobSHA, fileName); err != nil {
		return err
	}
	_ = p.setState(ctx, blobSHA, stateRunning, "", "", "", nil, 0)

	result, err := p.client.ParseFile(ctx, fileName, data)
	if err != nil {
		msg := err.Error()
		var retry *mineru.RetryableError
		attempts := 1
		if rec, _ := p.getRecord(ctx, blobSHA); rec != nil {
			attempts = rec.Attempts + 1
		}
		_ = p.bumpAttempts(ctx, blobSHA, attempts, msg)
		if errors.As(err, &retry) && attempts < maxAttempts {
			return err
		}
		_ = p.setState(ctx, blobSHA, stateFailed, "", "", msg, nil, 0)
		return err
	}
	return p.setState(ctx, blobSHA, stateDone, result.BatchID, result.Markdown, "", result.ContentList, result.TotalPages)
}

func (p *Pipeline) waitUntilDone(ctx context.Context, blobSHA string) (bool, string) {
	ticker := time.NewTicker(2 * time.Second)
	defer ticker.Stop()
	for {
		rec, _ := p.getRecord(ctx, blobSHA)
		if rec != nil {
			if rec.State == stateDone && strings.TrimSpace(rec.Markdown) != "" {
				return true, rec.Markdown
			}
			if rec.State == stateFailed {
				return false, ""
			}
		}
		select {
		case <-ctx.Done():
			return false, ""
		case <-ticker.C:
		}
	}
}

func (p *Pipeline) getRecord(ctx context.Context, blobSHA string) (*ParseRecord, error) {
	if p.pool == nil {
		return nil, nil
	}
	row := p.pool.QueryRow(ctx, `
		SELECT id, blob_sha256, file_name, engine, state, task_id, markdown, content_list, total_pages, error_message, attempts
		FROM asset_parses WHERE blob_sha256 = $1 AND engine = $2
	`, blobSHA, engineMinerU)
	var rec ParseRecord
	var contentList []byte
	err := row.Scan(&rec.ID, &rec.BlobSHA256, &rec.FileName, &rec.Engine, &rec.State, &rec.TaskID,
		&rec.Markdown, &contentList, &rec.TotalPages, &rec.ErrorMessage, &rec.Attempts)
	if errors.Is(err, pgx.ErrNoRows) {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}
	rec.ContentList = contentList
	return &rec, nil
}

func (p *Pipeline) upsertPending(ctx context.Context, blobSHA, fileName string) error {
	if p.pool == nil {
		return nil
	}
	_, err := p.pool.Exec(ctx, `
		INSERT INTO asset_parses (blob_sha256, file_name, engine, state, attempts, updated_at)
		VALUES ($1, $2, $3, $4, 0, now())
		ON CONFLICT (blob_sha256, engine) DO UPDATE SET
			file_name = EXCLUDED.file_name,
			updated_at = now()
		WHERE asset_parses.state NOT IN ('done', 'running')
	`, blobSHA, fileName, engineMinerU, statePending)
	return err
}

func (p *Pipeline) setState(ctx context.Context, blobSHA, state, taskID, markdown, errMsg string, contentList json.RawMessage, pages int) error {
	if p.pool == nil {
		return nil
	}
	cl := []byte("[]")
	if len(contentList) > 0 {
		cl = contentList
	}
	_, err := p.pool.Exec(ctx, `
		UPDATE asset_parses SET
			state = $2,
			task_id = CASE WHEN $3 = '' THEN task_id ELSE $3 END,
			markdown = CASE WHEN $4 = '' THEN markdown ELSE $4 END,
			content_list = CASE WHEN $5::jsonb = '[]'::jsonb THEN content_list ELSE $5::jsonb END,
			total_pages = CASE WHEN $6 = 0 THEN total_pages ELSE $6 END,
			error_message = $7,
			updated_at = now()
		WHERE blob_sha256 = $1 AND engine = $8
	`, blobSHA, state, taskID, markdown, cl, pages, errMsg, engineMinerU)
	return err
}

func (p *Pipeline) bumpAttempts(ctx context.Context, blobSHA string, attempts int, errMsg string) error {
	if p.pool == nil {
		return nil
	}
	_, err := p.pool.Exec(ctx, `
		UPDATE asset_parses SET attempts = $2, error_message = $3, updated_at = now()
		WHERE blob_sha256 = $1 AND engine = $4
	`, blobSHA, attempts, errMsg, engineMinerU)
	return err
}

// IndexChunks splits markdown and stores embeddings into material_chunks.
func (p *Pipeline) IndexChunks(ctx context.Context, projectID int64, blobSHA, fileName, markdown string) error {
	if p == nil || p.pool == nil || p.embedder == nil {
		return nil
	}
	sections := splitMaterialSections(markdown)
	if len(sections) == 0 {
		return nil
	}
	texts := make([]string, len(sections))
	for i, s := range sections {
		texts[i] = fmt.Sprintf("%s %s\n%s", fileName, s.Title, s.Content)
	}
	vecs, err := p.embedder.Embed(ctx, texts)
	if err != nil {
		return err
	}
	for i, sec := range sections {
		var vec []float32
		if i < len(vecs) {
			vec = vecs[i]
		}
		hash := fmt.Sprintf("%x", sha256.Sum256([]byte(fmt.Sprintf("%s:%d:%s:%s", blobSHA, i, sec.Title, sec.Content))))
		meta, _ := json.Marshal(map[string]any{
			"file_name":  fileName,
			"page_start": sec.PageStart,
			"page_end":   sec.PageEnd,
		})
		vecLit := kb.VectorLiteral(vec)
		_, err := p.pool.Exec(ctx, `
			INSERT INTO material_chunks (project_id, blob_sha256, file_name, chunk_index, title, content, page_start, page_end, chunk_hash, embedding, metadata)
			VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10::vector, $11)
			ON CONFLICT (chunk_hash) DO UPDATE SET
				title = EXCLUDED.title,
				content = EXCLUDED.content,
				embedding = EXCLUDED.embedding,
				metadata = EXCLUDED.metadata
		`, projectID, blobSHA, fileName, i, sec.Title, sec.Content, sec.PageStart, sec.PageEnd, hash, vecLit, meta)
		if err != nil {
			return err
		}
	}
	return nil
}

// Search retrieves top chunks for a project+optional blob, with optional rerank.
func (p *Pipeline) Search(ctx context.Context, projectID int64, blobSHA, query string, limit int) ([]ChunkHit, error) {
	if p == nil || p.pool == nil || p.embedder == nil || strings.TrimSpace(query) == "" {
		return nil, nil
	}
	if limit <= 0 {
		limit = p.RAGTopK()
	}
	qVecs, err := p.embedder.Embed(ctx, []string{query})
	if err != nil || len(qVecs) == 0 {
		return nil, err
	}
	vecLit := kb.VectorLiteral(qVecs[0])
	fetch := limit * 3
	if fetch < limit {
		fetch = limit
	}
	rows, err := p.pool.Query(ctx, `
		SELECT file_name, title, content, page_start, page_end,
		       1 - (embedding <=> $1::vector) AS score
		FROM material_chunks
		WHERE project_id = $2
		  AND ($3 = '' OR blob_sha256 = $3)
		ORDER BY embedding <=> $1::vector ASC
		LIMIT $4
	`, vecLit, projectID, blobSHA, fetch)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var hits []ChunkHit
	var docs []string
	for rows.Next() {
		var h ChunkHit
		if err := rows.Scan(&h.FileName, &h.Title, &h.Content, &h.PageStart, &h.PageEnd, &h.Score); err != nil {
			continue
		}
		hits = append(hits, h)
		docs = append(docs, h.Title+"\n"+h.Content)
	}
	if p.reranker != nil && len(docs) > 0 {
		scores, rerr := p.reranker.Rerank(ctx, query, docs)
		if rerr == nil && len(scores) == len(hits) {
			for i := range hits {
				hits[i].Score = scores[i]
			}
			sortHitsByScore(hits)
		}
	}
	if len(hits) > limit {
		hits = hits[:limit]
	}
	return hits, nil
}

// BuildTOC extracts heading outline with page hints from markdown.
func BuildTOC(fileName, markdown string) string {
	var sb strings.Builder
	sb.WriteString(fmt.Sprintf("# 材料目录: %s\n\n", filepath.Base(fileName)))
	sb.WriteString("本文档体积较大，已改为目录 + 按任务检索片段注入。请结合 `.rag.md` 引用出处。\n\n")
	currentPage := 1
	count := 0
	for _, line := range strings.Split(markdown, "\n") {
		trim := strings.TrimSpace(line)
		if m := pageAnchorRe.FindStringSubmatch(trim); len(m) == 2 {
			fmt.Sscanf(m[1], "%d", &currentPage)
			continue
		}
		if strings.HasPrefix(trim, "#") {
			level := 0
			for level < len(trim) && trim[level] == '#' {
				level++
			}
			title := strings.TrimSpace(trim[level:])
			if title == "" {
				continue
			}
			sb.WriteString(fmt.Sprintf("- [p.%d] %s %s\n", currentPage, strings.Repeat("#", min(level, 3)), title))
			count++
			if count >= 80 {
				sb.WriteString("\n…（目录已截断）\n")
				break
			}
		}
	}
	if count == 0 {
		sb.WriteString("- （未识别到标题；请直接查看检索片段）\n")
	}
	return sb.String()
}

// FormatRAGMarkdown formats retrieved hits for workspace injection.
func FormatRAGMarkdown(fileName string, hits []ChunkHit) string {
	var sb strings.Builder
	sb.WriteString(fmt.Sprintf("# 材料检索片段: %s\n\n", filepath.Base(fileName)))
	if len(hits) == 0 {
		sb.WriteString("（未命中相关片段；可缩小问题范围或上传更小文档）\n")
		return sb.String()
	}
	for i, h := range hits {
		title := h.Title
		if title == "" {
			title = "片段"
		}
		page := h.PageStart
		if page <= 0 {
			page = 1
		}
		cite := fmt.Sprintf("〔材料: %s §%s p.%d〕", filepath.Base(fileName), title, page)
		sb.WriteString(fmt.Sprintf("## %d. %s\n\n", i+1, title))
		sb.WriteString(fmt.Sprintf("- 出处: %s\n", cite))
		if h.PageEnd > h.PageStart && h.PageEnd > 0 {
			sb.WriteString(fmt.Sprintf("- 页码范围: p.%d–p.%d\n", h.PageStart, h.PageEnd))
		}
		sb.WriteString("\n")
		sb.WriteString(strings.TrimSpace(h.Content))
		sb.WriteString("\n\n---\n\n")
	}
	return sb.String()
}

type section struct {
	Title     string
	Content   string
	PageStart int
	PageEnd   int
}

func splitMaterialSections(md string) []section {
	lines := strings.Split(md, "\n")
	var sections []section
	var cur section
	cur.PageStart = 1
	cur.PageEnd = 1
	currentPage := 1
	flush := func() {
		cur.Content = strings.TrimSpace(cur.Content)
		if cur.Content == "" && cur.Title == "" {
			return
		}
		if cur.Title == "" {
			cur.Title = "正文"
		}
		cur.PageEnd = currentPage
		// further split oversized sections
		parts := splitByRuneLimit(cur, chunkRuneLimit)
		sections = append(sections, parts...)
		cur = section{PageStart: currentPage, PageEnd: currentPage}
	}
	for _, line := range lines {
		trim := strings.TrimSpace(line)
		if m := pageAnchorRe.FindStringSubmatch(trim); len(m) == 2 {
			fmt.Sscanf(m[1], "%d", &currentPage)
			if cur.PageStart == 0 {
				cur.PageStart = currentPage
			}
			cur.Content += line + "\n"
			continue
		}
		if strings.HasPrefix(trim, "## ") {
			flush()
			cur.Title = strings.TrimSpace(strings.TrimPrefix(trim, "## "))
			cur.PageStart = currentPage
			continue
		}
		cur.Content += line + "\n"
	}
	flush()
	return sections
}

func splitByRuneLimit(sec section, limit int) []section {
	if utf8.RuneCountInString(sec.Content) <= limit {
		return []section{sec}
	}
	paras := strings.Split(sec.Content, "\n\n")
	var out []section
	var buf strings.Builder
	part := 1
	flush := func() {
		if buf.Len() == 0 {
			return
		}
		title := sec.Title
		if part > 1 {
			title = fmt.Sprintf("%s (%d)", sec.Title, part)
		}
		out = append(out, section{
			Title:     title,
			Content:   strings.TrimSpace(buf.String()),
			PageStart: sec.PageStart,
			PageEnd:   sec.PageEnd,
		})
		buf.Reset()
		part++
	}
	for _, p := range paras {
		if utf8.RuneCountInString(buf.String())+utf8.RuneCountInString(p) > limit && buf.Len() > 0 {
			flush()
		}
		if buf.Len() > 0 {
			buf.WriteString("\n\n")
		}
		buf.WriteString(p)
	}
	flush()
	if len(out) == 0 {
		return []section{sec}
	}
	return out
}

func sortHitsByScore(hits []ChunkHit) {
	for i := 0; i < len(hits); i++ {
		for j := i + 1; j < len(hits); j++ {
			if hits[j].Score > hits[i].Score {
				hits[i], hits[j] = hits[j], hits[i]
			}
		}
	}
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}