alt/services/worker/internal/backtest/bar_source.go

package backtest

import (
	"context"
	"fmt"
	"sort"
	"strings"
	"time"

	"git.toki-labs.com/toki/alt/packages/domain/backtest"
	"git.toki-labs.com/toki/alt/packages/domain/market"
	"git.toki-labs.com/toki/alt/services/worker/internal/storage"
)

// StorageBarSource adapts worker market-data storage ports to the backtest
// engine's RunSpec input boundary. It also keeps a market/timeframe/date-range
// GetBars method so paper trading can reuse the same storage adapter.
type StorageBarSource struct {
	instruments storage.InstrumentStore
	bars        storage.BarStore
}

func NewStorageBarSource(instruments storage.InstrumentStore, bars storage.BarStore) *StorageBarSource {
	return &StorageBarSource{instruments: instruments, bars: bars}
}

// GetBarsForRun resolves the run's input selector against stored instruments and
// returns only the selected instruments' bars. An empty selector falls back to
// the full market so existing whole-market runs keep working.
func (s *StorageBarSource) GetBarsForRun(ctx context.Context, spec backtest.RunSpec) ([]market.Bar, error) {
	if s == nil || s.instruments == nil || s.bars == nil {
		return nil, fmt.Errorf("backtest storage bar source: storage is not configured")
	}

	instruments, err := s.instruments.ListInstruments(ctx)
	if err != nil {
		return nil, fmt.Errorf("backtest storage bar source: list instruments: %w", err)
	}

	selected, err := selectInstruments(instruments, spec.Market, spec.Selector)
	if err != nil {
		return nil, err
	}

	return s.collectBars(ctx, selected, spec.Timeframe, spec.From, spec.To)
}

// GetBars returns every stored bar for the given market and date range. It is
// retained for paper trading compatibility, which has no run-level selector.
func (s *StorageBarSource) GetBars(ctx context.Context, mkt market.Market, timeframe market.Timeframe, from, to time.Time) ([]market.Bar, error) {
	if s == nil || s.instruments == nil || s.bars == nil {
		return nil, fmt.Errorf("backtest storage bar source: storage is not configured")
	}

	instruments, err := s.instruments.ListInstruments(ctx)
	if err != nil {
		return nil, fmt.Errorf("backtest storage bar source: list instruments: %w", err)
	}

	var selected []market.Instrument
	for _, inst := range instruments {
		if inst.Market == mkt {
			selected = append(selected, inst)
		}
	}

	return s.collectBars(ctx, selected, timeframe, from, to)
}

// collectBars reads bars for each instrument and returns them sorted by
// timestamp with instrument id as a stable tie-breaker.
func (s *StorageBarSource) collectBars(ctx context.Context, instruments []market.Instrument, timeframe market.Timeframe, from, to time.Time) ([]market.Bar, error) {
	var out []market.Bar
	for _, inst := range instruments {
		bars, err := s.bars.GetBars(ctx, inst.ID, timeframe, from, to)
		if err != nil {
			return nil, fmt.Errorf("backtest storage bar source: get bars for %q: %w", inst.ID, err)
		}
		out = append(out, bars...)
	}

	sort.Slice(out, func(i, j int) bool {
		if out[i].Timestamp.Equal(out[j].Timestamp) {
			return out[i].InstrumentID < out[j].InstrumentID
		}
		return out[i].Timestamp.Before(out[j].Timestamp)
	})
	return out, nil
}

// selectInstruments resolves the input selector within the run's market. An empty
// selector returns the whole market. Each requested instrument id and symbol must
// resolve to a stored instrument; an unresolved selector entry returns an error so
// a backtest never runs silently on the wrong input.
func selectInstruments(instruments []market.Instrument, mkt market.Market, selector backtest.InputSelector) ([]market.Instrument, error) {
	inMarket := make([]market.Instrument, 0, len(instruments))
	for _, inst := range instruments {
		if inst.Market == mkt {
			inMarket = append(inMarket, inst)
		}
	}

	wantIDs := make([]string, 0, len(selector.InstrumentIDs))
	for _, id := range selector.InstrumentIDs {
		trimmed := strings.TrimSpace(string(id))
		if trimmed == "" {
			return nil, fmt.Errorf("backtest storage bar source: selector instrument id is empty")
		}
		wantIDs = append(wantIDs, trimmed)
	}
	wantSymbols := make([]string, 0, len(selector.Symbols))
	for _, sym := range selector.Symbols {
		trimmed := strings.TrimSpace(sym)
		if trimmed == "" {
			return nil, fmt.Errorf("backtest storage bar source: selector symbol is empty")
		}
		wantSymbols = append(wantSymbols, trimmed)
	}

	// Empty selector keeps the historical whole-market behavior.
	if len(wantIDs) == 0 && len(wantSymbols) == 0 {
		return inMarket, nil
	}

	selected := make([]market.Instrument, 0, len(wantIDs)+len(wantSymbols))
	seen := make(map[market.InstrumentID]bool)
	add := func(inst market.Instrument) {
		if !seen[inst.ID] {
			seen[inst.ID] = true
			selected = append(selected, inst)
		}
	}

	for _, want := range wantIDs {
		found := false
		for _, inst := range inMarket {
			if string(inst.ID) == want {
				add(inst)
				found = true
				break
			}
		}
		if !found {
			return nil, fmt.Errorf("backtest storage bar source: selector instrument id %q not found in market %q", want, mkt)
		}
	}

	for _, want := range wantSymbols {
		found := false
		for _, inst := range inMarket {
			if instrumentMatchesSymbol(inst, want) {
				add(inst)
				found = true
				break
			}
		}
		if !found {
			return nil, fmt.Errorf("backtest storage bar source: selector symbol %q not found in market %q", want, mkt)
		}
	}

	return selected, nil
}

// instrumentMatchesSymbol matches a selector symbol against the instrument's
// canonical symbol and any provider symbol value.
func instrumentMatchesSymbol(inst market.Instrument, symbol string) bool {
	if strings.TrimSpace(inst.Symbol) == symbol {
		return true
	}
	for _, providerSymbol := range inst.ProviderSymbols {
		if strings.TrimSpace(providerSymbol) == symbol {
			return true
		}
	}
	return false
}