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

package socket

import (
	"context"
	"errors"
	"time"

	protoSocket "git.toki-labs.com/toki/proto-socket/go"

	altv1 "git.toki-labs.com/toki/alt/packages/contracts/gen/go/alt/v1"
	"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/marketdata/importer"
	"git.toki-labs.com/toki/alt/services/worker/internal/scheduler"
	"git.toki-labs.com/toki/alt/services/worker/internal/storage"
)

// handlerTimeout bounds a single worker-side backtest request. Execution itself
// is detached through the starter, so this only guards the synchronous record
// and read work against a wedged store call.
const handlerTimeout = 10 * time.Second

const (
	backtestErrorInvalidRequest = "invalid_request"
	backtestErrorUnavailable    = "unavailable"
	backtestErrorNotFound       = "not_found"
	backtestErrorTimeout        = "timeout"
	backtestErrorInternal       = "internal"
)

// ErrMonthlySourceBarsNotFound marks a monthly aggregation request whose
// selected instruments have no stored daily source bars in the requested range.
var ErrMonthlySourceBarsNotFound = errors.New("monthly aggregation source daily bars not found")

// BacktestStarter records a pending run and schedules its execution. The worker
// jobs package provides the concrete implementation; the interface keeps the
// socket layer testable with a fake.
type BacktestStarter interface {
	StartBacktest(ctx context.Context, spec backtest.RunSpec) (backtest.Run, error)
}

// DailyBarImporter runs a daily bar import command. The marketdata importer
// provides the concrete implementation; the narrow interface keeps the socket
// handler testable with a fake and decoupled from provider/storage wiring.
type DailyBarImporter interface {
	ImportDailyBars(ctx context.Context, request importer.DailyBarRequest) (importer.Result, error)
}

// MonthlyBarAggregator aggregates daily bars into monthly bars and stores them.
// The contract layer provides the concrete implementation; the narrow interface
// keeps the socket handler testable with a fake.
type MonthlyBarAggregator interface {
	AggregateMonthlyBars(ctx context.Context, request AggregateMonthlyRequest) (AggregateMonthlyResult, error)
}

// RefreshStatusQuerier provides read access to the runtime refresh status
// recorded by the scheduler after each tick. The socket handler calls All or
// Get to answer operator queries without importing the scheduler package.
type RefreshStatusQuerier interface {
	All() []scheduler.RefreshStatus
	Get(name string) (scheduler.RefreshStatus, bool)
}

// AggregateMonthlyRequest carries the filter/range for a monthly aggregation.
type AggregateMonthlyRequest struct {
	Provider market.Provider
	Selector market.UniverseSelector
	From     time.Time
	To       time.Time
}

// AggregateMonthlyResult holds the outcome of a monthly bar aggregation.
type AggregateMonthlyResult struct {
	Instruments      int
	SourceDailyCount int
	MonthlyBarCount  int
	Provenance       []MonthlyProvenanceEntry
}

// MonthlyProvenanceEntry records provenance for one instrument's monthly bars.
type MonthlyProvenanceEntry struct {
	InstrumentID      market.InstrumentID
	SourceDailyCount  int
	MonthlyBarCount   int
	SourceStart       time.Time
	SourceEnd         time.Time
	AggregationRuleID string
}

// Deps bundles the worker-owned capabilities socket handlers need. A nil field
// disables only its related handler surface so the server can still run a
// reduced surface (e.g. hello-only) without panicking.
type Deps struct {
	Starter           BacktestStarter
	Analysis          storage.BacktestAnalysisStore
	Results           storage.BacktestResultStore
	Instruments       storage.InstrumentStore
	Bars              storage.BarStore
	DailyBarImporter  DailyBarImporter
	MonthlyAggregator MonthlyBarAggregator
	RefreshStatus     RefreshStatusQuerier
	Paper             PaperService
	Live              LiveService
}

// BacktestDeps is retained for existing call sites while the socket dependency
// bundle grows beyond backtest into market reads.
type BacktestDeps = Deps

