Agent-System/internal/storage/cas.go

package storage

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"math"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

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

// Embedder 向量生成抽象接口
type Embedder interface {
	EmbedOne(ctx context.Context, text string) ([]float32, error)
}

// CAS 内容寻址存储与切片防爆炸去重管理器
type CAS struct {
	baseDir  string
	pool     *pgxpool.Pool
	embedder Embedder

	// 离线/降级内存索引 (当数据库未连接时)
	mu           sync.RWMutex
	memoryBlobs  map[string]*BlobInfo
	memoryRefs   map[int64]*FileReference
	memoryChunks []*MemoryChunk
	nextRefID    int64
	nextChunkID  int64
}

// NewCAS 创建 CAS 存储管理器
func NewCAS(baseDir string, pool *pgxpool.Pool, embedder Embedder) (*CAS, error) {
	if strings.TrimSpace(baseDir) == "" {
		return nil, errors.New("storage: baseDir is required (runtime roots come from config, not package defaults)")
	}
	blobDir := filepath.Join(baseDir, "blobs")
	if err := os.MkdirAll(blobDir, 0755); err != nil {
		return nil, fmt.Errorf("failed to create blob dir: %w", err)
	}

	return &CAS{
		baseDir:      baseDir,
		pool:         pool,
		embedder:     embedder,
		memoryBlobs:  make(map[string]*BlobInfo),
		memoryRefs:   make(map[int64]*FileReference),
		memoryChunks: make([]*MemoryChunk, 0),
		nextRefID:    1,
		nextChunkID:  1,
	}, nil
}

// StoreFile 写入物理文件并进行 CAS Hash 去重
func (c *CAS) StoreFile(ctx context.Context, projectID, userID int64, fileName, fileCategory string, data []byte, mimeType string) (*FileReference, error) {
	if len(data) == 0 {
		return nil, errors.New("cannot store empty file")
	}
	if fileCategory == "" {
		fileCategory = "material"
	}
	if mimeType == "" {
		mimeType = "application/octet-stream"
	}

	hashBytes := sha256.Sum256(data)
	hashStr := hex.EncodeToString(hashBytes[:])
	sizeBytes := int64(len(data))

	relPath := fmt.Sprintf("blobs/%s/%s/%s.bin", hashStr[:2], hashStr[2:4], hashStr)
	fullPath := filepath.Join(c.baseDir, relPath)

	if c.pool != nil {
		return c.storeFileDB(ctx, projectID, userID, fileName, fileCategory, mimeType, hashStr, relPath, fullPath, sizeBytes, data)
	}

	return c.storeFileMemory(projectID, userID, fileName, fileCategory, mimeType, hashStr, relPath, fullPath, sizeBytes, data)
}

