core/server/workspace/indexer.go

package workspace

import (
	"context"
	"fmt"
	"io"
	"log"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"switchboard-core/knowledge"
	"switchboard-core/models"
	"switchboard-core/notifications"
	"switchboard-core/store"
)

// ── Configuration ────────────────────────────

const (
	// indexTimeout is the max time for indexing a single file.
	indexTimeout = 5 * time.Minute

	// maxIndexTextLength caps text extracted from a single file (5 MB).
	maxIndexTextLength = 5 * 1024 * 1024
)

// indexableExtensions defines source code and text file extensions
// eligible for chunking and embedding. MIME detection covers text/*,
// but code files often have application/* types — this map catches them.
var indexableExtensions = map[string]bool{
	".go": true, ".py": true, ".js": true, ".ts": true, ".jsx": true,
	".tsx": true, ".rs": true, ".c": true, ".cpp": true, ".h": true,
	".hpp": true, ".java": true, ".rb": true, ".php": true, ".swift": true,
	".kt": true, ".scala": true, ".sh": true, ".bash": true, ".zsh": true,
	".sql": true, ".yaml": true, ".yml": true, ".toml": true, ".json": true,
	".xml": true, ".css": true, ".scss": true, ".less": true,
	".md": true, ".txt": true, ".csv": true, ".html": true, ".htm": true,
	".dockerfile": true, ".tf": true, ".hcl": true, ".proto": true,
	".graphql": true, ".vue": true, ".svelte": true,
	".pl": true, ".pm": true, // Perl
	".lua": true, ".zig": true, ".nim": true, ".ex": true, ".exs": true,
	// Config / data
	".ini": true, ".cfg": true, ".conf": true, ".env": true,
	".mk": true, ".cmake": true, ".makefile": true,
	".r": true, ".jl": true, // R, Julia
}

// CodeChunkConfig returns chunking settings optimized for source code.
// Larger chunks (1500 chars) because functions can be long, with
// separators tuned to split on function/class/type boundaries.
func CodeChunkConfig() knowledge.ChunkConfig {
	return knowledge.ChunkConfig{
		ChunkSize:    1500,
		ChunkOverlap: 200,
		Separators: []string{
			"\n\nfunc ", "\n\nfn ",     // Go, Rust
			"\n\nclass ", "\n\ndef ",    // Python, Ruby
			"\n\ntype ", "\n\nsub ",     // Go types, Perl
			"\n\npub ", "\n\nimpl ",     // Rust
			"\n\n", "\n", " ",           // fallback hierarchy
		},
	}
}

// isIndexable returns true if the file is eligible for text indexing.
func isIndexable(path string, contentType string) bool {
	// Check extension first (catches source code with non-text MIME types)
	ext := strings.ToLower(filepath.Ext(path))
	if indexableExtensions[ext] {
		return true
	}
	// Also index any text/* content type
	if strings.HasPrefix(contentType, "text/") {
		return true
	}
	// Specific filenames without extension
	base := strings.ToLower(filepath.Base(path))
	switch base {
	case "makefile", "dockerfile", "jenkinsfile", "gemfile",
		"rakefile", "procfile", "vagrantfile", ".gitignore",
		".dockerignore", ".editorconfig":
		return true
	}
	return false
}

// isCodeFile returns true if the file is a source code file (uses code chunking).
func isCodeFile(path string) bool {
	ext := strings.ToLower(filepath.Ext(path))
	switch ext {
	case ".md", ".txt", ".csv", ".html", ".htm", ".xml", ".json", ".yaml", ".yml":
		return false // prose / data — use default chunking
	}
	return indexableExtensions[ext]
}

// ── Indexer ──────────────────────────────────

// Indexer processes workspace files through the chunking + embedding pipeline.
// It reuses the knowledge.Embedder for vector generation and shares the
// concurrency semaphore with KB ingestion to avoid overwhelming the
// embedding provider.
type Indexer struct {
	stores   store.Stores
	embedder *knowledge.Embedder
	wfs      *FS
	sem      chan struct{} // shared semaphore with KB ingestion
	wg       sync.WaitGroup
	enabled  bool // global kill switch (WORKSPACE_INDEXING_ENABLED)
}

