nofx/store/trader.go

package store

import (
	"database/sql"
	"encoding/json"
	"nofx/logger"
	"nofx/market"
	"slices"
	"strings"
	"time"
)

// TraderStore 交易员存储
type TraderStore struct {
	db          *sql.DB
	decryptFunc func(string) string
}

// Trader 交易员配置
type Trader struct {
	ID                   string    `json:"id"`
	UserID               string    `json:"user_id"`
	Name                 string    `json:"name"`
	AIModelID            string    `json:"ai_model_id"`
	ExchangeID           string    `json:"exchange_id"`
	InitialBalance       float64   `json:"initial_balance"`
	ScanIntervalMinutes  int       `json:"scan_interval_minutes"`
	IsRunning            bool      `json:"is_running"`
	BTCETHLeverage       int       `json:"btc_eth_leverage"`
	AltcoinLeverage      int       `json:"altcoin_leverage"`
	TradingSymbols       string    `json:"trading_symbols"`
	UseCoinPool          bool      `json:"use_coin_pool"`
	UseOITop             bool      `json:"use_oi_top"`
	CustomPrompt         string    `json:"custom_prompt"`
	OverrideBasePrompt   bool      `json:"override_base_prompt"`
	SystemPromptTemplate string    `json:"system_prompt_template"`
	IsCrossMargin        bool      `json:"is_cross_margin"`
	CreatedAt            time.Time `json:"created_at"`
	UpdatedAt            time.Time `json:"updated_at"`
}

// TraderFullConfig 交易员完整配置（包含AI模型和交易所）
type TraderFullConfig struct {
	Trader   *Trader
	AIModel  *AIModel
	Exchange *Exchange
}

func (s *TraderStore) initTables() error {
	_, err := s.db.Exec(`
		CREATE TABLE IF NOT EXISTS traders (
			id TEXT PRIMARY KEY,
			user_id TEXT NOT NULL DEFAULT 'default',
			name TEXT NOT NULL,
			ai_model_id TEXT NOT NULL,
			exchange_id TEXT NOT NULL,
			initial_balance REAL NOT NULL,
			scan_interval_minutes INTEGER DEFAULT 3,
			is_running BOOLEAN DEFAULT 0,
			btc_eth_leverage INTEGER DEFAULT 5,
			altcoin_leverage INTEGER DEFAULT 5,
			trading_symbols TEXT DEFAULT '',
			use_coin_pool BOOLEAN DEFAULT 0,
			use_oi_top BOOLEAN DEFAULT 0,
			custom_prompt TEXT DEFAULT '',
			override_base_prompt BOOLEAN DEFAULT 0,
			system_prompt_template TEXT DEFAULT 'default',
			is_cross_margin BOOLEAN DEFAULT 1,
			created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
			updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
			FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE
		)
	`)
	if err != nil {
		return err
	}

	// 触发器
	_, err = s.db.Exec(`
		CREATE TRIGGER IF NOT EXISTS update_traders_updated_at
		AFTER UPDATE ON traders
		BEGIN
			UPDATE traders SET updated_at = CURRENT_TIMESTAMP WHERE id = NEW.id;
		END
	`)
	if err != nil {
		return err
	}

	// 向后兼容
	alterQueries := []string{
		`ALTER TABLE traders ADD COLUMN custom_prompt TEXT DEFAULT ''`,
		`ALTER TABLE traders ADD COLUMN override_base_prompt BOOLEAN DEFAULT 0`,
		`ALTER TABLE traders ADD COLUMN is_cross_margin BOOLEAN DEFAULT 1`,
		`ALTER TABLE traders ADD COLUMN btc_eth_leverage INTEGER DEFAULT 5`,
		`ALTER TABLE traders ADD COLUMN altcoin_leverage INTEGER DEFAULT 5`,
		`ALTER TABLE traders ADD COLUMN trading_symbols TEXT DEFAULT ''`,
		`ALTER TABLE traders ADD COLUMN use_coin_pool BOOLEAN DEFAULT 0`,
		`ALTER TABLE traders ADD COLUMN use_oi_top BOOLEAN DEFAULT 0`,
		`ALTER TABLE traders ADD COLUMN system_prompt_template TEXT DEFAULT 'default'`,
	}
	for _, q := range alterQueries {
		s.db.Exec(q)
	}

	return nil
}