func (c *CAS) storeFileDB(ctx context.Context, projectID, userID int64, fileName, fileCategory, mimeType, hashStr, relPath, fullPath string, sizeBytes int64, data []byte) (*FileReference, error) {
	tx, err := c.pool.Begin(ctx)
	if err != nil {
		return nil, fmt.Errorf("begin tx failed: %w", err)
	}
	defer tx.Rollback(ctx)

	var existingHash string
	var refCount int
	err = tx.QueryRow(ctx, "SELECT sha256, ref_count FROM storage_blobs WHERE sha256 = $1 FOR UPDATE", hashStr).Scan(&existingHash, &refCount)

	isDedup := false
	if err == nil && existingHash != "" {
		// 命中去重：物理文件已存在，递增引用计数
		isDedup = true
		_, err = tx.Exec(ctx, "UPDATE storage_blobs SET ref_count = ref_count + 1 WHERE sha256 = $1", hashStr)
		if err != nil {
			return nil, fmt.Errorf("update ref_count failed: %w", err)
		}
	} else {
		// 新文件：写入物理磁盘 (采用原子写入)
		dir := filepath.Dir(fullPath)
		if err := os.MkdirAll(dir, 0755); err != nil {
			return nil, fmt.Errorf("mkdir failed: %w", err)
		}
		tmpPath := fullPath + ".tmp"
		if err := os.WriteFile(tmpPath, data, 0644); err != nil {
			return nil, fmt.Errorf("write tmp file failed: %w", err)
		}
		if err := os.Rename(tmpPath, fullPath); err != nil {
			return nil, fmt.Errorf("rename file failed: %w", err)
		}

		_, err = tx.Exec(ctx,
			"INSERT INTO storage_blobs (sha256, size_bytes, mime_type, storage_path, ref_count, created_at) VALUES ($1, $2, $3, $4, 1, now())",
			hashStr, sizeBytes, mimeType, relPath)
		if err != nil {
			return nil, fmt.Errorf("insert blob failed: %w", err)
		}
	}

	// 插入引用记录
	var refID int64
	var createdAt time.Time
	err = tx.QueryRow(ctx, `
		INSERT INTO file_references (project_id, user_id, blob_sha256, file_name, file_category, summary, metadata, created_at)
		VALUES ($1, $2, $3, $4, $5, '', '{}'::jsonb, now())
		RETURNING id, created_at
	`, projectID, userID, hashStr, fileName, fileCategory).Scan(&refID, &createdAt)
	if err != nil {
		return nil, fmt.Errorf("insert file reference failed: %w", err)
	}

	if err := tx.Commit(ctx); err != nil {
		return nil, fmt.Errorf("commit tx failed: %w", err)
	}

	return &FileReference{
		ID:             refID,
		ProjectID:      projectID,
		UserID:         userID,
		BlobSHA256:     hashStr,
		FileName:       fileName,
		FileCategory:   fileCategory,
		CreatedAt:      createdAt,
		IsDeduplicated: isDedup,
	}, nil
}

func (c *CAS) storeFileMemory(projectID, userID int64, fileName, fileCategory, mimeType, hashStr, relPath, fullPath string, sizeBytes int64, data []byte) (*FileReference, error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	isDedup := false
	blob, exists := c.memoryBlobs[hashStr]
	if exists {
		isDedup = true
		blob.RefCount++
	} else {
		dir := filepath.Dir(fullPath)
		if err := os.MkdirAll(dir, 0755); err != nil {
			return nil, fmt.Errorf("mkdir failed: %w", err)
		}
		if err := os.WriteFile(fullPath, data, 0644); err != nil {
			return nil, fmt.Errorf("write file failed: %w", err)
		}
		c.memoryBlobs[hashStr] = &BlobInfo{
			SHA256:      hashStr,
			SizeBytes:   sizeBytes,
			MimeType:    mimeType,
			StoragePath: relPath,
			RefCount:    1,
			CreatedAt:   time.Now(),
		}
	}

	refID := c.nextRefID
	c.nextRefID++

	ref := &FileReference{
		ID:             refID,
		ProjectID:      projectID,
		UserID:         userID,
		BlobSHA256:     hashStr,
		FileName:       fileName,
		FileCategory:   fileCategory,
		CreatedAt:      time.Now(),
		IsDeduplicated: isDedup,
	}
	c.memoryRefs[refID] = ref
	return ref, nil
}