// backtestHandlers returns the session handlers for the backtest command and
// query surface, each closing over the shared dependencies.
func backtestHandlers(deps Deps) []sessionHandler {
	return []sessionHandler{
		{
			requestType: protoSocket.TypeNameOf(&altv1.StartBacktestRequest{}),
			register: func(client *protoSocket.WsClient) {
				protoSocket.AddRequestListenerTyped[*altv1.StartBacktestRequest, *altv1.StartBacktestResponse](
					&client.Communicator,
					func(req *altv1.StartBacktestRequest) (*altv1.StartBacktestResponse, error) {
						return handleStartBacktest(deps, req)
					},
				)
			},
		},
		{
			requestType: protoSocket.TypeNameOf(&altv1.GetBacktestRunRequest{}),
			register: func(client *protoSocket.WsClient) {
				protoSocket.AddRequestListenerTyped[*altv1.GetBacktestRunRequest, *altv1.GetBacktestRunResponse](
					&client.Communicator,
					func(req *altv1.GetBacktestRunRequest) (*altv1.GetBacktestRunResponse, error) {
						return handleGetBacktestRun(deps, req)
					},
				)
			},
		},
		{
			requestType: protoSocket.TypeNameOf(&altv1.ListBacktestRunsRequest{}),
			register: func(client *protoSocket.WsClient) {
				protoSocket.AddRequestListenerTyped[*altv1.ListBacktestRunsRequest, *altv1.ListBacktestRunsResponse](
					&client.Communicator,
					func(req *altv1.ListBacktestRunsRequest) (*altv1.ListBacktestRunsResponse, error) {
						return handleListBacktestRuns(deps, req)
					},
				)
			},
		},
		{
			requestType: protoSocket.TypeNameOf(&altv1.GetBacktestRunDetailRequest{}),
			register: func(client *protoSocket.WsClient) {
				protoSocket.AddRequestListenerTyped[*altv1.GetBacktestRunDetailRequest, *altv1.GetBacktestRunDetailResponse](
					&client.Communicator,
					func(req *altv1.GetBacktestRunDetailRequest) (*altv1.GetBacktestRunDetailResponse, error) {
						return handleGetBacktestRunDetail(deps, req)
					},
				)
			},
		},
		{
			requestType: protoSocket.TypeNameOf(&altv1.GetBacktestResultRequest{}),
			register: func(client *protoSocket.WsClient) {
				protoSocket.AddRequestListenerTyped[*altv1.GetBacktestResultRequest, *altv1.GetBacktestResultResponse](
					&client.Communicator,
					func(req *altv1.GetBacktestResultRequest) (*altv1.GetBacktestResultResponse, error) {
						return handleGetBacktestResult(deps, req)
					},
				)
			},
		},
		{
			requestType: protoSocket.TypeNameOf(&altv1.CompareBacktestRunsRequest{}),
			register: func(client *protoSocket.WsClient) {
				protoSocket.AddRequestListenerTyped[*altv1.CompareBacktestRunsRequest, *altv1.CompareBacktestRunsResponse](
					&client.Communicator,
					func(req *altv1.CompareBacktestRunsRequest) (*altv1.CompareBacktestRunsResponse, error) {
						return handleCompareBacktestRuns(deps, req)
					},
				)
			},
		},
	}
}

func handleStartBacktest(deps Deps, req *altv1.StartBacktestRequest) (*altv1.StartBacktestResponse, error) {
	spec, err := runSpecFromProto(req.GetSpec())
	if err != nil {
		return &altv1.StartBacktestResponse{Error: invalidRequest(err.Error())}, nil
	}
	if deps.Starter == nil {
		return &altv1.StartBacktestResponse{Error: unavailableError("backtest start is not available")}, nil
	}

	ctx, cancel := context.WithTimeout(context.Background(), handlerTimeout)
	defer cancel()

	run, err := deps.Starter.StartBacktest(ctx, spec)
	if err != nil {
		return &altv1.StartBacktestResponse{Error: backendErrorInfo(err)}, nil
	}
	return &altv1.StartBacktestResponse{Run: runToProto(run)}, nil
}

// handleGetBacktestRun answers a single-run poll. It reads through the analysis
// store's run detail so polling reuses the same store-backed read path as
// list/detail instead of growing a second run-store dependency on Deps.
func handleGetBacktestRun(deps Deps, req *altv1.GetBacktestRunRequest) (*altv1.GetBacktestRunResponse, error) {
	if req.GetRunId() == "" {
		return &altv1.GetBacktestRunResponse{Error: invalidRequest("run_id is required")}, nil
	}
	if deps.Analysis == nil {
		return &altv1.GetBacktestRunResponse{Error: unavailableError("backtest analysis is not available")}, nil
	}

	ctx, cancel := context.WithTimeout(context.Background(), handlerTimeout)
	defer cancel()

	detail, err := deps.Analysis.GetRunDetail(ctx, backtest.RunID(req.GetRunId()))
	if err != nil {
		return &altv1.GetBacktestRunResponse{Error: backendErrorInfo(err)}, nil
	}
	return &altv1.GetBacktestRunResponse{Run: runToProto(detail.Run)}, nil
}

