poimen-workflows/internal/indexing/lessons_index.go

package indexing

import (
	"bufio"
	"encoding/json"
	"fmt"
	"os"
	"strings"
	"sync"
	"time"
)

// Lesson represents a learned lesson from a past failure
type Lesson struct {
	ID            string                 `json:"id"`
	TaskType      string                 `json:"task_type"`
	ActivityType  string                 `json:"activity_type"`
	FailureType   string                 `json:"failure_type"`
	FailureMsg    string                 `json:"failure_msg"`
	Resolution    string                 `json:"resolution"`
	Pattern       string                 `json:"pattern"`
	TimesSeen     int                    `json:"times_seen"`
	LastSeen      time.Time              `json:"last_seen"`
	FirstSeen     time.Time              `json:"first_seen"`
	Metadata      map[string]interface{} `json:"metadata,omitempty"`
}

// LessonIndex provides fast indexed access to lessons
type LessonIndex struct {
	mu              sync.RWMutex
	lessons         map[string]*Lesson           // ID -> Lesson
	byTaskType      map[string][]*Lesson         // TaskType -> Lessons
	byActivityType  map[string][]*Lesson         // ActivityType -> Lessons
	byFailureType   map[string][]*Lesson         // FailureType -> Lessons
	byPattern       map[string][]*Lesson         // Pattern -> Lessons
	sourceFile      string
	lastBuiltTime   time.Time
	lessonCount     int
	buildTime       time.Duration
}

// NewLessonIndex creates a new lesson index
func NewLessonIndex() *LessonIndex {
	return &LessonIndex{
		lessons:        make(map[string]*Lesson),
		byTaskType:     make(map[string][]*Lesson),
		byActivityType: make(map[string][]*Lesson),
		byFailureType:  make(map[string][]*Lesson),
		byPattern:      make(map[string][]*Lesson),
	}
}

// BuildFromFile loads lessons from a JSONL file and builds the index
func (li *LessonIndex) BuildFromFile(filePath string) error {
	li.mu.Lock()
	defer li.mu.Unlock()

	startTime := time.Now()

	// Clear existing index
	li.lessons = make(map[string]*Lesson)
	li.byTaskType = make(map[string][]*Lesson)
	li.byActivityType = make(map[string][]*Lesson)
	li.byFailureType = make(map[string][]*Lesson)
	li.byPattern = make(map[string][]*Lesson)

	// Open file
	file, err := os.Open(filePath)
	if err != nil {
		if os.IsNotExist(err) {
			li.sourceFile = filePath
			li.lastBuiltTime = time.Now()
			li.buildTime = time.Since(startTime)
			return nil // File doesn't exist yet
		}
		return err
	}
	defer file.Close()

	// Read JSONL lines
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		var lesson Lesson
		if err := json.Unmarshal(scanner.Bytes(), &lesson); err != nil {
			continue // Skip malformed lines
		}

		li.addLessonLocked(&lesson)
	}

	if err := scanner.Err(); err != nil {
		return err
	}

	li.sourceFile = filePath
	li.lastBuiltTime = time.Now()
	li.buildTime = time.Since(startTime)
	li.lessonCount = len(li.lessons)

	return nil
}

// addLessonLocked adds a lesson to all indexes (must be called with lock held)
func (li *LessonIndex) addLessonLocked(lesson *Lesson) {
	if lesson.ID == "" {
		return
	}

	li.lessons[lesson.ID] = lesson

	// Index by task type
	if lesson.TaskType != "" {
		li.byTaskType[lesson.TaskType] = append(li.byTaskType[lesson.TaskType], lesson)
	}

	// Index by activity type
	if lesson.ActivityType != "" {
		li.byActivityType[lesson.ActivityType] = append(li.byActivityType[lesson.ActivityType], lesson)
	}

	// Index by failure type
	if lesson.FailureType != "" {
		li.byFailureType[lesson.FailureType] = append(li.byFailureType[lesson.FailureType], lesson)
	}

	// Index by pattern
	if lesson.Pattern != "" {
		li.byPattern[lesson.Pattern] = append(li.byPattern[lesson.Pattern], lesson)
	}
}

// AddLesson adds a single lesson and updates indexes
func (li *LessonIndex) AddLesson(lesson *Lesson) {
	li.mu.Lock()
	defer li.mu.Unlock()

	li.addLessonLocked(lesson)
	li.lessonCount = len(li.lessons)
}

// FindByTaskType returns all lessons for a task type
func (li *LessonIndex) FindByTaskType(taskType string) []*Lesson {
	li.mu.RLock()
	defer li.mu.RUnlock()

	if lessons, exists := li.byTaskType[taskType]; exists {
		// Return a copy to prevent external modifications
		result := make([]*Lesson, len(lessons))
		copy(result, lessons)
		return result
	}

	return make([]*Lesson, 0)
}

// FindByActivityType returns all lessons for an activity type
func (li *LessonIndex) FindByActivityType(activityType string) []*Lesson {
	li.mu.RLock()
	defer li.mu.RUnlock()

	if lessons, exists := li.byActivityType[activityType]; exists {
		result := make([]*Lesson, len(lessons))
		copy(result, lessons)
		return result
	}

	return make([]*Lesson, 0)
}

// FindByFailureType returns all lessons for a failure type
func (li *LessonIndex) FindByFailureType(failureType string) []*Lesson {
	li.mu.RLock()
	defer li.mu.RUnlock()

	if lessons, exists := li.byFailureType[failureType]; exists {
		result := make([]*Lesson, len(lessons))
		copy(result, lessons)
		return result
	}

	return make([]*Lesson, 0)
}