// NewIndexer creates a workspace indexer.
// The semaphore should be shared with the KB Ingester to cap total
// embedding concurrency.
func NewIndexer(stores store.Stores, embedder *knowledge.Embedder, wfs *FS, sem chan struct{}, enabled bool) *Indexer {
	return &Indexer{
		stores:   stores,
		embedder: embedder,
		wfs:      wfs,
		sem:      sem,
		enabled:  enabled,
	}
}

// IsEnabled returns true if workspace indexing is globally enabled
// AND the embedding role is configured.
func (idx *Indexer) IsEnabled(ctx context.Context) bool {
	return idx.enabled && idx.embedder.IsConfigured(ctx)
}

// IndexFile starts async indexing for a single workspace file.
// Called from WorkspaceFS.WriteFile after a successful write.
// No-op if indexing is disabled, the file is non-indexable, or the
// sha256 hasn't changed (content-addressed skip).
func (idx *Indexer) IndexFile(w *models.Workspace, f *models.WorkspaceFile, userID string, teamID *string) {
	if !idx.enabled || !w.IndexingEnabled {
		return
	}

	// Check if indexable
	if f.IsDirectory || !isIndexable(f.Path, f.ContentType) {
		// Mark as skipped so the UI knows it's intentional
		ctx := context.Background()
		idx.stores.Workspaces.UpdateFileIndexStatus(ctx, f.ID, models.IndexStatusSkipped, 0)
		return
	}

	// Mark as pending
	ctx := context.Background()
	idx.stores.Workspaces.UpdateFileIndexStatus(ctx, f.ID, models.IndexStatusPending, 0)

	idx.wg.Add(1)
	go func() {
		defer idx.wg.Done()

		// Acquire semaphore
		idx.sem <- struct{}{}
		defer func() { <-idx.sem }()

		ictx, cancel := context.WithTimeout(context.Background(), indexTimeout)
		defer cancel()

		if err := idx.indexFile(ictx, w, f, userID, teamID); err != nil {
			log.Printf("❌ workspace index %s/%s: %v", w.ID[:8], f.Path, err)
			idx.stores.Workspaces.UpdateFileIndexStatus(ictx, f.ID, models.IndexStatusError, 0)
		}
	}()
}

// IndexBatch starts async indexing for multiple files (e.g. after archive extract).
// Files are indexed sequentially within a single goroutine to be friendly
// to rate limits.
func (idx *Indexer) IndexBatch(w *models.Workspace, files []models.WorkspaceFile, userID string, teamID *string) {
	if !idx.enabled || !w.IndexingEnabled {
		return
	}

	// Filter to indexable files only
	var indexable []models.WorkspaceFile
	ctx := context.Background()
	for _, f := range files {
		if f.IsDirectory || !isIndexable(f.Path, f.ContentType) {
			idx.stores.Workspaces.UpdateFileIndexStatus(ctx, f.ID, models.IndexStatusSkipped, 0)
			continue
		}
		idx.stores.Workspaces.UpdateFileIndexStatus(ctx, f.ID, models.IndexStatusPending, 0)
		indexable = append(indexable, f)
	}

	if len(indexable) == 0 {
		return
	}

	idx.wg.Add(1)
	go func() {
		defer idx.wg.Done()

		// One semaphore slot for the entire batch
		idx.sem <- struct{}{}
		defer func() { <-idx.sem }()

		var indexed, errors int
		for _, f := range indexable {
			f := f
			ictx, cancel := context.WithTimeout(context.Background(), indexTimeout)
			if err := idx.indexFile(ictx, w, &f, userID, teamID); err != nil {
				log.Printf("❌ workspace index %s/%s: %v", w.ID[:8], f.Path, err)
				idx.stores.Workspaces.UpdateFileIndexStatus(ictx, f.ID, models.IndexStatusError, 0)
				errors++
			} else {
				indexed++
			}
			cancel()
		}

		log.Printf("📑 workspace batch index %s: %d indexed, %d errors, %d skipped",
			w.ID[:8], indexed, errors, len(files)-len(indexable))

		// Notify on completion
		if svc := notifications.Default(); svc != nil && indexed > 0 {
			// Find workspace owner for notification target
			if w.OwnerType == models.WorkspaceOwnerUser {
				svc.Notify(context.Background(), &models.Notification{
					UserID:       w.OwnerID,
					Type:         "workspace.indexed",
					Title:        "Workspace indexed",
					Body:         fmt.Sprintf("%d files indexed in workspace %q", indexed, w.Name),
					ResourceType: "workspace",
					ResourceID:   w.ID,
				})
			}
		}
	}()
}