func (s *TraderStore) decrypt(encrypted string) string {
	if s.decryptFunc != nil {
		return s.decryptFunc(encrypted)
	}
	return encrypted
}

// Create 创建交易员
func (s *TraderStore) Create(trader *Trader) error {
	_, err := s.db.Exec(`
		INSERT INTO traders (id, user_id, name, ai_model_id, exchange_id, initial_balance, scan_interval_minutes,
		                     is_running, btc_eth_leverage, altcoin_leverage, trading_symbols, use_coin_pool,
		                     use_oi_top, custom_prompt, override_base_prompt, system_prompt_template, is_cross_margin)
		VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
	`, trader.ID, trader.UserID, trader.Name, trader.AIModelID, trader.ExchangeID, trader.InitialBalance,
		trader.ScanIntervalMinutes, trader.IsRunning, trader.BTCETHLeverage, trader.AltcoinLeverage,
		trader.TradingSymbols, trader.UseCoinPool, trader.UseOITop, trader.CustomPrompt,
		trader.OverrideBasePrompt, trader.SystemPromptTemplate, trader.IsCrossMargin)
	return err
}

// List 获取用户的交易员列表
func (s *TraderStore) List(userID string) ([]*Trader, error) {
	rows, err := s.db.Query(`
		SELECT id, user_id, name, ai_model_id, exchange_id, initial_balance, scan_interval_minutes, is_running,
		       COALESCE(btc_eth_leverage, 5), COALESCE(altcoin_leverage, 5), COALESCE(trading_symbols, ''),
		       COALESCE(use_coin_pool, 0), COALESCE(use_oi_top, 0), COALESCE(custom_prompt, ''),
		       COALESCE(override_base_prompt, 0), COALESCE(system_prompt_template, 'default'),
		       COALESCE(is_cross_margin, 1), created_at, updated_at
		FROM traders WHERE user_id = ? ORDER BY created_at DESC
	`, userID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var traders []*Trader
	for rows.Next() {
		var t Trader
		var createdAt, updatedAt string
		err := rows.Scan(
			&t.ID, &t.UserID, &t.Name, &t.AIModelID, &t.ExchangeID,
			&t.InitialBalance, &t.ScanIntervalMinutes, &t.IsRunning,
			&t.BTCETHLeverage, &t.AltcoinLeverage, &t.TradingSymbols,
			&t.UseCoinPool, &t.UseOITop, &t.CustomPrompt, &t.OverrideBasePrompt,
			&t.SystemPromptTemplate, &t.IsCrossMargin, &createdAt, &updatedAt,
		)
		if err != nil {
			return nil, err
		}
		t.CreatedAt, _ = time.Parse("2006-01-02 15:04:05", createdAt)
		t.UpdatedAt, _ = time.Parse("2006-01-02 15:04:05", updatedAt)
		traders = append(traders, &t)
	}
	return traders, nil
}

// UpdateStatus 更新交易员运行状态
func (s *TraderStore) UpdateStatus(userID, id string, isRunning bool) error {
	_, err := s.db.Exec(`UPDATE traders SET is_running = ? WHERE id = ? AND user_id = ?`, isRunning, id, userID)
	return err
}

// Update 更新交易员配置
func (s *TraderStore) Update(trader *Trader) error {
	_, err := s.db.Exec(`
		UPDATE traders SET
			name = ?, ai_model_id = ?, exchange_id = ?, scan_interval_minutes = ?,
			btc_eth_leverage = ?, altcoin_leverage = ?, trading_symbols = ?,
			custom_prompt = ?, override_base_prompt = ?, system_prompt_template = ?,
			is_cross_margin = ?, updated_at = CURRENT_TIMESTAMP
		WHERE id = ? AND user_id = ?
	`, trader.Name, trader.AIModelID, trader.ExchangeID, trader.ScanIntervalMinutes,
		trader.BTCETHLeverage, trader.AltcoinLeverage, trader.TradingSymbols,
		trader.CustomPrompt, trader.OverrideBasePrompt, trader.SystemPromptTemplate,
		trader.IsCrossMargin, trader.ID, trader.UserID)
	return err
}

// UpdateInitialBalance 更新初始余额
func (s *TraderStore) UpdateInitialBalance(userID, id string, newBalance float64) error {
	_, err := s.db.Exec(`UPDATE traders SET initial_balance = ? WHERE id = ? AND user_id = ?`, newBalance, id, userID)
	return err
}

// UpdateCustomPrompt 更新自定义提示词
func (s *TraderStore) UpdateCustomPrompt(userID, id string, customPrompt string, overrideBase bool) error {
	_, err := s.db.Exec(`UPDATE traders SET custom_prompt = ?, override_base_prompt = ? WHERE id = ? AND user_id = ?`,
		customPrompt, overrideBase, id, userID)
	return err
}

// Delete 删除交易员
func (s *TraderStore) Delete(userID, id string) error {
	_, err := s.db.Exec(`DELETE FROM traders WHERE id = ? AND user_id = ?`, id, userID)
	return err
}

// GetFullConfig 获取交易员完整配置
func (s *TraderStore) GetFullConfig(userID, traderID string) (*TraderFullConfig, error) {
	var trader Trader
	var aiModel AIModel
	var exchange Exchange
	var traderCreatedAt, traderUpdatedAt string
	var aiModelCreatedAt, aiModelUpdatedAt string
	var exchangeCreatedAt, exchangeUpdatedAt string

	err := s.db.QueryRow(`
		SELECT
			t.id, t.user_id, t.name, t.ai_model_id, t.exchange_id, t.initial_balance, t.scan_interval_minutes, t.is_running,
			COALESCE(t.btc_eth_leverage, 5), COALESCE(t.altcoin_leverage, 5), COALESCE(t.trading_symbols, ''),
			COALESCE(t.use_coin_pool, 0), COALESCE(t.use_oi_top, 0), COALESCE(t.custom_prompt, ''),
			COALESCE(t.override_base_prompt, 0), COALESCE(t.system_prompt_template, 'default'),
			COALESCE(t.is_cross_margin, 1), t.created_at, t.updated_at,
			a.id, a.user_id, a.name, a.provider, a.enabled, a.api_key,
			COALESCE(a.custom_api_url, ''), COALESCE(a.custom_model_name, ''), a.created_at, a.updated_at,
			e.id, e.user_id, e.name, e.type, e.enabled, e.api_key, e.secret_key, e.testnet,
			COALESCE(e.hyperliquid_wallet_addr, ''), COALESCE(e.aster_user, ''), COALESCE(e.aster_signer, ''),
			COALESCE(e.aster_private_key, ''), COALESCE(e.lighter_wallet_addr, ''), COALESCE(e.lighter_private_key, ''),
			COALESCE(e.lighter_api_key_private_key, ''), e.created_at, e.updated_at
		FROM traders t
		JOIN ai_models a ON t.ai_model_id = a.id AND t.user_id = a.user_id
		JOIN exchanges e ON t.exchange_id = e.id AND t.user_id = e.user_id
		WHERE t.id = ? AND t.user_id = ?
	`, traderID, userID).Scan(
		&trader.ID, &trader.UserID, &trader.Name, &trader.AIModelID, &trader.ExchangeID,
		&trader.InitialBalance, &trader.ScanIntervalMinutes, &trader.IsRunning,
		&trader.BTCETHLeverage, &trader.AltcoinLeverage, &trader.TradingSymbols,
		&trader.UseCoinPool, &trader.UseOITop, &trader.CustomPrompt, &trader.OverrideBasePrompt,
		&trader.SystemPromptTemplate, &trader.IsCrossMargin, &traderCreatedAt, &traderUpdatedAt,
		&aiModel.ID, &aiModel.UserID, &aiModel.Name, &aiModel.Provider, &aiModel.Enabled, &aiModel.APIKey,
		&aiModel.CustomAPIURL, &aiModel.CustomModelName, &aiModelCreatedAt, &aiModelUpdatedAt,
		&exchange.ID, &exchange.UserID, &exchange.Name, &exchange.Type, &exchange.Enabled,
		&exchange.APIKey, &exchange.SecretKey, &exchange.Testnet, &exchange.HyperliquidWalletAddr,
		&exchange.AsterUser, &exchange.AsterSigner, &exchange.AsterPrivateKey,
		&exchange.LighterWalletAddr, &exchange.LighterPrivateKey, &exchange.LighterAPIKeyPrivateKey,
		&exchangeCreatedAt, &exchangeUpdatedAt,
	)
	if err != nil {
		return nil, err
	}

	trader.CreatedAt, _ = time.Parse("2006-01-02 15:04:05", traderCreatedAt)
	trader.UpdatedAt, _ = time.Parse("2006-01-02 15:04:05", traderUpdatedAt)
	aiModel.CreatedAt, _ = time.Parse("2006-01-02 15:04:05", aiModelCreatedAt)
	aiModel.UpdatedAt, _ = time.Parse("2006-01-02 15:04:05", aiModelUpdatedAt)
	exchange.CreatedAt, _ = time.Parse("2006-01-02 15:04:05", exchangeCreatedAt)
	exchange.UpdatedAt, _ = time.Parse("2006-01-02 15:04:05", exchangeUpdatedAt)

	// 解密
	aiModel.APIKey = s.decrypt(aiModel.APIKey)
	exchange.APIKey = s.decrypt(exchange.APIKey)
	exchange.SecretKey = s.decrypt(exchange.SecretKey)
	exchange.AsterPrivateKey = s.decrypt(exchange.AsterPrivateKey)
	exchange.LighterPrivateKey = s.decrypt(exchange.LighterPrivateKey)
	exchange.LighterAPIKeyPrivateKey = s.decrypt(exchange.LighterAPIKeyPrivateKey)

	return &TraderFullConfig{
		Trader:   &trader,
		AIModel:  &aiModel,
		Exchange: &exchange,
	}, nil
}

// GetCustomCoins 获取所有交易员自定义币种
func (s *TraderStore) GetCustomCoins() []string {
	var symbol string
	var symbols []string
	_ = s.db.QueryRow(`
		SELECT GROUP_CONCAT(trading_symbols, ',') as symbol
		FROM traders WHERE trading_symbols != ''
	`).Scan(&symbol)

	// 如果没有自定义币种，返回默认币种
	if symbol == "" {
		var symbolJSON string
		_ = s.db.QueryRow(`SELECT value FROM system_config WHERE key = 'default_coins'`).Scan(&symbolJSON)
		if symbolJSON != "" {
			if err := json.Unmarshal([]byte(symbolJSON), &symbols); err != nil {
				logger.Warnf("⚠️  解析default_coins配置失败: %v，使用硬编码默认值", err)
				symbols = []string{"BTCUSDT", "ETHUSDT", "SOLUSDT", "BNBUSDT"}
			}
		} else {
			symbols = []string{"BTCUSDT", "ETHUSDT", "SOLUSDT", "BNBUSDT"}
		}
		return symbols
	}

	// 处理并去重币种列表
	for _, s := range strings.Split(symbol, ",") {
		if s == "" {
			continue
		}
		coin := market.Normalize(s)
		if !slices.Contains(symbols, coin) {
			symbols = append(symbols, coin)
		}
	}
	return symbols
}

// ListAll 获取所有用户的交易员列表
func (s *TraderStore) ListAll() ([]*Trader, error) {
	rows, err := s.db.Query(`
		SELECT id, user_id, name, ai_model_id, exchange_id, initial_balance, scan_interval_minutes, is_running,
		       COALESCE(btc_eth_leverage, 5), COALESCE(altcoin_leverage, 5), COALESCE(trading_symbols, ''),
		       COALESCE(use_coin_pool, 0), COALESCE(use_oi_top, 0), COALESCE(custom_prompt, ''),
		       COALESCE(override_base_prompt, 0), COALESCE(system_prompt_template, 'default'),
		       COALESCE(is_cross_margin, 1), created_at, updated_at
		FROM traders ORDER BY created_at DESC
	`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var traders []*Trader
	for rows.Next() {
		var t Trader
		var createdAt, updatedAt string
		err := rows.Scan(
			&t.ID, &t.UserID, &t.Name, &t.AIModelID, &t.ExchangeID,
			&t.InitialBalance, &t.ScanIntervalMinutes, &t.IsRunning,
			&t.BTCETHLeverage, &t.AltcoinLeverage, &t.TradingSymbols,
			&t.UseCoinPool, &t.UseOITop, &t.CustomPrompt, &t.OverrideBasePrompt,
			&t.SystemPromptTemplate, &t.IsCrossMargin, &createdAt, &updatedAt,
		)
		if err != nil {
			return nil, err
		}
		t.CreatedAt, _ = time.Parse("2006-01-02 15:04:05", createdAt)
		t.UpdatedAt, _ = time.Parse("2006-01-02 15:04:05", updatedAt)
		traders = append(traders, &t)
	}
	return traders, nil
}