Agent-System/internal/dialoguectx/collector.go

package dialoguectx

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"time"

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

	"chat-agent-system/internal/ctxmgr"
)

// Collector loads project dialogue and enriches speaker metadata.
type Collector struct {
	Pool *pgxpool.Pool
	// MemoryFallback is used when Pool is nil (tests / in-memory).
	MemoryFallback *ctxmgr.ContextManager
}

// CollectAndPackage builds a layered dialogue package for a project.
func (c *Collector) CollectAndPackage(ctx context.Context, opts CollectOptions) (*Package, error) {
	now := opts.Now
	if now.IsZero() {
		now = time.Now()
	}
	pkg := &Package{
		AssembledAt:      now,
		InjectionEnabled: InjectionEnabled(),
		Facts:            []DialogueFact{},
	}
	if !pkg.InjectionEnabled || opts.ProjectID <= 0 {
		return pkg, nil
	}

	maxDays := opts.MaxDays
	if maxDays <= 0 {
		maxDays = MaxDaysFromEnv()
	}
	maxTurns := opts.MaxTurns
	if maxTurns <= 0 {
		maxTurns = MaxTurnsFromEnv()
	}

	turns, excludedPrivate, excludedNoAccess, excludedNoise, err := c.collectTurns(ctx, opts, maxDays, maxTurns)
	if err != nil {
		return nil, err
	}
	pkg.ExcludedPrivate = excludedPrivate
	pkg.ExcludedNoAccess = excludedNoAccess
	pkg.ExcludedNoise = excludedNoise
	pkg.TurnCount = len(turns)
	if len(turns) > 0 {
		pkg.CoveredFrom = turns[0].CreatedAt
		pkg.CoveredTo = turns[len(turns)-1].CreatedAt
	}

	pkg.TurnsMD = renderTurnsMD(turns, pkg.CoveredFrom, pkg.CoveredTo)
	pkg.DigestMD = renderDigestMD(turns, pkg)
	pkg.Facts = extractFacts(turns)
	return pkg, nil
}

type rawRow struct {
	rec         ctxmgr.ChatMessageRecord
	speakerName string
	systemRole  string
	projectRole string
	docDuty     string
}

func (c *Collector) collectTurns(ctx context.Context, opts CollectOptions, maxDays, maxTurns int) ([]RichDialogueTurn, int, int, int, error) {
	cutoff := time.Now().AddDate(0, 0, -maxDays)
	if !opts.Now.IsZero() {
		cutoff = opts.Now.AddDate(0, 0, -maxDays)
	}

	var rows []rawRow
	var err error
	if c != nil && c.Pool != nil {
		rows, err = c.queryPG(ctx, opts.ProjectID, cutoff, maxTurns*3)
	} else if c != nil && c.MemoryFallback != nil {
		rows, err = c.queryMemory(ctx, opts.ProjectID, cutoff, maxTurns*3)
	}
	if err != nil {
		return nil, 0, 0, 0, err
	}

	var (
		out              []RichDialogueTurn
		excludedPrivate  int
		excludedNoAccess int
		excludedNoise    int
		seen             = map[string]struct{}{}
	)
	for _, row := range rows {
		stageKey := row.rec.StageKey
		if stageKey == "" {
			stageKey = stageKeyFromMeta(row.rec.Metadata)
		}
		cr := cleanMessage(row.rec.Role, row.rec.Content, row.rec.ContentKind, row.rec.ExportPolicy, stageKey, opts.AllowedStageKey, row.rec.Metadata)
		if cr.Private {
			excludedPrivate++
			continue
		}
		if cr.NoAccess {
			excludedNoAccess++
			continue
		}
		if cr.Noise || !cr.Include {
			excludedNoise++
			continue
		}
		fp := contentFingerprint(row.rec.Role, cr.Content)
		if _, ok := seen[fp]; ok {
			excludedNoise++
			continue
		}
		seen[fp] = struct{}{}

		name := strings.TrimSpace(row.speakerName)
		if name == "" {
			if row.rec.Role == "assistant" {
				name = "助手"
			} else {
				name = fmt.Sprintf("用户#%d", row.rec.UserID)
			}
		}
		out = append(out, RichDialogueTurn{
			MessageID:   row.rec.ID,
			SessionID:   row.rec.SessionID,
			SpeakerID:   row.rec.UserID,
			SpeakerName: name,
			SystemRole:  row.systemRole,
			ProjectRole: row.projectRole,
			DocDuty:     row.docDuty,
			Role:        row.rec.Role,
			Content:     truncateRunes(cr.Content, MaxTurnContentRunes),
			StageKey:    stageKey,
			DayKey:      row.rec.DayKey,
			Category:    row.rec.Category,
			ContentKind: row.rec.ContentKind,
			CreatedAt:   row.rec.CreatedAt,
		})
	}

	if len(out) > maxTurns {
		out = out[len(out)-maxTurns:]
	}
	return out, excludedPrivate, excludedNoAccess, excludedNoise, nil
}

