aitrade/pkg/app/trader/trader.go

package trader

import (
	"context"
	"fmt"
	"log/slog"
	"math/rand"
	"strings"
	"time"

	"github.com/pheinrich/aitrade/pkg/app/client"
	"github.com/pheinrich/aitrade/pkg/app/strategy"
	"github.com/pheinrich/aitrade/pkg/config"
	"github.com/pheinrich/aitrade/pkg/db"
	"github.com/pheinrich/aitrade/pkg/model"
)

type Trader struct {
	client        *client.IBClient
	tradeRepo     *db.TradeRepository
	newsRepo      *db.NewsRepository
	balanceRepo   *db.BalanceRepository
	whitelistRepo *db.WhitelistRepository
	positionRepo  *db.PositionRepository
	strategy      strategy.Strategy
	executor      *Executor
	dryRunExec    *DryRunExecutor
	limiter       *RateLimiter
	stopLossMgr   *StopLossManager
	cfg           *config.TradingConfig
	logger        *slog.Logger
}

func NewTrader(
	client *client.IBClient,
	tradeRepo *db.TradeRepository,
	newsRepo *db.NewsRepository,
	balanceRepo *db.BalanceRepository,
	whitelistRepo *db.WhitelistRepository,
	positionRepo *db.PositionRepository,
	strategy strategy.Strategy,
	cfg *config.TradingConfig,
	logger *slog.Logger,
) *Trader {
	riskParams := strategy.GetRiskParams()
	limiter := NewRateLimiter(riskParams.MaxTradesPerHour, riskParams.MaxParallelTrades)
	stopLossMgr := NewStopLossManager(client, tradeRepo, cfg.StopLossEnabled, logger)
	executor := NewExecutor(client, tradeRepo, stopLossMgr, logger)
	dryRunExec := NewDryRunExecutor(tradeRepo, balanceRepo, positionRepo, cfg.DryRunBalance, logger)

	return &Trader{
		client:        client,
		tradeRepo:     tradeRepo,
		newsRepo:      newsRepo,
		balanceRepo:   balanceRepo,
		whitelistRepo: whitelistRepo,
		positionRepo:  positionRepo,
		strategy:      strategy,
		executor:      executor,
		dryRunExec:    dryRunExec,
		limiter:       limiter,
		stopLossMgr:   stopLossMgr,
		cfg:           cfg,
		logger:        logger,
	}
}

func (t *Trader) Run(ctx context.Context) error {
	mode := "LIVE"
	if t.cfg.DryRun {
		mode = "DRY-RUN"
	}

	t.logger.Info("trader starting",
		slog.String("mode", mode),
		slog.String("strategy", t.strategy.Name()),
		slog.Duration("pending_time", t.cfg.PendingTime.Duration),
		slog.Bool("trading_enabled", t.cfg.TradingEnabled),
		slog.Duration("trading_interval", t.cfg.TradingInterval.Duration),
		slog.Int("watch_symbols", len(t.cfg.WatchSymbols)),
	)

	// Wait for IB Gateway connection in live mode (max 10 seconds)
	if !t.cfg.DryRun {
		t.logger.Info("waiting for IB Gateway connection")
		timeout := time.After(10 * time.Second)
		ticker := time.NewTicker(100 * time.Millisecond)
		defer ticker.Stop()

	waitLoop:
		for {
			if t.client.IsConnected() {
				t.logger.Info("IB Gateway connected and ready")
				break waitLoop
			}

			select {
			case <-timeout:
				t.logger.Warn("IB Gateway connection timeout - starting anyway")
				break waitLoop
			case <-ctx.Done():
				return ctx.Err()
			case <-ticker.C:
				// Continue waiting
			}
		}
	}

	// Start stop-loss manager only in live mode
	if !t.cfg.DryRun {
		go t.stopLossMgr.Run(ctx)
	}

	// Start auto-trading loop if enabled
	if t.cfg.TradingEnabled {
		go t.runTradingLoop(ctx)
	}

	// Process pending trades ticker
	ticker := time.NewTicker(5 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			t.logger.Info("trader stopping")
			return ctx.Err()
		case <-ticker.C:
			if err := t.processPendingTrades(ctx); err != nil {
				t.logger.Error("failed to process pending trades", slog.Any("error", err))
			}
		}
	}
}