// Wait blocks until all in-flight indexing goroutines complete.
func (idx *Indexer) Wait() {
	idx.wg.Wait()
}

// ── Internal Pipeline ────────────────────────

func (idx *Indexer) indexFile(ctx context.Context, w *models.Workspace, f *models.WorkspaceFile, userID string, teamID *string) error {
	start := time.Now()

	// Mark as indexing
	idx.stores.Workspaces.UpdateFileIndexStatus(ctx, f.ID, models.IndexStatusIndexing, 0)

	// Check if embedding is configured (might have changed since enqueue)
	if !idx.embedder.IsConfigured(ctx) {
		return fmt.Errorf("embedding role not configured")
	}

	// ── Step 1: Read file content ───────────
	rc, _, err := idx.wfs.ReadFile(ctx, w, f.Path)
	if err != nil {
		return fmt.Errorf("read: %w", err)
	}
	defer rc.Close()

	// Cap at max text length
	limited := io.LimitReader(rc, maxIndexTextLength)
	data, err := io.ReadAll(limited)
	if err != nil {
		return fmt.Errorf("read content: %w", err)
	}
	text := string(data)

	if strings.TrimSpace(text) == "" {
		idx.stores.Workspaces.UpdateFileIndexStatus(ctx, f.ID, models.IndexStatusSkipped, 0)
		return nil
	}

	// ── Step 2: Chunk ───────────────────────
	var cfg knowledge.ChunkConfig
	if isCodeFile(f.Path) {
		cfg = CodeChunkConfig()
	} else {
		cfg = knowledge.DefaultChunkConfig()
	}

	chunks := knowledge.SplitText(text, cfg)
	if len(chunks) == 0 {
		idx.stores.Workspaces.UpdateFileIndexStatus(ctx, f.ID, models.IndexStatusSkipped, 0)
		return nil
	}

	// ── Step 3: Embed ───────────────────────
	texts := make([]string, len(chunks))
	for i, c := range chunks {
		texts[i] = c.Content
	}

	embedResult, err := idx.embedder.EmbedChunks(ctx, userID, teamID, texts)
	if err != nil {
		return fmt.Errorf("embed: %w", err)
	}

	// Track usage
	idx.embedder.LogUsage(ctx, userID, nil, embedResult)

	// ── Step 4: Delete old chunks + insert new ─
	if err := idx.stores.Workspaces.DeleteChunksByFile(ctx, f.ID); err != nil {
		return fmt.Errorf("delete old chunks: %w", err)
	}

	dbChunks := make([]models.WorkspaceChunk, len(chunks))
	for i, c := range chunks {
		// Compute approximate line number from character offset
		lineStart := 0
		if i > 0 {
			charOffset := i * (cfg.ChunkSize - cfg.ChunkOverlap)
			if charOffset < len(text) {
				lineStart = strings.Count(text[:charOffset], "\n") + 1
			}
		}

		dbChunks[i] = models.WorkspaceChunk{
			WorkspaceID: w.ID,
			FileID:      f.ID,
			ChunkIndex:  c.Index,
			Content:     c.Content,
			TokenCount:  c.TokenCount,
			Embedding:   embedResult.Vectors[i],
			Metadata: models.JSONMap{
				"line_start": lineStart,
				"language":   inferLanguage(f.Path),
			},
		}
	}

	if err := idx.stores.Workspaces.InsertChunks(ctx, dbChunks); err != nil {
		return fmt.Errorf("store chunks: %w", err)
	}

	// ── Step 5: Update file index status ────
	if err := idx.stores.Workspaces.UpdateFileIndexStatus(ctx, f.ID, models.IndexStatusReady, len(chunks)); err != nil {
		log.Printf("⚠ failed to update index status for %s: %v", f.Path, err)
	}

	log.Printf("  📑 indexed %s → %d chunks (%s)", f.Path, len(chunks), time.Since(start).Round(time.Millisecond))
	return nil
}

// inferLanguage extracts a simple language identifier from file extension.
func inferLanguage(path string) string {
	ext := strings.ToLower(filepath.Ext(path))
	if ext == "" {
		return ""
	}
	// Strip leading dot
	return ext[1:]
}