alt/services/worker/internal/papertrading/engine.go

package papertrading

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

	"git.toki-labs.com/toki/alt/packages/domain/backtest"
	"git.toki-labs.com/toki/alt/packages/domain/market"
)

// BarSource resolves daily bars for a paper trading run.
type BarSource interface {
	GetBars(ctx context.Context, mkt market.Market, timeframe market.Timeframe, from, to time.Time) ([]market.Bar, error)
}

// StrategyPort resolves strategies for a paper trading run.
type StrategyPort interface {
	GetStrategy(ctx context.Context, id backtest.StrategyID) (backtest.Strategy, error)
}

// Engine implements the daily paper execution loop. Orders decided in one bar
// fill on the next bar: market orders at next-bar open, limit orders when the
// bar's high/low crosses the limit price.
type Engine struct {
	barSource    BarSource
	strategyPort StrategyPort
}

// NewEngine creates a new paper trading Engine instance.
func NewEngine(barSource BarSource, strategyPort StrategyPort) *Engine {
	return &Engine{
		barSource:    barSource,
		strategyPort: strategyPort,
	}
}

// RunRequest carries all inputs for one paper execution run.
type RunRequest struct {
	Account backtest.PaperAccount
	Run     backtest.Run
}

// RejectedOrder records an order that was denied by a risk gate or could not
// be applied to the portfolio.
type RejectedOrder struct {
	Order      backtest.OrderIntent
	Reason     string
	BarTime    time.Time
	Instrument market.InstrumentID
	CashBefore backtest.PaperAccount
}

// Snapshot captures the terminal state of a paper run.
type Snapshot struct {
	Account     backtest.PaperAccount
	Fills       []backtest.Fill
	Rejected    []RejectedOrder
	EquityCurve []backtest.EquityPoint
}

// Run executes a daily paper trading simulation and returns the final snapshot.
func (e *Engine) Run(ctx context.Context, req RunRequest) (*Snapshot, error) {
	strategy, err := e.strategyPort.GetStrategy(ctx, req.Run.Spec.StrategyID)
	if err != nil {
		return nil, fmt.Errorf("failed to get strategy %s: %w", req.Run.Spec.StrategyID, err)
	}

	bars, err := e.barSource.GetBars(ctx, req.Run.Spec.Market, req.Run.Spec.Timeframe, req.Run.Spec.From, req.Run.Spec.To)
	if err != nil {
		return nil, fmt.Errorf("failed to get bars: %w", err)
	}

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

	if len(bars) == 0 {
		return &Snapshot{
			Account:     req.Account,
			Fills:       nil,
			Rejected:    nil,
			EquityCurve: nil,
		}, nil
	}

	// Clone starting account state so we don't mutate the input
	positionsCopy := make(map[market.InstrumentID]backtest.Position, len(req.Account.Portfolio.Positions))
	for k, v := range req.Account.Portfolio.Positions {
		positionsCopy[k] = v
	}
	account := req.Account
	account.Portfolio = backtest.PortfolioState{
		Cash:      req.Account.Portfolio.Cash,
		Positions: positionsCopy,
	}
	var equityCurve []backtest.EquityPoint
	var fills []backtest.Fill
	var rejected []RejectedOrder

	// Instrument-scoped pending: an order is only eligible to fill on a later bar
	// of the *same* instrument. This prevents orders for instrument A from
	// accidentally matching on the same timestamp of instrument B when
	// StorageBarSource returns multi-instrument bars sorted by time.
	pendingOrders := make(map[market.InstrumentID][]backtest.OrderIntent)

	for _, bar := range bars {
		// --- Phase 1: Fill pending orders for this instrument ---
		instrumentPending, hasPending := pendingOrders[bar.InstrumentID]
		if hasPending {
			delete(pendingOrders, bar.InstrumentID)
		}

		var newFills []backtest.Fill
		var newRejected []RejectedOrder
		var nextForThisInst []backtest.OrderIntent

		for _, order := range instrumentPending {
			decision := backtest.CheckRisk(account, order)
			if !decision.Allowed {
				newRejected = append(newRejected, RejectedOrder{
					Order:      order,
					Reason:     decision.Reason,
					BarTime:    bar.Timestamp,
					Instrument: order.InstrumentID,
					CashBefore: account,
				})
				continue
			}

			fill, ok, err := backtest.FillOrderOnDailyBar(order, bar)
			if err != nil {
				return nil, fmt.Errorf("fill order on bar at %s: %w", bar.Timestamp, err)
			}
			if !ok {
				// Limit not crossed, carry forward for this instrument
				nextForThisInst = append(nextForThisInst, order)
				continue
			}

			nextPortfolio, err := account.Portfolio.ApplyFill(fill)
			if err != nil {
				// ApplyFill failed (e.g. insufficient cash for buy, insufficient
				// position for sell). Record as rejected and discard from pending
				// to avoid infinite retry across bars.
				newRejected = append(newRejected, RejectedOrder{
					Order:      order,
					Reason:     err.Error(),
					BarTime:    bar.Timestamp,
					Instrument: order.InstrumentID,
					CashBefore: account,
				})
				continue
			}
			account.Portfolio = nextPortfolio
			newFills = append(newFills, fill)
		}

		rejected = append(rejected, newRejected...)
		// Only merge remaining pending orders for this instrument back if any
		if len(nextForThisInst) > 0 {
			pendingOrders[bar.InstrumentID] = nextForThisInst
		}
		fills = append(fills, newFills...)

		// --- Phase 2: Run strategy with current account state (after fills) ---
		input := backtest.StrategyInput{
			Run:       req.Run,
			Bar:       bar,
			Portfolio: account.Portfolio,
			History:   nil,
		}

		strategyOrders, err := strategy.Decide(input)
		if err != nil {
			return nil, fmt.Errorf("strategy decide failed at %s: %w", bar.Timestamp, err)
		}

		// Store strategy decisions in instrument-scoped pending keyed by each
		// order's own InstrumentID so that a strategy can return orders for an
		// instrument different from the current bar. Empty InstrumentID is treated
		// as a strategy bug: the order is rejected immediately with no retry so the
		// fault is visible early.
		for _, order := range strategyOrders {
			if order.InstrumentID == "" {
				rejected = append(rejected, RejectedOrder{
					Order:      order,
					Reason:     "empty order.InstrumentID",
					BarTime:    bar.Timestamp,
					Instrument: bar.InstrumentID,
					CashBefore: account,
				})
				continue
			}
			pendingOrders[order.InstrumentID] = append(pendingOrders[order.InstrumentID], order)
		}

		// --- Phase 3: Mark current position price and record equity ---
		if _, ok := account.Portfolio.Position(bar.InstrumentID); ok {
			account.Portfolio, err = account.Portfolio.MarkPrice(bar.InstrumentID, bar.Close)
			if err != nil {
				return nil, fmt.Errorf("failed to mark price for %s: %w", bar.InstrumentID, err)
			}
		}

		equity, err := account.Portfolio.Equity()
		if err != nil {
			return nil, fmt.Errorf("failed to calculate equity at %s: %w", bar.Timestamp, err)
		}
		equityCurve = append(equityCurve, backtest.EquityPoint{
			Timestamp: bar.Timestamp,
			Equity:    equity,
		})

		account.UpdatedAt = bar.Timestamp
	}

	return &Snapshot{
		Account:     account, // req.Account.ID is now preserved
		Fills:       fills,
		Rejected:    rejected,
		EquityCurve: equityCurve,
	}, nil
}