func (t *Trader) processPendingTrades(ctx context.Context) error {
	// Get expired pending trades
	trades, err := t.tradeRepo.GetExpiredPendingTrades(ctx)
	if err != nil {
		return fmt.Errorf("failed to get expired pending trades: %w", err)
	}

	for _, trade := range trades {
		if err := t.executePendingTrade(ctx, trade); err != nil {
			t.logger.Error("failed to execute pending trade",
				slog.Int64("trade_id", trade.ID),
				slog.Any("error", err),
			)
		}
	}

	return nil
}

func (t *Trader) runTradingLoop(ctx context.Context) {
	t.logger.Info("auto-trading loop starting")

	ticker := time.NewTicker(t.cfg.TradingInterval.Duration)
	defer ticker.Stop()

	// Run immediately on start
	if err := t.analyzeAndTrade(ctx); err != nil {
		t.logger.Error("trading analysis failed", slog.Any("error", err))
	}

	for {
		select {
		case <-ctx.Done():
			t.logger.Info("auto-trading loop stopping")
			return
		case <-ticker.C:
			if err := t.analyzeAndTrade(ctx); err != nil {
				t.logger.Error("trading analysis failed", slog.Any("error", err))
			}
		}
	}
}

func (t *Trader) analyzeAndTrade(ctx context.Context) error {
	t.logger.Info("analyzing market for trading opportunities")

	// Get recent news (last 24 hours)
	news, err := t.newsRepo.GetRecent(ctx, 100)
	if err != nil {
		return fmt.Errorf("failed to get news: %w", err)
	}

	t.logger.Info("fetched news articles", slog.Int("count", len(news)))

	// Check existing positions for SELL opportunities
	if err := t.checkPositionsForSell(ctx, news); err != nil {
		t.logger.Error("failed to check positions for sell", slog.Any("error", err))
	}

	// Analyze each watched symbol for BUY opportunities
	for _, symbol := range t.cfg.WatchSymbols {
		// Skip if we already have a position
		position, err := t.positionRepo.GetBySymbol(ctx, symbol)
		if err != nil {
			t.logger.Error("failed to check position", slog.String("symbol", symbol), slog.Any("error", err))
			continue
		}
		if position != nil {
			t.logger.Debug("skipping symbol - position already open", slog.String("symbol", symbol))
			continue
		}

		// Filter news relevant to this symbol
		symbolNews := filterNewsBySymbol(news, symbol)

		t.logger.Info("analyzing symbol",
			slog.String("symbol", symbol),
			slog.Int("relevant_news", len(symbolNews)),
		)

		// Get market data
		marketData, err := t.getMarketData(ctx, symbol)
		if err != nil {
			t.logger.Warn("failed to get market data",
				slog.String("symbol", symbol),
				slog.Any("error", err),
			)
			continue
		}

		// Run strategy analysis
		signal, err := t.strategy.Analyze(ctx, marketData, symbolNews)
		if err != nil {
			t.logger.Warn("strategy analysis failed",
				slog.String("symbol", symbol),
				slog.Any("error", err),
			)
			continue
		}

		// No signal means no trade opportunity
		if signal == nil {
			t.logger.Debug("no trade signal generated", slog.String("symbol", symbol))
			continue
		}

		// Only process BUY signals here (SELL is handled in checkPositionsForSell)
		if signal.Action != model.ActionBuy {
			continue
		}

		// Calculate position size based on confidence and available capital
		availableCapital := t.getAvailableCapital(ctx)
		signal.Quantity = t.strategy.CalculatePositionSize(marketData.LastPrice, signal.Confidence, availableCapital)

		// Apply absolute maximum trade value limit
		validatedQuantity, err := strategy.ValidateTradeValue(signal.Quantity, marketData.LastPrice, t.cfg.MaxTradeValue)
		if err != nil {
			t.logger.Warn("trade rejected - exceeds max trade value",
				slog.String("symbol", signal.Symbol),
				slog.Int("calculated_quantity", signal.Quantity),
				slog.Float64("price", marketData.LastPrice),
				slog.Float64("max_trade_value", t.cfg.MaxTradeValue),
				slog.Any("error", err),
			)
			continue
		}

		signal.Quantity = validatedQuantity
		tradeValue := marketData.LastPrice * float64(signal.Quantity)

		// Create pending trade
		t.logger.Info("trade signal generated",
			slog.String("symbol", signal.Symbol),
			slog.String("action", string(signal.Action)),
			slog.Int("quantity", signal.Quantity),
			slog.Float64("confidence", signal.Confidence),
			slog.Float64("trade_value", tradeValue),
		)

		if err := t.CreatePendingTrade(ctx, signal); err != nil {
			t.logger.Error("failed to create pending trade",
				slog.String("symbol", signal.Symbol),
				slog.Any("error", err),
			)
		}
	}

	return nil
}