func (c *Collector) queryPG(ctx context.Context, projectID int64, cutoff time.Time, limit int) ([]rawRow, error) {
	q := `
		SELECT cm.id, cm.session_id, COALESCE(cm.user_id, 0), cm.role, cm.content,
		       cm.intent, cm.branch, cm.metadata, COALESCE(cm.day_key::text, ''),
		       COALESCE(cm.category, ''), COALESCE(cm.content_kind, ''),
		       COALESCE(cm.export_policy, 'none'),
		       COALESCE(cm.stage_key, ''), cm.created_at,
		       COALESCE(u.name, ''), COALESCE(u.role, ''),
		       COALESCE(pm.project_role, ''), COALESCE(pm.document_duty, '')
		FROM chat_messages cm
		JOIN chat_sessions cs ON cs.session_id = cm.session_id
		LEFT JOIN users u ON u.id = cm.user_id
		LEFT JOIN project_members pm ON pm.project_id = cs.project_id AND pm.user_id = cm.user_id AND pm.status = 'active'
		WHERE cs.project_id = $1 AND cm.deleted_at IS NULL AND cm.created_at >= $2
		ORDER BY cm.created_at ASC
		LIMIT $3
	`
	pgRows, err := c.Pool.Query(ctx, q, projectID, cutoff, limit)
	if err != nil {
		// Fallback without document_duty / stage_key columns if schema older.
		return c.queryPGLegacy(ctx, projectID, cutoff, limit)
	}
	defer pgRows.Close()

	var out []rawRow
	for pgRows.Next() {
		var rec ctxmgr.ChatMessageRecord
		var metaJSON []byte
		var name, sysRole, pRole, duty string
		if err := pgRows.Scan(
			&rec.ID, &rec.SessionID, &rec.UserID, &rec.Role, &rec.Content,
			&rec.Intent, &rec.Branch, &metaJSON, &rec.DayKey, &rec.Category,
			&rec.ContentKind, &rec.ExportPolicy, &rec.StageKey, &rec.CreatedAt,
			&name, &sysRole, &pRole, &duty,
		); err != nil {
			return nil, err
		}
		if len(metaJSON) > 0 {
			_ = json.Unmarshal(metaJSON, &rec.Metadata)
		}
		out = append(out, rawRow{rec: rec, speakerName: name, systemRole: sysRole, projectRole: pRole, docDuty: duty})
	}
	return out, pgRows.Err()
}

func (c *Collector) queryPGLegacy(ctx context.Context, projectID int64, cutoff time.Time, limit int) ([]rawRow, error) {
	q := `
		SELECT cm.id, cm.session_id, COALESCE(cm.user_id, 0), cm.role, cm.content,
		       cm.intent, cm.branch, cm.metadata, COALESCE(cm.day_key::text, ''),
		       COALESCE(cm.category, ''), COALESCE(cm.content_kind, ''),
		       COALESCE(cm.export_policy, 'none'), cm.created_at,
		       COALESCE(u.name, ''), COALESCE(u.role, ''),
		       COALESCE(pm.project_role, '')
		FROM chat_messages cm
		JOIN chat_sessions cs ON cs.session_id = cm.session_id
		LEFT JOIN users u ON u.id = cm.user_id
		LEFT JOIN project_members pm ON pm.project_id = cs.project_id AND pm.user_id = cm.user_id AND pm.status = 'active'
		WHERE cs.project_id = $1 AND cm.deleted_at IS NULL AND cm.created_at >= $2
		ORDER BY cm.created_at ASC
		LIMIT $3
	`
	pgRows, err := c.Pool.Query(ctx, q, projectID, cutoff, limit)
	if err != nil {
		return nil, fmt.Errorf("dialoguectx collect: %w", err)
	}
	defer pgRows.Close()
	var out []rawRow
	for pgRows.Next() {
		var rec ctxmgr.ChatMessageRecord
		var metaJSON []byte
		var name, sysRole, pRole string
		if err := pgRows.Scan(
			&rec.ID, &rec.SessionID, &rec.UserID, &rec.Role, &rec.Content,
			&rec.Intent, &rec.Branch, &metaJSON, &rec.DayKey, &rec.Category,
			&rec.ContentKind, &rec.ExportPolicy, &rec.CreatedAt,
			&name, &sysRole, &pRole,
		); err != nil {
			return nil, err
		}
		if len(metaJSON) > 0 {
			_ = json.Unmarshal(metaJSON, &rec.Metadata)
		}
		out = append(out, rawRow{rec: rec, speakerName: name, systemRole: sysRole, projectRole: pRole})
	}
	return out, pgRows.Err()
}

