aitrade/pkg/db/trades.go

package db

import (
	"context"
	"fmt"
	"time"

	"github.com/pheinrich/aitrade/pkg/model"
)

type TradeRepository struct {
	db Database
}

func NewTradeRepository(db Database) *TradeRepository {
	return &TradeRepository{db: db}
}

func (r *TradeRepository) Create(ctx context.Context, trade *model.Trade) error {
	query := `
		INSERT INTO trades (
			symbol, action, quantity, status, confidence, reasoning,
			target_price, stop_loss_price, pending_until, created_at,
			is_dry_run
		) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
	`

	result, err := r.db.ExecContext(ctx, query,
		trade.Symbol,
		trade.Action,
		trade.Quantity,
		trade.Status,
		trade.Confidence,
		trade.Reasoning,
		trade.TargetPrice,
		trade.StopLossPrice,
		trade.PendingUtil,
		trade.CreatedAt,
		trade.IsDryRun,
	)
	if err != nil {
		return fmt.Errorf("failed to insert trade: %w", err)
	}

	id, err := result.LastInsertId()
	if err != nil {
		return fmt.Errorf("failed to get last insert id: %w", err)
	}

	trade.ID = id
	return nil
}

func (r *TradeRepository) Update(ctx context.Context, trade *model.Trade) error {
	query := `
		UPDATE trades SET
			status = ?,
			executed_price = ?,
			ib_order_id = ?,
			approved_at = ?,
			rejected_at = ?,
			submitted_at = ?,
			filled_at = ?,
			completed_at = ?,
			rejection_reason = ?,
			forced_by_user = ?,
			dry_run_pnl = ?
		WHERE id = ?
	`

	_, err := r.db.ExecContext(ctx, query,
		trade.Status,
		trade.ExecutedPrice,
		trade.IBOrderID,
		trade.ApprovedAt,
		trade.RejectedAt,
		trade.SubmittedAt,
		trade.FilledAt,
		trade.CompletedAt,
		trade.RejectionReason,
		trade.ForcedByUser,
		trade.DryRunPnL,
		trade.ID,
	)

	if err != nil {
		return fmt.Errorf("failed to update trade: %w", err)
	}

	return nil
}

func (r *TradeRepository) GetByID(ctx context.Context, id int64) (*model.Trade, error) {
	query := `
		SELECT id, symbol, action, quantity, status, confidence, reasoning,
			target_price, executed_price, stop_loss_price, ib_order_id,
			created_at, pending_until, approved_at, rejected_at, submitted_at,
			filled_at, completed_at, rejection_reason, forced_by_user,
			is_dry_run, dry_run_pnl
		FROM trades
		WHERE id = ?
	`

	var trade model.Trade
	err := r.db.QueryRowContext(ctx, query, id).Scan(
		&trade.ID,
		&trade.Symbol,
		&trade.Action,
		&trade.Quantity,
		&trade.Status,
		&trade.Confidence,
		&trade.Reasoning,
		&trade.TargetPrice,
		&trade.ExecutedPrice,
		&trade.StopLossPrice,
		&trade.IBOrderID,
		&trade.CreatedAt,
		&trade.PendingUtil,
		&trade.ApprovedAt,
		&trade.RejectedAt,
		&trade.SubmittedAt,
		&trade.FilledAt,
		&trade.CompletedAt,
		&trade.RejectionReason,
		&trade.ForcedByUser,
		&trade.IsDryRun,
		&trade.DryRunPnL,
	)

	if err != nil {
		return nil, fmt.Errorf("failed to get trade: %w", err)
	}

	return &trade, nil
}

func (r *TradeRepository) GetPendingTrades(ctx context.Context) ([]*model.Trade, error) {
	query := `
		SELECT id, symbol, action, quantity, status, confidence, reasoning,
			target_price, executed_price, stop_loss_price, ib_order_id,
			created_at, pending_until, approved_at, rejected_at, submitted_at,
			filled_at, completed_at, rejection_reason, forced_by_user,
			is_dry_run, dry_run_pnl
		FROM trades
		WHERE status = ?
		ORDER BY created_at ASC
	`

	return r.queryTrades(ctx, query, model.TradePending)
}