func (t *Trader) checkPositionsForSell(ctx context.Context, news []*model.NewsArticle) error {
	positions, err := t.positionRepo.GetAll(ctx)
	if err != nil {
		return fmt.Errorf("failed to get positions: %w", err)
	}

	for _, position := range positions {
		// Get current market data
		marketData, err := t.getMarketData(ctx, position.Symbol)
		if err != nil {
			t.logger.Warn("failed to get market data for position",
				slog.String("symbol", position.Symbol),
				slog.Any("error", err),
			)
			continue
		}

		// Update position with current price
		position.CurrentPrice = &marketData.LastPrice
		pnl := (marketData.LastPrice - position.EntryPrice) * float64(position.Quantity)
		position.UnrealizedPnL = &pnl

		if err := t.positionRepo.Update(ctx, position); err != nil {
			t.logger.Error("failed to update position", slog.Any("error", err))
		}

		// Check hold time
		holdTime := time.Since(position.OpenedAt)
		if holdTime < time.Duration(t.cfg.HoldTimeMinutes)*time.Minute {
			t.logger.Debug("position not old enough to sell",
				slog.String("symbol", position.Symbol),
				slog.Duration("hold_time", holdTime),
			)
			continue
		}

		// Calculate profit percentage
		profitPct := ((marketData.LastPrice - position.EntryPrice) / position.EntryPrice) * 100

		// SELL trigger 1: Take profit target reached
		if profitPct >= t.cfg.TakeProfitPercent {
			t.logger.Info("take profit triggered",
				slog.String("symbol", position.Symbol),
				slog.Float64("profit_pct", profitPct),
				slog.Float64("target_pct", t.cfg.TakeProfitPercent),
			)

			signal := &strategy.TradeSignal{
				Symbol:     position.Symbol,
				Action:     model.ActionSell,
				Quantity:   position.Quantity,
				Confidence: 0.95,
				Reasoning:  fmt.Sprintf("Take profit: %.2f%% gain (target: %.2f%%)", profitPct, t.cfg.TakeProfitPercent),
			}

			if err := t.CreatePendingTrade(ctx, signal); err != nil {
				t.logger.Error("failed to create sell trade", slog.Any("error", err))
			}
			continue
		}

		// SELL trigger 2: Negative sentiment (if enabled)
		if t.cfg.NegativeSentiment {
			symbolNews := filterNewsBySymbol(news, position.Symbol)
			if len(symbolNews) >= 2 {
				signal, err := t.strategy.Analyze(ctx, marketData, symbolNews)
				if err == nil && signal != nil && signal.Action == model.ActionSell {
					t.logger.Info("negative sentiment sell triggered",
						slog.String("symbol", position.Symbol),
						slog.String("reasoning", signal.Reasoning),
					)

					signal.Quantity = position.Quantity
					if err := t.CreatePendingTrade(ctx, signal); err != nil {
						t.logger.Error("failed to create sell trade", slog.Any("error", err))
					}
					continue
				}
			}
		}
	}

	return nil
}

func (t *Trader) getAvailableCapital(ctx context.Context) float64 {
	var totalCapital float64

	if t.cfg.DryRun {
		totalCapital = t.dryRunExec.GetBalance()
	} else {
		// In live mode, get from balance repository
		balance, err := t.balanceRepo.GetLatest(ctx)
		if err != nil || balance == nil {
			t.logger.Warn("failed to get balance, using default", slog.Any("error", err))
			totalCapital = 100000.0 // Fallback
		} else {
			// Use buying power (includes margin) for available capital
			totalCapital = balance.BuyingPower
		}
	}

	// Get current open positions count
	positions, err := t.positionRepo.GetAll(ctx)
	if err != nil {
		t.logger.Error("failed to get positions", slog.Any("error", err))
		positions = []*model.Position{} // Empty on error
	}
	openPositions := len(positions)

	// Get pending trades count
	pendingTrades, err := t.tradeRepo.GetPendingTrades(ctx)
	if err != nil {
		t.logger.Error("failed to get pending trades", slog.Any("error", err))
		pendingTrades = []*model.Trade{} // Empty on error
	}
	pendingCount := len(pendingTrades)

	// Calculate slots: max parallel trades - (open positions + pending trades)
	maxParallel := t.limiter.maxParallel
	usedSlots := openPositions + pendingCount
	availableSlots := maxParallel - usedSlots

	if availableSlots <= 0 {
		t.logger.Warn("no available trade slots",
			slog.Int("max_parallel", maxParallel),
			slog.Int("open_positions", openPositions),
			slog.Int("pending_trades", pendingCount),
		)
		return 0.0 // No capital available if no slots
	}

	// Divide total capital by max parallel trades to reserve capital for other trades
	// This ensures we don't use all capital on the first trade
	capitalPerSlot := totalCapital / float64(maxParallel)

	t.logger.Debug("calculated available capital",
		slog.Float64("total_capital", totalCapital),
		slog.Int("max_parallel", maxParallel),
		slog.Int("available_slots", availableSlots),
		slog.Float64("capital_per_slot", capitalPerSlot),
	)

	return capitalPerSlot
}

