email_agent/internal/classify/classify.go

// Classifier orchestration: read the live category list and rules, ask JEV about
// every pending message, apply the user's rules as overrides, and store the
// result (FR-3, FR-4, FR-5).
package classify

import (
	"context"
	"fmt"
	"log/slog"
	"strings"
	"sync"
	"time"

	"gitea.lab.audasmedia.com.au/sam/email_agent/internal/store"
)

// ActorAgent is the actor recorded in the action log for classifier writes.
const ActorAgent = "agent"

// Summary describes one classification run.
type Summary struct {
	Considered   int
	Classified   int
	Failed       int
	Model        string
	RulesVersion int
	Remaining    int
	Duration     time.Duration
}

// Classifier ties the JEV client to the store.
type Classifier struct {
	client      *Client
	store       *store.Store
	logger      *slog.Logger
	concurrency int
	maxPerRun   int
}

// NewClassifier builds a Classifier. concurrency and maxPerRun fall back to safe
// defaults when they are not positive.
func NewClassifier(client *Client, st *store.Store, logger *slog.Logger, concurrency, maxPerRun int) *Classifier {
	if concurrency < 1 {
		concurrency = 4
	}
	if concurrency > 16 {
		concurrency = 16
	}
	if maxPerRun < 1 {
		maxPerRun = 200
	}
	if logger == nil {
		logger = slog.New(slog.DiscardHandler)
	}
	return &Classifier{
		client:      client,
		store:       st,
		logger:      logger,
		concurrency: concurrency,
		maxPerRun:   maxPerRun,
	}
}

// Model returns the configured JEV model alias.
func (c *Classifier) Model() string { return c.client.Model() }

// Configured reports whether JEV_API_KEY is present.
func (c *Classifier) Configured() bool { return c.client.Configured() }

// Run classifies the pending messages, up to maxPerRun. It returns a Summary
// even when individual messages fail; only setup failures return an error.
func (c *Classifier) Run(ctx context.Context) (Summary, error) {
	return c.run(ctx, false)
}

// RunAll re-classifies every stored message. It is used after the category list
// or the rules change, so the new settings take effect on existing mail.
func (c *Classifier) RunAll(ctx context.Context) (Summary, error) {
	return c.run(ctx, true)
}

func (c *Classifier) run(ctx context.Context, all bool) (Summary, error) {
	if !c.client.Configured() {
		return Summary{}, ErrNotConfigured
	}

	start := time.Now()
	mode := "pending"
	if all {
		mode = "all"
	}
	categories, err := c.store.ListCategories(ctx)
	if err != nil {
		return Summary{}, err
	}
	if len(categories) == 0 {
		return Summary{}, fmt.Errorf("no categories configured")
	}
	options := make([]CategoryOption, 0, len(categories))
	for _, cat := range categories {
		options = append(options, CategoryOption{Name: cat.Name, Description: cat.Description})
	}
	valid := make(map[string]string, len(categories)) // lower(name) -> name
	for _, cat := range categories {
		valid[strings.ToLower(cat.Name)] = cat.Name
	}

	rulesText, rulesVersion, err := c.store.CurrentRules(ctx)
	if err != nil {
		return Summary{}, err
	}
	rules := ParseRules(rulesText)
	for _, perr := range rules.Errors {
		c.logger.Warn("rules line skipped", "error", perr.Error())
	}

	pending, err := c.pending(ctx, all)
	if err != nil {
		return Summary{}, err
	}

	sum := Summary{
		Considered:   len(pending),
		Model:        c.client.Model(),
		RulesVersion: rulesVersion,
	}
	if len(pending) == 0 {
		sum.Remaining = 0
		sum.Duration = time.Since(start)
		return sum, nil
	}

	type outcome struct {
		msg    store.Message
		class  store.Classification
		rule   *Rule
		failed bool
	}

	work := make(chan store.Message)
	results := make(chan outcome)
	var wg sync.WaitGroup
	for i := 0; i < c.concurrency; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for msg := range work {
				class, rule, err := c.classifyOne(ctx, msg, options, valid, rules)
				if err != nil {
					c.logger.Warn("classify message failed",
						"message_id", msg.ID, "subject", msg.Subject, "error", err)
					results <- outcome{msg: msg, failed: true}
					continue
				}
				results <- outcome{msg: msg, class: class, rule: rule}
			}
		}()
	}

	go func() {
		defer close(work)
		for _, msg := range pending {
			select {
			case <-ctx.Done():
				return
			case work <- msg:
			}
		}
	}()

	go func() {
		wg.Wait()
		close(results)
	}()

	for res := range results {
		if res.failed {
			sum.Failed++
			continue
		}
		if _, err := c.store.SaveClassification(ctx, res.class); err != nil {
			c.logger.Error("save classification failed", "message_id", res.msg.ID, "error", err)
			sum.Failed++
			continue
		}
		detail := fmt.Sprintf("category=%s importance=%s action=%s confidence=%.2f model=%s rules_version=%d",
			res.class.Category, res.class.Importance, res.class.Action,
			res.class.Confidence, res.class.Model, rulesVersion)
		if res.rule != nil {
			detail += fmt.Sprintf(" rule_line=%d rule=%q", res.rule.Line, res.rule.Raw)
		}
		if err := c.store.LogMessageAction(ctx, res.msg.ID, "classify", ActorAgent, detail); err != nil {
			c.logger.Warn("log classify action failed", "message_id", res.msg.ID, "error", err)
		}
		sum.Classified++
	}

	if cov, err := c.store.ClassificationCoverage(ctx); err == nil {
		sum.Remaining = cov.Missing()
	}
	sum.Duration = time.Since(start)

	if err := c.store.LogAction(ctx, "classify_run", ActorAgent,
		fmt.Sprintf("mode=%s considered=%d classified=%d failed=%d remaining=%d model=%s rules_version=%d",
			mode, sum.Considered, sum.Classified, sum.Failed, sum.Remaining, sum.Model, rulesVersion)); err != nil {
		c.logger.Warn("log classify run failed", "error", err)
	}

	c.logger.Info("classification run complete",
		"considered", sum.Considered,
		"classified", sum.Classified,
		"failed", sum.Failed,
		"remaining", sum.Remaining,
		"model", sum.Model,
		"duration_ms", sum.Duration.Milliseconds(),
	)
	return sum, nil
}