func (r *TradeRepository) GetExpiredPendingTrades(ctx context.Context) ([]*model.Trade, error) {
	query := `
		SELECT id, symbol, action, quantity, status, confidence, reasoning,
			target_price, executed_price, stop_loss_price, ib_order_id,
			created_at, pending_until, approved_at, rejected_at, submitted_at,
			filled_at, completed_at, rejection_reason, forced_by_user,
			is_dry_run, dry_run_pnl
		FROM trades
		WHERE status = ? AND pending_until <= ?
		ORDER BY created_at ASC
	`

	return r.queryTrades(ctx, query, model.TradePending, time.Now())
}

func (r *TradeRepository) GetActiveTrades(ctx context.Context) ([]*model.Trade, error) {
	query := `
		SELECT id, symbol, action, quantity, status, confidence, reasoning,
			target_price, executed_price, stop_loss_price, ib_order_id,
			created_at, pending_until, approved_at, rejected_at, submitted_at,
			filled_at, completed_at, rejection_reason, forced_by_user,
			is_dry_run, dry_run_pnl
		FROM trades
		WHERE status IN (?, ?, ?)
		ORDER BY created_at ASC
	`

	return r.queryTrades(ctx, query, model.TradeSubmitted, model.TradeFilled, model.TradeCompleted)
}

func (r *TradeRepository) GetTradesSince(ctx context.Context, since time.Time) ([]*model.Trade, error) {
	query := `
		SELECT id, symbol, action, quantity, status, confidence, reasoning,
			target_price, executed_price, stop_loss_price, ib_order_id,
			created_at, pending_until, approved_at, rejected_at, submitted_at,
			filled_at, completed_at, rejection_reason, forced_by_user,
			is_dry_run, dry_run_pnl
		FROM trades
		WHERE created_at >= ?
		ORDER BY created_at DESC
	`

	return r.queryTrades(ctx, query, since)
}

func (r *TradeRepository) CountTradesSince(ctx context.Context, since time.Time) (int, error) {
	query := `SELECT COUNT(*) FROM trades WHERE created_at >= ? AND status NOT IN (?, ?)`

	var count int
	err := r.db.QueryRowContext(ctx, query, since, model.TradePending, model.TradeRejected).Scan(&count)
	if err != nil {
		return 0, fmt.Errorf("failed to count trades: %w", err)
	}

	return count, nil
}

func (r *TradeRepository) queryTrades(ctx context.Context, query string, args ...any) ([]*model.Trade, error) {
	rows, err := r.db.QueryContext(ctx, query, args...)
	if err != nil {
		return nil, fmt.Errorf("failed to query trades: %w", err)
	}
	defer rows.Close()

	var trades []*model.Trade
	for rows.Next() {
		var trade model.Trade
		if err := rows.Scan(
			&trade.ID,
			&trade.Symbol,
			&trade.Action,
			&trade.Quantity,
			&trade.Status,
			&trade.Confidence,
			&trade.Reasoning,
			&trade.TargetPrice,
			&trade.ExecutedPrice,
			&trade.StopLossPrice,
			&trade.IBOrderID,
			&trade.CreatedAt,
			&trade.PendingUtil,
			&trade.ApprovedAt,
			&trade.RejectedAt,
			&trade.SubmittedAt,
			&trade.FilledAt,
			&trade.CompletedAt,
			&trade.RejectionReason,
			&trade.ForcedByUser,
			&trade.IsDryRun,
			&trade.DryRunPnL,
		); err != nil {
			return nil, fmt.Errorf("failed to scan trade: %w", err)
		}
		trades = append(trades, &trade)
	}

	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("error iterating trade rows: %w", err)
	}

	return trades, nil
}