func handleListBacktestRuns(deps Deps, req *altv1.ListBacktestRunsRequest) (*altv1.ListBacktestRunsResponse, error) {
	if deps.Analysis == nil {
		return &altv1.ListBacktestRunsResponse{Error: unavailableError("backtest analysis is not available")}, nil
	}

	ctx, cancel := context.WithTimeout(context.Background(), handlerTimeout)
	defer cancel()

	runs, err := deps.Analysis.ListRuns(ctx, runStatusFromProto(req.GetStatus()))
	if err != nil {
		return &altv1.ListBacktestRunsResponse{Error: backendErrorInfo(err)}, nil
	}
	return &altv1.ListBacktestRunsResponse{Runs: runsToProto(runs)}, nil
}

func handleGetBacktestRunDetail(deps Deps, req *altv1.GetBacktestRunDetailRequest) (*altv1.GetBacktestRunDetailResponse, error) {
	if req.GetRunId() == "" {
		return &altv1.GetBacktestRunDetailResponse{Error: invalidRequest("run_id is required")}, nil
	}
	if deps.Analysis == nil {
		return &altv1.GetBacktestRunDetailResponse{Error: unavailableError("backtest analysis is not available")}, nil
	}

	ctx, cancel := context.WithTimeout(context.Background(), handlerTimeout)
	defer cancel()

	detail, err := deps.Analysis.GetRunDetail(ctx, backtest.RunID(req.GetRunId()))
	if err != nil {
		return &altv1.GetBacktestRunDetailResponse{Error: backendErrorInfo(err)}, nil
	}

	resp := &altv1.GetBacktestRunDetailResponse{Run: runToProto(detail.Run)}
	// A run can exist before it produces a result; leave the result field nil in
	// that case so callers can distinguish "no result yet" from an empty result.
	if detail.HasResult {
		resp.Result = resultToProto(detail.Result)
	}
	return resp, nil
}

func handleGetBacktestResult(deps Deps, req *altv1.GetBacktestResultRequest) (*altv1.GetBacktestResultResponse, error) {
	if req.GetRunId() == "" {
		return &altv1.GetBacktestResultResponse{Error: invalidRequest("run_id is required")}, nil
	}
	if deps.Results == nil {
		return &altv1.GetBacktestResultResponse{Error: unavailableError("backtest results are not available")}, nil
	}

	ctx, cancel := context.WithTimeout(context.Background(), handlerTimeout)
	defer cancel()

	result, err := deps.Results.GetResult(ctx, backtest.RunID(req.GetRunId()))
	if err != nil {
		return &altv1.GetBacktestResultResponse{Error: backendErrorInfo(err)}, nil
	}
	return &altv1.GetBacktestResultResponse{Result: resultToProto(result)}, nil
}

func handleCompareBacktestRuns(deps Deps, req *altv1.CompareBacktestRunsRequest) (*altv1.CompareBacktestRunsResponse, error) {
	ids := make([]backtest.RunID, 0, len(req.GetRunIds()))
	for _, id := range req.GetRunIds() {
		if id == "" {
			return &altv1.CompareBacktestRunsResponse{Error: invalidRequest("run_ids must not contain empty values")}, nil
		}
		ids = append(ids, backtest.RunID(id))
	}
	// An empty compare request is a no-op rather than an error: the response just
	// carries an empty result list.
	if len(ids) == 0 {
		return &altv1.CompareBacktestRunsResponse{}, nil
	}
	if deps.Analysis == nil {
		return &altv1.CompareBacktestRunsResponse{Error: unavailableError("backtest analysis is not available")}, nil
	}

	ctx, cancel := context.WithTimeout(context.Background(), handlerTimeout)
	defer cancel()

	results, err := deps.Analysis.CompareResults(ctx, ids)
	if err != nil {
		return &altv1.CompareBacktestRunsResponse{Error: backendErrorInfo(err)}, nil
	}
	return &altv1.CompareBacktestRunsResponse{Results: resultsToProto(results)}, nil
}

func invalidRequest(reason string) *altv1.ErrorInfo {
	return errorInfo(backtestErrorInvalidRequest, "invalid backtest request: "+reason)
}

func unavailableError(message string) *altv1.ErrorInfo {
	return errorInfo(backtestErrorUnavailable, message)
}

func backendErrorInfo(err error) *altv1.ErrorInfo {
	if err == nil {
		return nil
	}
	switch {
	case errors.Is(err, storage.ErrRunNotFound), errors.Is(err, storage.ErrResultNotFound):
		return errorInfo(backtestErrorNotFound, err.Error())
	case errors.Is(err, context.DeadlineExceeded):
		return errorInfo(backtestErrorTimeout, err.Error())
	default:
		return errorInfo(backtestErrorInternal, err.Error())
	}
}

func errorInfo(code, message string) *altv1.ErrorInfo {
	return &altv1.ErrorInfo{Code: code, Message: message}
}