// classifyOne asks JEV about one message and overlays the matching user rule.
func (c *Classifier) classifyOne(ctx context.Context, msg store.Message, options []CategoryOption, valid map[string]string, rules RuleSet) (store.Classification, *Rule, error) {
	result, err := c.client.Classify(ctx, MessageInput{
		From:     msg.FromAddr,
		To:       msg.ToAddr,
		Subject:  msg.Subject,
		Source:   msg.SourceAddress,
		Body:     msg.BodyText,
		Received: receivedString(msg),
	}, options)
	if err != nil {
		return store.Classification{}, nil, err
	}

	rule := rules.Match(map[string]string{
		"from":    msg.FromAddr,
		"to":      msg.ToAddr,
		"subject": msg.Subject,
		"source":  msg.SourceAddress,
		"body":    msg.BodyText,
	})

	category := result.Category
	importance := result.Importance
	action := result.Action
	if rule != nil {
		if v, ok := rule.Effects["category"]; ok {
			if name, found := valid[strings.ToLower(v)]; found {
				category = name
			} else {
				c.logger.Warn("rule names an unknown category; ignoring effect",
					"message_id", msg.ID, "rule_line", rule.Line, "category", v)
			}
		}
		if v, ok := rule.Effects["importance"]; ok {
			importance = v
		}
		if v, ok := rule.Effects["action"]; ok {
			action = v
		}
	}

	// The classifier must always store a usable category, importance, and
	// action (Phase 3 DoD). JEV can only pick from the supplied options, but a
	// broken response or a missing category list must not leave NULLs.
	if name, found := valid[strings.ToLower(category)]; found {
		category = name
	} else {
		if category != "" {
			c.logger.Warn("classifier returned an unknown category; using fallback",
				"message_id", msg.ID, "category", category)
		}
		category = options[0].Name
	}
	if !allowedImportance[strings.ToLower(importance)] {
		importance = "medium"
	} else {
		importance = strings.ToLower(importance)
	}
	if !allowedAction[strings.ToLower(action)] {
		action = "keep"
	} else {
		action = strings.ToLower(action)
	}

	// A rule override invalidates the JEV confidence: report the confidence of
	// the answer the category actually came from.
	confidence := result.CategoryConfidence
	if rule != nil {
		if _, ok := rule.Effects["category"]; ok {
			confidence = 1
		}
	}

	return store.Classification{
		MessageID:  msg.ID,
		Category:   category,
		Importance: importance,
		Action:     action,
		Confidence: confidence,
		Model:      result.Model,
	}, rule, nil
}

// pending returns the messages to classify for this run.
func (c *Classifier) pending(ctx context.Context, all bool) ([]store.Message, error) {
	if all {
		return c.store.AllMessagesForClassification(ctx, c.maxPerRun)
	}
	return c.store.PendingClassifications(ctx, c.maxPerRun)
}

func receivedString(msg store.Message) string {
	t := msg.DateSent
	if t.IsZero() {
		t = msg.ReceivedAt
	}
	if t.IsZero() {
		return ""
	}
	return t.UTC().Format(time.RFC3339)
}