// OpenFileByRef 根据文件引用 ID 打开物理文件流
func (c *CAS) OpenFileByRef(ctx context.Context, refID int64) (io.ReadCloser, *FileReference, error) {
	if c.pool != nil {
		var ref FileReference
		var storagePath string
		err := c.pool.QueryRow(ctx, `
			SELECT r.id, r.project_id, r.user_id, r.blob_sha256, r.file_name, r.file_category, r.summary, r.created_at, b.storage_path
			FROM file_references r
			JOIN storage_blobs b ON r.blob_sha256 = b.sha256
			WHERE r.id = $1
		`, refID).Scan(&ref.ID, &ref.ProjectID, &ref.UserID, &ref.BlobSHA256, &ref.FileName, &ref.FileCategory, &ref.Summary, &ref.CreatedAt, &storagePath)
		if err != nil {
			return nil, nil, fmt.Errorf("file reference not found: %w", err)
		}

		fullPath := filepath.Join(c.baseDir, storagePath)
		f, err := os.Open(fullPath)
		if err != nil {
			return nil, nil, fmt.Errorf("open physical file failed: %w", err)
		}
		return f, &ref, nil
	}

	c.mu.RLock()
	defer c.mu.RUnlock()
	ref, ok := c.memoryRefs[refID]
	if !ok {
		return nil, nil, errors.New("file reference not found in memory")
	}
	blob, ok := c.memoryBlobs[ref.BlobSHA256]
	if !ok {
		return nil, nil, errors.New("blob not found in memory")
	}
	fullPath := filepath.Join(c.baseDir, blob.StoragePath)
	f, err := os.Open(fullPath)
	if err != nil {
		return nil, nil, fmt.Errorf("open physical file failed: %w", err)
	}
	return f, ref, nil
}

// StoreMemoryChunk 存储记忆切片并自动去重向量化计算
func (c *CAS) StoreMemoryChunk(ctx context.Context, in MemoryChunkInput) (*MemoryChunk, error) {
	if strings.TrimSpace(in.Content) == "" {
		return nil, errors.New("chunk content cannot be empty")
	}
	if in.ImportanceWeight <= 0 {
		in.ImportanceWeight = 1.00
	}

	normalized := strings.TrimSpace(in.Title) + "\n" + strings.TrimSpace(in.Content)
	h := sha256.Sum256([]byte(normalized))
	chunkHash := hex.EncodeToString(h[:])

	var embedding []float32

	// 1. 检查是否存在相同切片 Hash，复用向量结果以避免重复调用大模型 API 与浪费存储
	if c.pool != nil {
		var vecStr *string
		err := c.pool.QueryRow(ctx, "SELECT embedding::text FROM context_memory_chunks WHERE chunk_hash = $1 AND embedding IS NOT NULL LIMIT 1", chunkHash).Scan(&vecStr)
		if err == nil && vecStr != nil {
			embedding = parseVectorText(*vecStr)
		}
	}

	// 2. 若未命中现有切片，且配置了 Embedder，则调用生成向量
	if len(embedding) == 0 && c.embedder != nil {
		if vec, err := c.embedder.EmbedOne(ctx, in.Content); err == nil && len(vec) > 0 {
			embedding = vec
		}
	}

	metaJSON, _ := json.Marshal(in.Metadata)
	if len(metaJSON) == 0 {
		metaJSON = []byte("{}")
	}

	if c.pool != nil {
		var chunkID int64
		var createdAt time.Time

		var vecLiteral any
		if len(embedding) > 0 {
			vecLiteral = vectorLiteral(embedding)
		}

		err := c.pool.QueryRow(ctx, `
			INSERT INTO context_memory_chunks
				(domain, project_id, user_id, reference_id, title, content, chunk_hash, embedding, tsv, importance_weight, metadata, created_at)
			VALUES
				($1, $2, $3, $4, $5, $6, $7, $8::vector, to_tsvector('simple', $5 || ' ' || $6), $9, $10, now())
			RETURNING id, created_at
		`, in.Domain, in.ProjectID, in.UserID, in.ReferenceID, in.Title, in.Content, chunkHash, vecLiteral, in.ImportanceWeight, metaJSON).Scan(&chunkID, &createdAt)

		if err != nil {
			return nil, fmt.Errorf("insert context memory chunk failed: %w", err)
		}

		return &MemoryChunk{
			ID:               chunkID,
			Domain:           in.Domain,
			ProjectID:        in.ProjectID,
			UserID:           in.UserID,
			ReferenceID:      in.ReferenceID,
			Title:            in.Title,
			Content:          in.Content,
			ChunkHash:        chunkHash,
			Embedding:        embedding,
			ImportanceWeight: in.ImportanceWeight,
			Metadata:         in.Metadata,
			CreatedAt:        createdAt,
		}, nil
	}

	// 降级内存存储
	c.mu.Lock()
	defer c.mu.Unlock()
	chunkID := c.nextChunkID
	c.nextChunkID++
	chunk := &MemoryChunk{
		ID:               chunkID,
		Domain:           in.Domain,
		ProjectID:        in.ProjectID,
		UserID:           in.UserID,
		ReferenceID:      in.ReferenceID,
		Title:            in.Title,
		Content:          in.Content,
		ChunkHash:        chunkHash,
		Embedding:        embedding,
		ImportanceWeight: in.ImportanceWeight,
		Metadata:         in.Metadata,
		CreatedAt:        time.Now(),
	}
	c.memoryChunks = append(c.memoryChunks, chunk)
	return chunk, nil
}