func (t *Trader) getMarketData(ctx context.Context, symbol string) (*strategy.MarketData, error) {
	// In dry-run mode, generate simulated market data
	if t.cfg.DryRun {
		return t.generateSimulatedMarketData(symbol), nil
	}

	// In live mode, get real market data from IB Gateway
	ibData, err := t.client.GetMarketData(ctx, symbol)
	if err != nil {
		t.logger.Warn("failed to get real market data from IB, using simulated data",
			slog.String("symbol", symbol),
			slog.Any("error", err))
		return t.generateSimulatedMarketData(symbol), nil
	}

	// Convert IB market data to strategy market data format
	return &strategy.MarketData{
		Symbol:    ibData.Symbol,
		LastPrice: ibData.LastPrice,
		BidPrice:  ibData.BidPrice,
		AskPrice:  ibData.AskPrice,
		Volume:    ibData.Volume,
	}, nil
}

func (t *Trader) generateSimulatedMarketData(symbol string) *strategy.MarketData {
	// Generate realistic simulated prices based on symbol
	basePrice := 150.0
	switch symbol {
	case "AAPL":
		basePrice = 180.0
	case "MSFT":
		basePrice = 420.0
	case "GOOGL":
		basePrice = 175.0
	case "TSLA":
		basePrice = 250.0
	case "AMZN":
		basePrice = 190.0
	}

	// Add some randomness
	change := (rand.Float64() - 0.5) * 10.0
	lastPrice := basePrice + change
	changePercent := (change / basePrice) * 100.0

	return &strategy.MarketData{
		Symbol:        symbol,
		LastPrice:     lastPrice,
		BidPrice:      lastPrice - 0.5,
		AskPrice:      lastPrice + 0.5,
		Volume:        int64(rand.Intn(10000000) + 1000000),
		Change:        change,
		ChangePercent: changePercent,
	}
}

func filterNewsBySymbol(news []*model.NewsArticle, symbol string) []*model.NewsArticle {
	var filtered []*model.NewsArticle
	for _, article := range news {
		if article.Symbols != "" && strings.Contains(article.Symbols, symbol) {
			filtered = append(filtered, article)
		}
	}
	return filtered
}

func (t *Trader) executePendingTrade(ctx context.Context, trade *model.Trade) error {
	// Check whitelist before execution
	whitelisted, err := t.whitelistRepo.IsSymbolWhitelisted(ctx, trade.Symbol)
	if err != nil {
		t.logger.Error("failed to check whitelist",
			slog.Int64("trade_id", trade.ID),
			slog.String("symbol", trade.Symbol),
			slog.Any("error", err),
		)
		return fmt.Errorf("failed to check whitelist: %w", err)
	}

	if !whitelisted {
		t.logger.Warn("trade rejected - symbol not whitelisted",
			slog.Int64("trade_id", trade.ID),
			slog.String("symbol", trade.Symbol),
		)

		// Mark trade as rejected
		now := time.Now()
		trade.Status = model.TradeRejected
		trade.RejectedAt = &now
		reason := fmt.Sprintf("Symbol %s not in whitelist", trade.Symbol)
		trade.RejectionReason = &reason

		if err := t.tradeRepo.Update(ctx, trade); err != nil {
			return fmt.Errorf("failed to update rejected trade: %w", err)
		}

		return nil
	}

	// Check rate limits
	if !t.limiter.CanTrade() {
		hourly, active := t.limiter.GetStats()
		t.logger.Warn("rate limit exceeded, postponing trade",
			slog.Int64("trade_id", trade.ID),
			slog.Int("hourly_trades", hourly),
			slog.Int("active_trades", active),
		)

		// Extend pending time by 1 minute
		newPendingUntil := time.Now().Add(1 * time.Minute)
		trade.PendingUtil = &newPendingUntil
		return t.tradeRepo.Update(ctx, trade)
	}

	// Execute the trade (dry run or live)
	t.limiter.RecordTrade()

	if t.cfg.DryRun {
		err = t.dryRunExec.ExecuteTrade(ctx, trade)
	} else {
		err = t.executor.ExecuteTrade(ctx, trade)
	}

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

	return nil
}

