core/server/knowledge/ingest.go

package knowledge

import (
	"context"
	"fmt"
	"io"
	"log"
	"strings"
	"sync"
	"time"

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

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

const (
	// DefaultConcurrency limits parallel ingestion goroutines.
	DefaultConcurrency = 3

	// MaxTextLength is the upper bound for extracted text (10 MB).
	// Documents larger than this are truncated with a warning.
	MaxTextLength = 10 * 1024 * 1024

	// contextTimeout for the entire ingestion of one document.
	ingestTimeout = 10 * time.Minute
)

// textMIMETypes are content types we can extract text from directly.
var textMIMETypes = map[string]bool{
	"text/plain":    true,
	"text/markdown": true,
	"text/csv":      true,
	"text/html":     true,
}

// ── Ingester ─────────────────────────────────

// Ingester manages the document ingestion pipeline.
// It runs chunk + embed as background goroutines with a concurrency
// semaphore to avoid overwhelming the embedding provider.
type Ingester struct {
	stores   store.Stores
	embedder *Embedder
	objStore storage.ObjectStore
	sem      chan struct{}
	wg       sync.WaitGroup
}

// NewIngester creates an ingestion pipeline.
func NewIngester(stores store.Stores, embedder *Embedder, objStore storage.ObjectStore, concurrency int) *Ingester {
	if concurrency <= 0 {
		concurrency = DefaultConcurrency
	}
	return &Ingester{
		stores:   stores,
		embedder: embedder,
		objStore: objStore,
		sem:      make(chan struct{}, concurrency),
	}
}

// IngestDocument starts async ingestion for a document.
// The document must already exist in kb_documents with status "pending"
// and the file must be stored at doc.StorageKey in the object store.
//
// The goroutine progresses through: pending → extracting → chunking →
// embedding → ready (or → error on failure).
func (ing *Ingester) IngestDocument(kb *models.KnowledgeBase, doc *models.KBDocument, userID string, teamID *string) {
	ing.wg.Add(1)
	go func() {
		defer ing.wg.Done()

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

		ctx, cancel := context.WithTimeout(context.Background(), ingestTimeout)
		defer cancel()

		if err := ing.ingest(ctx, kb, doc, userID, teamID); err != nil {
			errMsg := err.Error()
			log.Printf("❌ ingest doc %s (%s): %v", doc.ID, doc.Filename, err)
			ing.stores.KnowledgeBases.UpdateDocumentStatus(ctx, doc.ID, "error", &errMsg)
			ing.stores.KnowledgeBases.Update(ctx, kb.ID, map[string]interface{}{"status": "active"})
			// Notify document owner of ingestion failure (v0.20.0)
			if svc := notifications.Default(); svc != nil {
				notifications.NotifyKBError(svc, userID, kb.ID, kb.Name, errMsg)
			}
		}
	}()
}

// Wait blocks until all in-flight ingestion goroutines complete.
// Call during graceful shutdown.
func (ing *Ingester) Wait() {
	ing.wg.Wait()
}

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

func (ing *Ingester) ingest(ctx context.Context, kb *models.KnowledgeBase, doc *models.KBDocument, userID string, teamID *string) error {
	start := time.Now()
	log.Printf("▶ ingest %s (%s, %d bytes)", doc.Filename, doc.ContentType, doc.SizeBytes)

	// Mark KB as processing
	ing.stores.KnowledgeBases.Update(ctx, kb.ID, map[string]interface{}{"status": "processing"})

	// ── Step 1: Extract text ────────────────
	ing.stores.KnowledgeBases.UpdateDocumentStatus(ctx, doc.ID, "extracting", nil)

	text, err := ing.extractText(ctx, doc)
	if err != nil {
		return fmt.Errorf("extract: %w", err)
	}

	if strings.TrimSpace(text) == "" {
		return fmt.Errorf("extract: no text content found in %s", doc.Filename)
	}

	// Truncate very large documents
	if len(text) > MaxTextLength {
		log.Printf("  ⚠ truncating %s from %d to %d chars", doc.Filename, len(text), MaxTextLength)
		text = text[:MaxTextLength]
	}

	// Store extracted text for potential re-chunking later
	if err := ing.stores.KnowledgeBases.UpdateDocumentText(ctx, doc.ID, text, 0); err != nil {
		log.Printf("  ⚠ failed to store extracted text for %s: %v", doc.ID, err)
	}

	// ── Step 2: Chunk ───────────────────────
	ing.stores.KnowledgeBases.UpdateDocumentStatus(ctx, doc.ID, "chunking", nil)

	cfg := DefaultChunkConfig()
	chunks := SplitText(text, cfg)
	if len(chunks) == 0 {
		return fmt.Errorf("chunk: produced zero chunks from %s", doc.Filename)
	}

	log.Printf("  ✓ chunked %s → %d chunks", doc.Filename, len(chunks))

	// ── Step 3: Embed ───────────────────────
	ing.stores.KnowledgeBases.UpdateDocumentStatus(ctx, doc.ID, "embedding", nil)

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

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

	log.Printf("  ✓ embedded %d chunks (model=%s, dim=%d, tokens=%d)",
		len(chunks), embedResult.Model, embedResult.Dimensions, embedResult.InputTokens)

	// Track embedding usage (role = "embedding")
	ing.embedder.LogUsage(ctx, userID, nil, embedResult)

	// ── Step 4: Store chunks with vectors ───
	dbChunks := make([]models.KBChunk, len(chunks))
	for i, c := range chunks {
		dbChunks[i] = models.KBChunk{
			KBID:       kb.ID,
			DocumentID: doc.ID,
			ChunkIndex: c.Index,
			Content:    c.Content,
			TokenCount: c.TokenCount,
			Embedding:  embedResult.Vectors[i],
			Metadata: models.JSONMap{
				"char_start": i * (cfg.ChunkSize - cfg.ChunkOverlap), // approximate
			},
		}
	}

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

	// ── Step 5: Update stats ────────────────
	// Update document status + chunk count
	if err := ing.stores.KnowledgeBases.UpdateDocumentText(ctx, doc.ID, text, len(chunks)); err != nil {
		log.Printf("  ⚠ failed to update chunk count for %s: %v", doc.ID, err)
	}
	ing.stores.KnowledgeBases.UpdateDocumentStatus(ctx, doc.ID, "ready", nil)

	// Update KB-level stats
	if err := ing.stores.KnowledgeBases.UpdateStats(ctx, kb.ID); err != nil {
		log.Printf("  ⚠ failed to update KB stats: %v", err)
	}

	// Store embedding model info in KB config if not set
	if kb.EmbeddingConfig == nil || kb.EmbeddingConfig["model"] == nil {
		ing.stores.KnowledgeBases.Update(ctx, kb.ID, map[string]interface{}{
			"embedding_config": models.JSONMap{
				"model":      embedResult.Model,
				"dimensions": embedResult.Dimensions,
			},
		})
	}

	// Mark KB as active
	ing.stores.KnowledgeBases.Update(ctx, kb.ID, map[string]interface{}{"status": "active"})

	log.Printf("  ✓ ingest complete: %s (%d chunks, %s)",
		doc.Filename, len(chunks), time.Since(start).Round(time.Millisecond))

	// Notify document owner of successful ingestion (v0.20.0)
	if svc := notifications.Default(); svc != nil {
		notifications.NotifyKBReady(svc, userID, kb.ID, kb.Name, len(chunks))
	}

	return nil
}

// ── Text Extraction ──────────────────────────

// extractText reads the document file and returns its text content.
// For text-based files, reads directly from object store.
// For complex formats (PDF, docx), returns an error — the extraction
// sidecar integration is planned for a future phase.
func (ing *Ingester) extractText(ctx context.Context, doc *models.KBDocument) (string, error) {
	// If extracted text was already stored (e.g. from a rebuild), use it
	if doc.ExtractedText != nil && *doc.ExtractedText != "" {
		return *doc.ExtractedText, nil
	}

	// Text-based files: read directly
	if textMIMETypes[doc.ContentType] {
		return ing.readTextFile(ctx, doc.StorageKey)
	}

	// Complex document types are not yet supported for inline extraction.
	// The extraction sidecar (v0.12.0) handles files but isn't yet
	// wired into the KB pipeline. Planned for a future phase.
	return "", fmt.Errorf("unsupported content type %q — text-based files only for now (txt, md, csv, html)", doc.ContentType)
}

// readTextFile reads a text file from the object store and returns its contents.
func (ing *Ingester) readTextFile(ctx context.Context, storageKey string) (string, error) {
	rc, size, _, err := ing.objStore.Get(ctx, storageKey)
	if err != nil {
		return "", fmt.Errorf("read %s: %w", storageKey, err)
	}
	defer rc.Close()

	// Guard against absurdly large files
	if size > MaxTextLength+1024 {
		// Read up to limit
		limited := io.LimitReader(rc, MaxTextLength)
		data, err := io.ReadAll(limited)
		if err != nil {
			return "", fmt.Errorf("read %s: %w", storageKey, err)
		}
		return string(data), nil
	}

	data, err := io.ReadAll(rc)
	if err != nil {
		return "", fmt.Errorf("read %s: %w", storageKey, err)
	}
	return string(data), nil
}