// SearchMemory 统一多维度记忆语义检索 (向量相似度 + 全文检索 + 时间衰减权重)
func (c *CAS) SearchMemory(ctx context.Context, opts MemorySearchOptions) ([]*MemoryChunkHit, error) {
	if opts.TopK <= 0 {
		opts.TopK = 5
	}
	if opts.DecayHalfLifeDays <= 0 {
		opts.DecayHalfLifeDays = 7.0 // 默认 7 天半衰期
	}

	if c.pool != nil {
		return c.searchMemoryDB(ctx, opts)
	}

	return c.searchMemoryFallback(opts), nil
}

func (c *CAS) searchMemoryDB(ctx context.Context, opts MemorySearchOptions) ([]*MemoryChunkHit, error) {
	var queryVector []float32
	if len(opts.QueryVector) > 0 {
		queryVector = opts.QueryVector
	} else if c.embedder != nil && strings.TrimSpace(opts.QueryText) != "" {
		if vec, err := c.embedder.EmbedOne(ctx, opts.QueryText); err == nil {
			queryVector = vec
		}
	}

	var sb strings.Builder
	var args []any
	argIdx := 1

	sb.WriteString(`
		SELECT id, domain, project_id, user_id, reference_id, title, content, chunk_hash, importance_weight, metadata, created_at,
	`)

	if len(queryVector) > 0 {
		args = append(args, vectorLiteral(queryVector))
		sb.WriteString(fmt.Sprintf(" (1 - (embedding <=> $%d::vector)) AS raw_score ", argIdx))
		argIdx++
	} else {
		sb.WriteString(" 1.0 AS raw_score ")
	}

	sb.WriteString(" FROM context_memory_chunks WHERE 1=1 ")

	if opts.Domain != "" {
		sb.WriteString(fmt.Sprintf(" AND domain = $%d ", argIdx))
		args = append(args, opts.Domain)
		argIdx++
	}
	if opts.ProjectID > 0 {
		sb.WriteString(fmt.Sprintf(" AND project_id = $%d ", argIdx))
		args = append(args, opts.ProjectID)
		argIdx++
	}
	if opts.UserID > 0 {
		sb.WriteString(fmt.Sprintf(" AND user_id = $%d ", argIdx))
		args = append(args, opts.UserID)
		argIdx++
	}
	if strings.TrimSpace(opts.QueryText) != "" {
		if len(queryVector) > 0 {
			sb.WriteString(fmt.Sprintf(" AND (embedding IS NOT NULL OR tsv @@ plainto_tsquery('simple', $%d) OR title ILIKE '%%' || $%d || '%%' OR content ILIKE '%%' || $%d || '%%') ", argIdx, argIdx, argIdx))
		} else {
			sb.WriteString(fmt.Sprintf(" AND (tsv @@ plainto_tsquery('simple', $%d) OR title ILIKE '%%' || $%d || '%%' OR content ILIKE '%%' || $%d || '%%') ", argIdx, argIdx, argIdx))
		}
		args = append(args, opts.QueryText)
		argIdx++
	}

	sb.WriteString(" ORDER BY created_at DESC LIMIT 50 ")

	rows, err := c.pool.Query(ctx, sb.String(), args...)
	if err != nil {
		return nil, fmt.Errorf("search memory failed: %w", err)
	}
	defer rows.Close()

	var hits []*MemoryChunkHit
	lambda := math.Ln2 / opts.DecayHalfLifeDays

	for rows.Next() {
		var chunk MemoryChunk
		var metaJSON []byte
		var rawScore float64

		err := rows.Scan(
			&chunk.ID, &chunk.Domain, &chunk.ProjectID, &chunk.UserID, &chunk.ReferenceID,
			&chunk.Title, &chunk.Content, &chunk.ChunkHash, &chunk.ImportanceWeight,
			&metaJSON, &chunk.CreatedAt, &rawScore,
		)
		if err != nil {
			continue
		}
		_ = json.Unmarshal(metaJSON, &chunk.Metadata)

		timeDecay := 1.0
		if opts.ApplyTimeDecay {
			daysDiff := time.Since(chunk.CreatedAt).Hours() / 24.0
			if daysDiff > 0 {
				timeDecay = math.Exp(-lambda * daysDiff)
			}
		}

		finalScore := rawScore * timeDecay * chunk.ImportanceWeight
		hits = append(hits, &MemoryChunkHit{
			Chunk:      &chunk,
			RawScore:   rawScore,
			TimeDecay:  timeDecay,
			FinalScore: finalScore,
		})
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("scan cas chunks failed: %w", err)
	}

	// 按 FinalScore 降序排列
	sortChunkHits(hits)
	if len(hits) > opts.TopK {
		hits = hits[:opts.TopK]
	}

	return hits, nil
}

