aitrade/pkg/app/news/aggregator.go

package news

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

	"github.com/pheinrich/aitrade/pkg/config"
	"github.com/pheinrich/aitrade/pkg/db"
	"golang.org/x/sync/errgroup"
)

type Aggregator struct {
	sources       []Source
	newsRepo      *db.NewsRepository
	analyzer      *Analyzer
	llmScorer     *LLMScorer
	pollInterval  time.Duration
	fetchTimeout  time.Duration
	logger        *slog.Logger
	onNewsUpdated func() // Callback for SSE notifications
}

func NewAggregator(
	newsRepo *db.NewsRepository,
	pollInterval time.Duration,
	fetchTimeout time.Duration,
	llmScorer *LLMScorer,
	logger *slog.Logger,
) *Aggregator {
	agg := &Aggregator{
		sources:      make([]Source, 0),
		newsRepo:     newsRepo,
		analyzer:     NewAnalyzer(),
		llmScorer:    llmScorer,
		pollInterval: pollInterval,
		fetchTimeout: fetchTimeout,
		logger:       logger,
	}

	return agg
}

// SetNewsUpdateCallback sets a callback that's called when news are updated
func (a *Aggregator) SetNewsUpdateCallback(callback func()) {
	a.onNewsUpdated = callback
}

// AddSourcesFromConfig adds sources from config
func (a *Aggregator) AddSourcesFromConfig(sources []config.NewsSource, defaultRateLimit *config.RateLimit) {
	for _, src := range sources {
		if !src.Enabled || src.Name == "" || src.URL == "" {
			continue
		}

		var source Source

		switch src.Type {
		case "rss":
			rssSource := NewRSSSource(src.Name, src.URL, a.logger)

			// Add auth if provided
			if src.Auth != nil && src.Auth.Type == "basic" {
				rssSource.SetBasicAuth(src.Auth.Username, src.Auth.Password)
			}

			// Add custom headers if provided
			for key, value := range src.Headers {
				rssSource.AddHeader(key, value)
			}

			source = rssSource

		case "alphavantage":
			source = NewAlphaVantageSource(src.Name, src.URL, a.logger)

		case "finnhub":
			source = NewFinnhubSource(src.Name, src.URL, a.logger)

		default:
			a.logger.Warn("unsupported source type", slog.String("type", src.Type), slog.String("name", src.Name))
			continue
		}

		// Set rate limits: source-specific OR default
		rateLimit := src.RateLimit
		if rateLimit == nil {
			rateLimit = defaultRateLimit
		}

		if rateLimit != nil {
			// All source types implement SetRateLimit via their embedded rate limiter
			switch s := source.(type) {
			case *RSSSource:
				s.SetRateLimit(rateLimit.MaxPerHour, rateLimit.MaxPerDay)
			case *AlphaVantageSource:
				s.SetRateLimit(rateLimit.MaxPerHour, rateLimit.MaxPerDay)
			case *FinnhubSource:
				s.SetRateLimit(rateLimit.MaxPerHour, rateLimit.MaxPerDay)
			}

			a.logger.Info("rate limit configured",
				slog.String("source", src.Name),
				slog.Int("max_per_hour", rateLimit.MaxPerHour),
				slog.Int("max_per_day", rateLimit.MaxPerDay))
		}

		a.AddSource(source)
	}
}

func (a *Aggregator) AddSource(source Source) {
	a.sources = append(a.sources, source)
	a.logger.Info("added news source", slog.String("source", source.Name()))
}

func (a *Aggregator) Run(ctx context.Context) error {
	a.logger.Info("news aggregator starting",
		slog.Int("sources", len(a.sources)),
		slog.Duration("poll_interval", a.pollInterval),
	)

	// Fetch immediately on start
	if err := a.fetchAllSources(ctx); err != nil {
		a.logger.Error("initial news fetch failed", slog.Any("error", err))
	}

	ticker := time.NewTicker(a.pollInterval)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			a.logger.Info("news aggregator stopping")
			return ctx.Err()
		case <-ticker.C:
			if err := a.fetchAllSources(ctx); err != nil {
				a.logger.Error("news fetch failed", slog.Any("error", err))
			}
		}
	}
}

func (a *Aggregator) fetchAllSources(ctx context.Context) error {
	a.logger.Debug("fetching from all news sources")

	var totalFetched, totalStored atomic.Int64

	// Use errgroup for parallel fetching
	g, gCtx := errgroup.WithContext(ctx)

	// Launch goroutines for each source
	for _, source := range a.sources {
		src := source // Capture loop variable
		g.Go(func() error {
			// Create a timeout context for this specific source
			sourceCtx, cancel := context.WithTimeout(gCtx, a.fetchTimeout)
			defer cancel()

			articles, err := src.Fetch(sourceCtx)
			if err != nil {
				a.logger.Warn("failed to fetch from source",
					slog.String("source", src.Name()),
					slog.Any("error", err),
				)
				return nil // Don't fail entire operation on single source error
			}

			totalFetched.Add(int64(len(articles)))

			// Analyze sentiment and store
			for _, article := range articles {
				// Always use keyword analyzer first (fast)
				a.analyzer.Analyze(article)

				// If LLM scorer enabled, use it (may fallback to keyword)
				if a.llmScorer != nil && a.llmScorer.Enabled() {
					if err := a.llmScorer.Analyze(sourceCtx, article); err != nil {
						a.logger.Warn("LLM scoring failed",
							slog.String("url", article.URL),
							slog.Any("error", err))
						// Article already has keyword sentiment, continue
					}
				} else {
					// No LLM scorer, mark as keyword-only
					article.SentimentMethod = "keyword"
				}

				if err := a.newsRepo.Create(sourceCtx, article); err != nil {
					a.logger.Error("failed to store article",
						slog.String("url", article.URL),
						slog.Any("error", err),
					)
					continue
				}

				totalStored.Add(1)
			}

			return nil
		})
	}

	// Wait for all goroutines to complete
	if err := g.Wait(); err != nil {
		return err
	}

	fetched := int(totalFetched.Load())
	stored := int(totalStored.Load())

	a.logger.Info("news fetch completed",
		slog.Int("fetched", fetched),
		slog.Int("stored", stored),
		slog.Int("duplicates", fetched-stored),
	)

	// Notify listeners if new articles were stored
	if stored > 0 && a.onNewsUpdated != nil {
		a.onNewsUpdated()
	}

	return nil
}

func (a *Aggregator) GetRecent(ctx context.Context, limit int) ([]*db.NewsRepository, error) {
	return nil, fmt.Errorf("not implemented")
}