func (c *Collector) queryMemory(ctx context.Context, projectID int64, cutoff time.Time, limit int) ([]rawRow, error) {
	if c == nil || c.MemoryFallback == nil {
		return nil, nil
	}
	recs, err := c.MemoryFallback.CollectProjectDialogueRecords(ctx, projectID, limit)
	if err != nil {
		return nil, err
	}
	var out []rawRow
	for _, rec := range recs {
		if !rec.CreatedAt.IsZero() && rec.CreatedAt.Before(cutoff) {
			continue
		}
		out = append(out, rawRow{rec: rec, speakerName: "", systemRole: "", projectRole: ""})
	}
	return out, nil
}

// CollectFromRecords packages pre-loaded records (tests / memory mode).
func CollectFromRecords(recs []ctxmgr.ChatMessageRecord, opts CollectOptions, enrich map[int64]speakerMeta) (*Package, error) {
	now := opts.Now
	if now.IsZero() {
		now = time.Now()
	}
	pkg := &Package{AssembledAt: now, InjectionEnabled: true, Facts: []DialogueFact{}}
	maxDays := opts.MaxDays
	if maxDays <= 0 {
		maxDays = DefaultMaxDays
	}
	maxTurns := opts.MaxTurns
	if maxTurns <= 0 {
		maxTurns = DefaultMaxTurns
	}
	cutoff := now.AddDate(0, 0, -maxDays)

	var (
		turns            []RichDialogueTurn
		excludedPrivate  int
		excludedNoAccess int
		excludedNoise    int
		seen             = map[string]struct{}{}
	)
	for _, rec := range recs {
		if !rec.CreatedAt.IsZero() && rec.CreatedAt.Before(cutoff) {
			continue
		}
		stageKey := rec.StageKey
		if stageKey == "" {
			stageKey = stageKeyFromMeta(rec.Metadata)
		}
		cr := cleanMessage(rec.Role, rec.Content, rec.ContentKind, rec.ExportPolicy, stageKey, opts.AllowedStageKey, rec.Metadata)
		if cr.Private {
			excludedPrivate++
			continue
		}
		if cr.NoAccess {
			excludedNoAccess++
			continue
		}
		if cr.Noise || !cr.Include {
			excludedNoise++
			continue
		}
		fp := contentFingerprint(rec.Role, cr.Content)
		if _, ok := seen[fp]; ok {
			excludedNoise++
			continue
		}
		seen[fp] = struct{}{}
		meta := enrich[rec.UserID]
		name := meta.Name
		if name == "" {
			if rec.Role == "assistant" {
				name = "助手"
			} else {
				name = fmt.Sprintf("用户#%d", rec.UserID)
			}
		}
		turns = append(turns, RichDialogueTurn{
			MessageID:   rec.ID,
			SessionID:   rec.SessionID,
			SpeakerID:   rec.UserID,
			SpeakerName: name,
			SystemRole:  meta.SystemRole,
			ProjectRole: meta.ProjectRole,
			DocDuty:     meta.DocDuty,
			Role:        rec.Role,
			Content:     truncateRunes(cr.Content, MaxTurnContentRunes),
			StageKey:    stageKey,
			DayKey:      rec.DayKey,
			Category:    rec.Category,
			ContentKind: rec.ContentKind,
			CreatedAt:   rec.CreatedAt,
		})
	}
	if len(turns) > maxTurns {
		turns = turns[len(turns)-maxTurns:]
	}
	pkg.ExcludedPrivate = excludedPrivate
	pkg.ExcludedNoAccess = excludedNoAccess
	pkg.ExcludedNoise = excludedNoise
	pkg.TurnCount = len(turns)
	if len(turns) > 0 {
		pkg.CoveredFrom = turns[0].CreatedAt
		pkg.CoveredTo = turns[len(turns)-1].CreatedAt
	}
	pkg.TurnsMD = renderTurnsMD(turns, pkg.CoveredFrom, pkg.CoveredTo)
	pkg.DigestMD = renderDigestMD(turns, pkg)
	pkg.Facts = extractFacts(turns)
	return pkg, nil
}

type speakerMeta struct {
	Name        string
	SystemRole  string
	ProjectRole string
	DocDuty     string
}