func (t *Trader) CreatePendingTrade(ctx context.Context, signal *strategy.TradeSignal) error {
	mode := "LIVE"
	if t.cfg.DryRun {
		mode = "DRY-RUN"
	}

	t.logger.Info("creating pending trade",
		slog.String("mode", mode),
		slog.String("symbol", signal.Symbol),
		slog.String("action", string(signal.Action)),
		slog.Int("quantity", signal.Quantity),
		slog.Float64("confidence", signal.Confidence),
	)

	now := time.Now()
	pendingUntil := now.Add(t.cfg.PendingTime.Duration)

	trade := &model.Trade{
		Symbol:      signal.Symbol,
		Action:      signal.Action,
		Quantity:    signal.Quantity,
		Status:      model.TradePending,
		Confidence:  signal.Confidence,
		Reasoning:   signal.Reasoning,
		CreatedAt:   now,
		PendingUtil: &pendingUntil,
		IsDryRun:    t.cfg.DryRun,
	}

	// Set stop-loss price if enabled
	if t.strategy.GetRiskParams().StopLossPercent > 0 {
		// This will be calculated when we have the executed price
		stopLossPercent := t.strategy.GetRiskParams().StopLossPercent
		_ = stopLossPercent // Will be used after execution
	}

	if err := t.tradeRepo.Create(ctx, trade); err != nil {
		return fmt.Errorf("failed to create pending trade: %w", err)
	}

	t.logger.Info("pending trade created",
		slog.Int64("trade_id", trade.ID),
		slog.Time("pending_until", pendingUntil),
		slog.Bool("dry_run", t.cfg.DryRun),
	)

	return nil
}

func (t *Trader) ApproveTrade(ctx context.Context, tradeID int64, forceNow bool) error {
	trade, err := t.tradeRepo.GetByID(ctx, tradeID)
	if err != nil {
		return fmt.Errorf("failed to get trade: %w", err)
	}

	if trade.Status != model.TradePending {
		return fmt.Errorf("trade %d is not in pending status", tradeID)
	}

	now := time.Now()
	trade.Status = model.TradeApproved
	trade.ApprovedAt = &now
	trade.ForcedByUser = forceNow

	if forceNow {
		// Execute immediately
		trade.PendingUtil = &now
	}

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

	t.logger.Info("trade approved",
		slog.Int64("trade_id", tradeID),
		slog.Bool("force_now", forceNow),
	)

	return nil
}

func (t *Trader) RejectTrade(ctx context.Context, tradeID int64, reason string) error {
	trade, err := t.tradeRepo.GetByID(ctx, tradeID)
	if err != nil {
		return fmt.Errorf("failed to get trade: %w", err)
	}

	if trade.Status != model.TradePending {
		return fmt.Errorf("trade %d is not in pending status", tradeID)
	}

	now := time.Now()
	trade.Status = model.TradeRejected
	trade.RejectedAt = &now
	trade.RejectionReason = &reason

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

	t.logger.Info("trade rejected",
		slog.Int64("trade_id", tradeID),
		slog.String("reason", reason),
	)

	return nil
}

// GetLiveBalance fetches the current account balance from IB Gateway
func (t *Trader) GetLiveBalance(ctx context.Context) (*model.Balance, error) {
	summary, err := t.client.GetAccountSummary(ctx)
	if err != nil {
		return nil, fmt.Errorf("failed to get account summary from IB: %w", err)
	}

	balance := &model.Balance{
		Timestamp:     time.Now(),
		TotalValue:    summary.TotalValue,
		CashBalance:   summary.CashBalance,
		BuyingPower:   summary.BuyingPower,
		UnrealizedPnL: &summary.UnrealizedPnL,
		RealizedPnL:   &summary.RealizedPnL,
	}

	return balance, nil
}