func (c *CAS) searchMemoryFallback(opts MemorySearchOptions) []*MemoryChunkHit {
	c.mu.RLock()
	defer c.mu.RUnlock()

	var hits []*MemoryChunkHit
	lambda := math.Ln2 / opts.DecayHalfLifeDays

	for _, chunk := range c.memoryChunks {
		if opts.Domain != "" && chunk.Domain != opts.Domain {
			continue
		}
		if opts.ProjectID > 0 && chunk.ProjectID != opts.ProjectID {
			continue
		}
		if opts.UserID > 0 && chunk.UserID != opts.UserID {
			continue
		}

		rawScore := 1.0
		if opts.QueryText != "" {
			if !strings.Contains(chunk.Title, opts.QueryText) && !strings.Contains(chunk.Content, opts.QueryText) {
				continue
			}
		}

		timeDecay := 1.0
		if opts.ApplyTimeDecay {
			daysDiff := time.Since(chunk.CreatedAt).Hours() / 24.0
			if daysDiff > 0 {
				timeDecay = math.Exp(-lambda * daysDiff)
			}
		}

		finalScore := rawScore * timeDecay * chunk.ImportanceWeight
		hits = append(hits, &MemoryChunkHit{
			Chunk:      chunk,
			RawScore:   rawScore,
			TimeDecay:  timeDecay,
			FinalScore: finalScore,
		})
	}

	sortChunkHits(hits)
	if len(hits) > opts.TopK {
		hits = hits[:opts.TopK]
	}
	return hits
}

func sortChunkHits(hits []*MemoryChunkHit) {
	for i := 0; i < len(hits); i++ {
		for j := i + 1; j < len(hits); j++ {
			if hits[j].FinalScore > hits[i].FinalScore {
				hits[i], hits[j] = hits[j], hits[i]
			}
		}
	}
}

func vectorLiteral(v []float32) string {
	parts := make([]string, len(v))
	for i, x := range v {
		parts[i] = fmt.Sprintf("%.6f", x)
	}
	return "[" + strings.Join(parts, ",") + "]"
}

func parseVectorText(s string) []float32 {
	s = strings.Trim(s, "[] ")
	if s == "" {
		return nil
	}
	parts := strings.Split(s, ",")
	res := make([]float32, 0, len(parts))
	for _, p := range parts {
		var f float32
		if _, err := fmt.Sscanf(strings.TrimSpace(p), "%f", &f); err == nil {
			res = append(res, f)
		}
	}
	return res
}