// FindByPattern returns all lessons matching a pattern
func (li *LessonIndex) FindByPattern(pattern string) []*Lesson {
	li.mu.RLock()
	defer li.mu.RUnlock()

	if lessons, exists := li.byPattern[pattern]; exists {
		result := make([]*Lesson, len(lessons))
		copy(result, lessons)
		return result
	}

	return make([]*Lesson, 0)
}

// FindSimilar returns lessons containing a substring in failure message
func (li *LessonIndex) FindSimilar(substr string) []*Lesson {
	li.mu.RLock()
	defer li.mu.RUnlock()

	var results []*Lesson
	substr = strings.ToLower(substr)

	for _, lesson := range li.lessons {
		if strings.Contains(strings.ToLower(lesson.FailureMsg), substr) {
			results = append(results, lesson)
		}
	}

	return results
}

// GetLesson returns a specific lesson by ID
func (li *LessonIndex) GetLesson(id string) (*Lesson, bool) {
	li.mu.RLock()
	defer li.mu.RUnlock()

	lesson, exists := li.lessons[id]
	return lesson, exists
}

// GetStats returns index statistics
func (li *LessonIndex) GetStats() map[string]interface{} {
	li.mu.RLock()
	defer li.mu.RUnlock()

	return map[string]interface{}{
		"total_lessons":       len(li.lessons),
		"unique_task_types":   len(li.byTaskType),
		"unique_activity_types": len(li.byActivityType),
		"unique_failure_types": len(li.byFailureType),
		"unique_patterns":     len(li.byPattern),
		"last_built_time":     li.lastBuiltTime,
		"build_time":          li.buildTime,
		"source_file":         li.sourceFile,
	}
}

// GetAllLessons returns all lessons (for export/debugging)
func (li *LessonIndex) GetAllLessons() []*Lesson {
	li.mu.RLock()
	defer li.mu.RUnlock()

	result := make([]*Lesson, 0, len(li.lessons))
	for _, lesson := range li.lessons {
		result = append(result, lesson)
	}

	return result
}

// Count returns the total number of indexed lessons
func (li *LessonIndex) Count() int {
	li.mu.RLock()
	defer li.mu.RUnlock()

	return len(li.lessons)
}

// Clear clears all indexes
func (li *LessonIndex) Clear() {
	li.mu.Lock()
	defer li.mu.Unlock()

	li.lessons = make(map[string]*Lesson)
	li.byTaskType = make(map[string][]*Lesson)
	li.byActivityType = make(map[string][]*Lesson)
	li.byFailureType = make(map[string][]*Lesson)
	li.byPattern = make(map[string][]*Lesson)
	li.lessonCount = 0
}

// Rebuild rebuilds the index from the source file
func (li *LessonIndex) Rebuild() error {
	if li.sourceFile == "" {
		return fmt.Errorf("no source file set")
	}

	return li.BuildFromFile(li.sourceFile)
}

// QueryMultiple performs a multi-field query (AND logic)
func (li *LessonIndex) QueryMultiple(taskType, activityType, failureType string) []*Lesson {
	li.mu.RLock()
	defer li.mu.RUnlock()

	// Start with the most restrictive set
	var candidates []*Lesson

	// Choose the smallest set to iterate from
	if taskType != "" && activityType != "" && failureType != "" {
		// Use the smallest set
		sizes := []int{
			len(li.byTaskType[taskType]),
			len(li.byActivityType[activityType]),
			len(li.byFailureType[failureType]),
		}

		minIdx := 0
		for i, size := range sizes {
			if size < sizes[minIdx] {
				minIdx = i
			}
		}

		if minIdx == 0 {
			candidates = li.byTaskType[taskType]
		} else if minIdx == 1 {
			candidates = li.byActivityType[activityType]
		} else {
			candidates = li.byFailureType[failureType]
		}
	} else if taskType != "" && activityType != "" {
		if len(li.byTaskType[taskType]) <= len(li.byActivityType[activityType]) {
			candidates = li.byTaskType[taskType]
		} else {
			candidates = li.byActivityType[activityType]
		}
	} else if taskType != "" {
		candidates = li.byTaskType[taskType]
	} else if activityType != "" {
		candidates = li.byActivityType[activityType]
	} else if failureType != "" {
		candidates = li.byFailureType[failureType]
	}

	// Filter candidates
	var results []*Lesson
	for _, lesson := range candidates {
		if taskType != "" && lesson.TaskType != taskType {
			continue
		}
		if activityType != "" && lesson.ActivityType != activityType {
			continue
		}
		if failureType != "" && lesson.FailureType != failureType {
			continue
		}

		results = append(results, lesson)
	}

	return results
}

// GetByTimeRange returns lessons seen within a time range
func (li *LessonIndex) GetByTimeRange(startTime, endTime time.Time) []*Lesson {
	li.mu.RLock()
	defer li.mu.RUnlock()

	var results []*Lesson
	for _, lesson := range li.lessons {
		if !lesson.LastSeen.IsZero() &&
			lesson.LastSeen.After(startTime) &&
			lesson.LastSeen.Before(endTime) {
			results = append(results, lesson)
		}
	}

	return results
}

// GetMostFrequentFailures returns the most frequently seen failures
func (li *LessonIndex) GetMostFrequentFailures(limit int) []*Lesson {
	li.mu.RLock()
	defer li.mu.RUnlock()

	// Convert to slice
	var lessons []*Lesson
	for _, lesson := range li.lessons {
		lessons = append(lessons, lesson)
	}

	// Simple bubble sort (in practice, use a proper sort)
	for i := 0; i < len(lessons); i++ {
		for j := i + 1; j < len(lessons); j++ {
			if lessons[j].TimesSeen > lessons[i].TimesSeen {
				lessons[i], lessons[j] = lessons[j], lessons[i]
			}
		}
	}

	if limit > len(lessons) {
		limit = len(lessons)
	}

	return lessons[:limit]
}