alt/services/worker/internal/livetrading/service.go

package livetrading

import (
	"context"
	"errors"
	"fmt"
	"math/big"
	"strings"
	"sync"
	"time"

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

// BrokerPort is the abstract boundary to a live broker adapter. When nil, the
// Service treats all order operations as unavailable/disabled.
type BrokerPort interface {
	GetCapabilities(ctx context.Context, accountID string) (trading.BrokerCapability, error)
	SubmitOrder(ctx context.Context, accountID string, intent trading.OrderIntent) (OrderResult, error)
	CancelOrder(ctx context.Context, accountID, brokerOrderID string) (OrderResult, error)
	GetOrderStatus(ctx context.Context, accountID, brokerOrderID string) (OrderResult, error)
	// FetchAccountSnapshot retrieves a sanitized account snapshot from the broker.
	// Raw account numbers and provider payload must be stripped inside the adapter
	// before returning; the snapshot AccountID must be an operator alias or internal id.
	FetchAccountSnapshot(ctx context.Context, accountID string) (trading.AccountSnapshot, error)
}

// OrderResult carries the broker-reported outcome of an order operation.
type OrderResult struct {
	BrokerOrderID   string
	Status          trading.OrderStatus
	BrokerStatus    trading.BrokerStatus
	RejectionReason string
}

// OperatorConfirmation is the explicit operator gate required before any live
// order reaches the broker.
type OperatorConfirmation struct {
	Confirmed     bool
	OperatorID    string
	Reason        string
	ConfirmedAtMs int64
}

// SubmitRequest carries the validated parameters for submitting a live order.
type SubmitRequest struct {
	AccountID      string
	Intent         trading.OrderIntent
	Confirmation   OperatorConfirmation
	IdempotencyKey string
}

// LiveOrder is the in-memory record for a submitted live order, holding both
// the ALT-domain status and the raw broker status so neither is discarded.
type LiveOrder struct {
	ID              string
	AccountID       string
	Intent          trading.OrderIntent
	Status          trading.OrderStatus
	BrokerStatus    trading.BrokerStatus
	BrokerOrderID   string
	RejectionReason string
	CreatedAt       time.Time
	UpdatedAt       time.Time
}

var (
	// ErrConfirmationRequired is returned when SubmitLiveOrder is called without
	// a confirmed OperatorConfirmation.
	ErrConfirmationRequired = errors.New("operator confirmation required")
	// ErrBrokerUnavailable is returned when the BrokerPort is nil (disabled) or
	// not reachable.
	ErrBrokerUnavailable = errors.New("live broker is not available")
	// ErrOrderNotFound is returned when the requested live order id does not
	// exist in the in-memory registry for the given account.
	ErrOrderNotFound = errors.New("live order not found")
	// ErrInvalidIntent is returned when the live order intent is malformed:
	// missing instrument_id, invalid side, non-positive quantity, or missing type.
	ErrInvalidIntent = errors.New("invalid live order intent")
	// ErrRiskBlocked is returned when the kill switch is halted or the order
	// intent violates the active risk policy. The broker is never called.
	ErrRiskBlocked = errors.New("live order blocked by risk policy")
	// ErrAccountSnapshotNotFound is returned when GetLiveAccountSnapshot is
	// called for an account that has not yet been synced.
	ErrAccountSnapshotNotFound = errors.New("live account snapshot not found")
	// ErrAuditFailed is returned when an audit event cannot be persisted. The
	// operation itself may have succeeded, but the audit trail is mandatory.
	ErrAuditFailed = errors.New("live audit event failed to persist")
)

// AuditStore is the narrow interface the Service uses to persist and query live
// operation audit events. storage.LiveAuditStore satisfies this interface.
type AuditStore interface {
	AppendLiveAuditEvent(ctx context.Context, event trading.AuditEvent) error
	ListLiveAuditEvents(ctx context.Context, filter trading.AuditFilter) ([]trading.AuditEvent, error)
}

// Service is the worker-owned live trading runtime. It keeps an in-memory order
// registry and delegates execution to a BrokerPort. A nil BrokerPort means live
// trading is disabled; all order operations return ErrBrokerUnavailable.
type Service struct {
	broker     BrokerPort
	mu         sync.RWMutex
	orders     map[string]*LiveOrder
	seq        int64
	policyMu   sync.RWMutex
	policy     trading.RiskPolicy
	killSwitch trading.KillSwitchState

	acctMuMu sync.Mutex             // protects acctMus
	acctMus  map[string]*sync.Mutex // per-account submit serialization

	snapshotMu sync.RWMutex
	snapshots  map[string]*trading.AccountSnapshot // keyed by accountID

	auditStore AuditStore // optional; nil means no-op audit trail
}

// NewService creates a Service backed by broker with the default risk policy
// and a halted kill switch. Pass nil broker to operate in disabled/unavailable
// mode (broker capability probe and order operations all return ErrBrokerUnavailable).
func NewService(broker BrokerPort) *Service {
	return &Service{
		broker:     broker,
		orders:     make(map[string]*LiveOrder),
		policy:     DefaultRiskPolicy(),
		killSwitch: trading.DefaultKillSwitch(),
		acctMus:    make(map[string]*sync.Mutex),
		snapshots:  make(map[string]*trading.AccountSnapshot),
	}
}

// SetAuditStore injects an optional audit store for durable audit trail. When
// nil (the default), audit events are silently discarded.
func (s *Service) SetAuditStore(store AuditStore) {
	s.auditStore = store
}

// newAuditEventID returns a process-unique audit event identifier.
func newAuditEventID() string {
	return fmt.Sprintf("aud-%d", time.Now().UnixNano())
}

// appendAudit persists ev to the audit store when one is configured. Returns
// ErrAuditFailed (wrapping the store error) on failure so callers can propagate
// it as a typed internal error. Forbidden payload keys (raw secrets, tokens,
// account numbers) are rejected at this boundary before reaching the store.
func (s *Service) appendAudit(ctx context.Context, ev trading.AuditEvent) error {
	if s.auditStore == nil {
		return nil
	}
	if trading.HasRawSecretKey(ev.Payload) {
		return fmt.Errorf("%w: audit payload contains forbidden key", ErrAuditFailed)
	}
	if err := s.auditStore.AppendLiveAuditEvent(ctx, ev); err != nil {
		return fmt.Errorf("%w: %v", ErrAuditFailed, err)
	}
	return nil
}

// ListLiveAuditEvents delegates to the configured AuditStore. Returns nil when
// no audit store is configured.
func (s *Service) ListLiveAuditEvents(ctx context.Context, filter trading.AuditFilter) ([]trading.AuditEvent, error) {
	if s.auditStore == nil {
		return nil, nil
	}
	return s.auditStore.ListLiveAuditEvents(ctx, filter)
}

// submitLockForAccount returns the per-account mutex for submit serialization,
// creating it if it does not yet exist.
func (s *Service) submitLockForAccount(accountID string) *sync.Mutex {
	s.acctMuMu.Lock()
	mu, ok := s.acctMus[accountID]
	if !ok {
		mu = &sync.Mutex{}
		s.acctMus[accountID] = mu
	}
	s.acctMuMu.Unlock()
	return mu
}

// GetRiskPolicy returns the current risk policy.
func (s *Service) GetRiskPolicy() trading.RiskPolicy {
	s.policyMu.RLock()
	defer s.policyMu.RUnlock()
	return s.policy
}

// SetRiskPolicy replaces the active risk policy.
func (s *Service) SetRiskPolicy(policy trading.RiskPolicy) {
	s.policyMu.Lock()
	s.policy = policy
	s.policyMu.Unlock()
}

// GetKillSwitch returns the current kill switch state.
func (s *Service) GetKillSwitch() trading.KillSwitchState {
	s.policyMu.RLock()
	defer s.policyMu.RUnlock()
	return s.killSwitch
}

// SetKillSwitch replaces the active kill switch state and returns the new state.
// The kill_switch_set audit event is emitted best-effort (failure is silently
// discarded) because the kill switch must succeed even when audit storage is
// unavailable.
func (s *Service) SetKillSwitch(ks trading.KillSwitchState) trading.KillSwitchState {
	s.policyMu.Lock()
	s.killSwitch = ks
	s.policyMu.Unlock()
	now := time.Now().UTC()
	_ = s.appendAudit(context.Background(), trading.AuditEvent{
		EventID:   newAuditEventID(),
		Type:      trading.AuditEventTypeKillSwitchSet,
		Timestamp: now,
		Status:    fmt.Sprintf("halted=%v", ks.Halted),
		Reason:    ks.Reason,
		Payload:   map[string]string{"halted": fmt.Sprintf("%v", ks.Halted)},
	})
	return ks
}

// GetLiveBrokerCapabilities proxies the capability probe to the BrokerPort.
func (s *Service) GetLiveBrokerCapabilities(ctx context.Context, accountID string) (trading.BrokerCapability, error) {
	if s.broker == nil {
		return trading.BrokerCapability{}, ErrBrokerUnavailable
	}
	return s.broker.GetCapabilities(ctx, accountID)
}

// validateIntent checks the minimum required fields of a live order intent so
// malformed requests never reach the broker port even when operator-confirmed.
func validateIntent(intent trading.OrderIntent) error {
	if intent.InstrumentID == "" {
		return fmt.Errorf("%w: instrument_id is required", ErrInvalidIntent)
	}
	if intent.Side != trading.OrderSideBuy && intent.Side != trading.OrderSideSell {
		return fmt.Errorf("%w: side must be buy or sell", ErrInvalidIntent)
	}
	qty := strings.TrimSpace(intent.Quantity.Amount.Value)
	if qty == "" {
		return fmt.Errorf("%w: quantity must be greater than zero", ErrInvalidIntent)
	}
	qtyRat := new(big.Rat)
	if _, ok := qtyRat.SetString(qty); !ok {
		return fmt.Errorf("%w: quantity is not a valid decimal", ErrInvalidIntent)
	}
	if qtyRat.Sign() <= 0 {
		return fmt.Errorf("%w: quantity must be greater than zero", ErrInvalidIntent)
	}
	if intent.Type == "" {
		return fmt.Errorf("%w: order type is required", ErrInvalidIntent)
	}
	if lp := strings.TrimSpace(intent.LimitPrice.Amount.Value); lp != "" {
		lpRat := new(big.Rat)
		if _, ok := lpRat.SetString(lp); !ok {
			return fmt.Errorf("%w: limit_price is not a valid decimal", ErrInvalidIntent)
		}
		if lpRat.Sign() < 0 {
			return fmt.Errorf("%w: limit_price must be non-negative", ErrInvalidIntent)
		}
	}
	return nil
}

// SubmitLiveOrder validates the operator confirmation gate and the order
// intent shape, then delegates order submission to the BrokerPort before
// recording the order in the in-memory registry. The order type in
// req.Intent.Type is preserved verbatim so custom broker strings are not
// narrowed to market/limit.
func (s *Service) SubmitLiveOrder(ctx context.Context, req SubmitRequest) (LiveOrder, error) {
	if !req.Confirmation.Confirmed {
		return LiveOrder{}, ErrConfirmationRequired
	}
	if req.AccountID == "" {
		return LiveOrder{}, fmt.Errorf("%w: account_id is required", ErrInvalidIntent)
	}
	if err := validateIntent(req.Intent); err != nil {
		return LiveOrder{}, err
	}

	// Nil broker check before acquiring account mutex so ErrBrokerUnavailable
	// takes precedence over risk evaluation regardless of kill switch state.
	if s.broker == nil {
		return LiveOrder{}, ErrBrokerUnavailable
	}

	// Serialize same-account submits from risk evaluation through broker call and
	// order storage. Different accounts proceed independently.
	acctMu := s.submitLockForAccount(req.AccountID)
	acctMu.Lock()
	defer acctMu.Unlock()

	// Snapshot policy and kill switch after acquiring the account mutex so
	// queued submits see the current state rather than a stale pre-wait copy.
	s.policyMu.RLock()
	policy := s.policy
	ks := s.killSwitch
	s.policyMu.RUnlock()

	now := time.Now().UTC()
	s.mu.RLock()
	usage := s.riskUsageForAccount(req.AccountID, now)
	s.mu.RUnlock()

	if decision := trading.EvaluateRisk(policy, ks, req.Intent, usage); !decision.Allowed {
		order, riskErr := s.recordRiskRejectedOrder(ctx, req, decision, now)
		if auditErr := s.appendAudit(ctx, trading.AuditEvent{
			EventID:   newAuditEventID(),
			Type:      trading.AuditEventTypeRiskBlocked,
			Timestamp: now,
			AccountID: req.AccountID,
			OrderID:   order.ID,
			Status:    string(trading.OrderStatusRejected),
			Reason:    decision.Reason,
			Actor:     req.Confirmation.OperatorID,
			Payload:   map[string]string{"instrument_id": string(req.Intent.InstrumentID), "reason": decision.Reason},
		}); auditErr != nil {
			return order, auditErr
		}
		return order, riskErr
	}

	// Emit submit_confirmed before calling the broker so the operator gate
	// decision is recorded even if the broker call later fails.
	if auditErr := s.appendAudit(ctx, trading.AuditEvent{
		EventID:   newAuditEventID(),
		Type:      trading.AuditEventTypeSubmitConfirmed,
		Timestamp: now,
		AccountID: req.AccountID,
		Actor:     req.Confirmation.OperatorID,
		Payload: map[string]string{
			"instrument_id": string(req.Intent.InstrumentID),
			"side":          string(req.Intent.Side),
			"type":          string(req.Intent.Type),
		},
	}); auditErr != nil {
		return LiveOrder{}, auditErr
	}

	result, err := s.broker.SubmitOrder(ctx, req.AccountID, req.Intent)
	if err != nil {
		return LiveOrder{}, err
	}

	s.mu.Lock()
	s.seq++
	id := fmt.Sprintf("lo-%s-%d", req.AccountID, s.seq)
	order := &LiveOrder{
		ID:              id,
		AccountID:       req.AccountID,
		Intent:          req.Intent,
		Status:          result.Status,
		BrokerStatus:    result.BrokerStatus,
		BrokerOrderID:   result.BrokerOrderID,
		RejectionReason: result.RejectionReason,
		CreatedAt:       now,
		UpdatedAt:       now,
	}
	s.orders[id] = order
	s.mu.Unlock()

	// Emit broker_accepted or broker_rejected to record the broker's decision.
	brokerEvType := trading.AuditEventTypeBrokerAccepted
	if order.Status == trading.OrderStatusRejected {
		brokerEvType = trading.AuditEventTypeBrokerRejected
	}
	if auditErr := s.appendAudit(ctx, trading.AuditEvent{
		EventID:   newAuditEventID(),
		Type:      brokerEvType,
		Timestamp: order.CreatedAt,
		AccountID: order.AccountID,
		OrderID:   order.ID,
		Status:    string(order.Status),
		Reason:    order.RejectionReason,
		Actor:     req.Confirmation.OperatorID,
		Payload: map[string]string{
			"broker_order_id": order.BrokerOrderID,
		},
	}); auditErr != nil {
		return *order, auditErr
	}
	return *order, nil
}

// isTerminalStatus returns true when status represents a final order state.
// Terminal orders are excluded from the open order count for risk evaluation.
func isTerminalStatus(status trading.OrderStatus) bool {
	switch status {
	case trading.OrderStatusFilled, trading.OrderStatusCanceled, trading.OrderStatusRejected:
		return true
	}
	return false
}

// riskUsageForAccount returns the current daily and open order counts for
// accountID. Callers must hold s.mu (at least RLock) while calling this.
func (s *Service) riskUsageForAccount(accountID string, now time.Time) trading.RiskUsage {
	todayUTC := now.Truncate(24 * time.Hour)
	var daily, open int
	for _, o := range s.orders {
		if o.AccountID != accountID {
			continue
		}
		if !o.CreatedAt.Before(todayUTC) {
			daily++
		}
		if !isTerminalStatus(o.Status) {
			open++
		}
	}
	return trading.RiskUsage{DailyOrderCount: daily, OpenOrderCount: open}
}

// recordRiskRejectedOrder creates an in-memory rejected live order for a
// risk-blocked submission, stores it, and returns both the order and the
// typed ErrRiskBlocked error so the socket layer can surface both.
func (s *Service) recordRiskRejectedOrder(ctx context.Context, req SubmitRequest, decision trading.RiskDecision, now time.Time) (LiveOrder, error) {
	s.mu.Lock()
	s.seq++
	id := fmt.Sprintf("lo-%s-%d", req.AccountID, s.seq)
	order := &LiveOrder{
		ID:              id,
		AccountID:       req.AccountID,
		Intent:          req.Intent,
		Status:          trading.OrderStatusRejected,
		RejectionReason: decision.Reason,
		CreatedAt:       now,
		UpdatedAt:       now,
	}
	s.orders[id] = order
	s.mu.Unlock()
	return *order, fmt.Errorf("%w: %s", ErrRiskBlocked, decision.Reason)
}

// CancelLiveOrder asks the BrokerPort to cancel the order identified by
// liveOrderID and updates the in-memory status with the broker response.
func (s *Service) CancelLiveOrder(ctx context.Context, accountID, liveOrderID string) (LiveOrder, error) {
	if s.broker == nil {
		return LiveOrder{}, ErrBrokerUnavailable
	}

	s.mu.RLock()
	order, ok := s.orders[liveOrderID]
	s.mu.RUnlock()

	if !ok || order.AccountID != accountID {
		return LiveOrder{}, ErrOrderNotFound
	}

	result, err := s.broker.CancelOrder(ctx, accountID, order.BrokerOrderID)
	if err != nil {
		return LiveOrder{}, err
	}

	s.mu.Lock()
	order.Status = result.Status
	order.BrokerStatus = result.BrokerStatus
	order.UpdatedAt = time.Now().UTC()
	updatedAt := order.UpdatedAt
	s.mu.Unlock()

	if auditErr := s.appendAudit(ctx, trading.AuditEvent{
		EventID:   newAuditEventID(),
		Type:      trading.AuditEventTypeCancelConfirmed,
		Timestamp: updatedAt,
		AccountID: accountID,
		OrderID:   liveOrderID,
		Status:    string(order.Status),
		Payload:   map[string]string{"broker_order_id": order.BrokerOrderID},
	}); auditErr != nil {
		return *order, auditErr
	}
	return *order, nil
}

// SyncLiveAccount fetches a fresh account snapshot from the broker, stores it
// in memory, and returns it. A nil broker returns ErrBrokerUnavailable.
// The snapshot is sanitized by the BrokerPort adapter; no raw account numbers
// or provider payload reach this layer.
func (s *Service) SyncLiveAccount(ctx context.Context, accountID string) (trading.AccountSnapshot, error) {
	if s.broker == nil {
		return trading.AccountSnapshot{}, ErrBrokerUnavailable
	}
	snap, err := s.broker.FetchAccountSnapshot(ctx, accountID)
	if err != nil {
		// Emit failure audit best-effort; don't mask the original broker error.
		_ = s.appendAudit(ctx, trading.AuditEvent{
			EventID:   newAuditEventID(),
			Type:      trading.AuditEventTypeAccountSyncFailure,
			Timestamp: time.Now().UTC(),
			AccountID: accountID,
			Status:    "failed",
			Reason:    err.Error(),
		})
		return trading.AccountSnapshot{}, err
	}
	s.snapshotMu.Lock()
	s.snapshots[accountID] = &snap
	s.snapshotMu.Unlock()
	if auditErr := s.appendAudit(ctx, trading.AuditEvent{
		EventID:   newAuditEventID(),
		Type:      trading.AuditEventTypeAccountSyncSuccess,
		Timestamp: snap.SyncedAt,
		Broker:    snap.Broker,
		AccountID: accountID,
		Status:    "synced",
		Payload:   map[string]string{"cash_count": fmt.Sprintf("%d", len(snap.Cash)), "position_count": fmt.Sprintf("%d", len(snap.Positions))},
	}); auditErr != nil {
		return trading.AccountSnapshot{}, auditErr
	}
	return snap, nil
}

// GetLiveAccountSnapshot returns the last synced account snapshot for
// accountID. Returns ErrAccountSnapshotNotFound when no sync has been
// performed yet for that account.
func (s *Service) GetLiveAccountSnapshot(ctx context.Context, accountID string) (trading.AccountSnapshot, error) {
	s.snapshotMu.RLock()
	snap, ok := s.snapshots[accountID]
	s.snapshotMu.RUnlock()
	if !ok {
		return trading.AccountSnapshot{}, ErrAccountSnapshotNotFound
	}
	return *snap, nil
}

// GetLiveOrder fetches the current broker status for a live order and refreshes
// the in-memory record with the broker response before returning it.
func (s *Service) GetLiveOrder(ctx context.Context, accountID, liveOrderID string) (LiveOrder, error) {
	if s.broker == nil {
		return LiveOrder{}, ErrBrokerUnavailable
	}

	s.mu.RLock()
	order, ok := s.orders[liveOrderID]
	s.mu.RUnlock()

	if !ok || order.AccountID != accountID {
		return LiveOrder{}, ErrOrderNotFound
	}

	result, err := s.broker.GetOrderStatus(ctx, accountID, order.BrokerOrderID)
	if err != nil {
		return LiveOrder{}, err
	}

	s.mu.Lock()
	order.Status = result.Status
	order.BrokerStatus = result.BrokerStatus
	order.UpdatedAt = time.Now().UTC()
	s.mu.Unlock()

	return *order, nil
}