KalBot/lib/paper/engine.js

import { db } from '../db.js';
import { notify } from '../notify.js';

/**
 * Per-Strategy Paper Trading Engine.
 * Each strategy gets its own isolated balance, PnL, and trade history.
 */
class StrategyPaperAccount {
  constructor(strategyName, initialBalance = 1000) {
    this.strategyName = strategyName;
    this.balance = initialBalance;
    this.initialBalance = initialBalance;
    this.openPositions = new Map(); // ticker -> [positions]
    this.totalPnL = 0;
    this.wins = 0;
    this.losses = 0;
  }

  getStats() {
    const openPositionsList = [];
    for (const [, positions] of this.openPositions) {
      openPositionsList.push(...positions);
    }
    return {
      strategy: this.strategyName,
      balance: parseFloat(this.balance.toFixed(2)),
      initialBalance: this.initialBalance,
      totalPnL: parseFloat(this.totalPnL.toFixed(2)),
      wins: this.wins,
      losses: this.losses,
      winRate: this.wins + this.losses > 0
        ? parseFloat(((this.wins / (this.wins + this.losses)) * 100).toFixed(1))
        : 0,
      openPositions: openPositionsList,
      totalTrades: this.wins + this.losses
    };
  }
}

export class PaperEngine {
  constructor(initialBalancePerStrategy = 1000) {
    this.initialBalancePerStrategy = initialBalancePerStrategy;
    this.accounts = new Map(); // strategyName -> StrategyPaperAccount
    this._resetting = false;
  }

  _getAccount(strategyName) {
    if (!this.accounts.has(strategyName)) {
      this.accounts.set(strategyName, new StrategyPaperAccount(strategyName, this.initialBalancePerStrategy));
    }
    return this.accounts.get(strategyName);
  }

  /**
   * Generate a short unique ID safe for SurrealDB v2 record keys.
   * Avoids colons so we control the full `table:id` format ourselves.
   */
  _genId() {
    return `pt_${Date.now()}_${Math.random().toString(36).slice(2, 8)}`;
  }

  /**
   * Extract the raw record key from a SurrealDB id.
   * e.g. "paper_positions:pt_123_abc" -> "pt_123_abc"
   *      "pt_123_abc" -> "pt_123_abc"
   */
  _rawId(id) {
    if (!id) return id;
    const str = typeof id === 'object' && id.id ? String(id.id) : String(id);
    const idx = str.indexOf(':');
    return idx >= 0 ? str.slice(idx + 1) : str;
  }

  /**
   * Build the full `table:id` string for use in raw queries.
   */
  _recordId(id) {
    const raw = this._rawId(id);
    return raw.startsWith('paper_positions:') ? raw : `paper_positions:⟨${raw}⟩`;
  }

  async init() {
    try {
      const states = await db.query('SELECT * FROM paper_strategy_state ORDER BY timestamp DESC');
      const rows = states[0] || [];
      const seen = new Set();
      for (const saved of rows) {
        if (!saved.strategyName || seen.has(saved.strategyName)) continue;
        seen.add(saved.strategyName);
        const acct = this._getAccount(saved.strategyName);
        acct.balance = saved.balance;
        acct.totalPnL = saved.totalPnL;
        acct.wins = saved.wins;
        acct.losses = saved.losses;
        console.log(`[Paper:${saved.strategyName}] Restored: $${acct.balance.toFixed(2)}, ${acct.wins}W/${acct.losses}L`);
      }

      const positions = await db.query('SELECT * FROM paper_positions WHERE settled = false');
      if (positions[0]) {
        for (const pos of positions[0]) {
          const acct = this._getAccount(pos.strategy);
          const list = acct.openPositions.get(pos.ticker) || [];
          list.push(pos);
          acct.openPositions.set(pos.ticker, list);
        }
        const totalOpen = positions[0].length;
        if (totalOpen > 0) {
          console.log(`[Paper] Loaded ${totalOpen} open position(s) from DB`);
        }
      }
    } catch (e) {
      console.error('[Paper] Init error (fresh start):', e.message);
    }
  }

  async executeTrade(signal, marketState) {
    if (this._resetting) return null;

    const acct = this._getAccount(signal.strategy);
    const cost = signal.size;

    if (acct.balance < cost) {
      console.log(`[Paper:${signal.strategy}] Insufficient balance ($${acct.balance.toFixed(2)}) for $${cost} trade`);
      return null;
    }

    const trade = {
      id: this._genId(),
      strategy: signal.strategy,
      ticker: signal.ticker,
      side: signal.side.toLowerCase(),
      price: signal.price,
      size: signal.size,
      cost,
      reason: signal.reason,
      entryTime: Date.now(),
      settled: false,
      result: null,
      pnl: null,
      marketState: {
        yesPct: marketState.yesPct,
        noPct: marketState.noPct,
        yesOdds: marketState.yesOdds,
        noOdds: marketState.noOdds
      }
    };

    acct.balance -= cost;

    const list = acct.openPositions.get(trade.ticker) || [];
    list.push(trade);
    acct.openPositions.set(trade.ticker, list);

    try {
      await db.create('paper_positions', trade);
      await this._saveState(acct);
    } catch (e) {
      console.error('[Paper] DB write error:', e.message);
    }

    const msg = `📝 PAPER [${trade.strategy}] ${trade.side.toUpperCase()} @ ${trade.price}¢ ($${cost}) | ${trade.reason}`;
    console.log(`[Paper] ${msg}`);
    await notify(msg, `Paper: ${trade.strategy}`, '1');

    return trade;
  }

  async settle(ticker, rawResult) {
    const result = String(rawResult || '').toLowerCase();

    if (result !== 'yes' && result !== 'no') {
      console.warn(`[Paper] Unknown settlement result "${rawResult}" for ${ticker}, skipping.`);
      return null;
    }

    const allSettled = [];

    for (const [strategyName, acct] of this.accounts) {
      const positions = acct.openPositions.get(ticker);
      if (!positions || positions.length === 0) continue;

      console.log(`[Paper:${strategyName}] Settling ${positions.length} position(s) for ${ticker}, result: ${result}`);

      for (const pos of positions) {
        const side = String(pos.side || '').toLowerCase();
        const won = side === result;

        const price = pos.price > 0 ? pos.price : 50;
        const payout = won ? (100 / price) * pos.cost : 0;
        const pnl = payout - pos.cost;

        pos.settled = true;
        pos.result = result;
        pos.pnl = parseFloat(pnl.toFixed(2));
        pos.settleTime = Date.now();

        acct.balance += payout;
        acct.totalPnL += pnl;

        if (won) acct.wins++;
        else acct.losses++;

        const recordId = this._recordId(pos.id);

        try {
          const updated = await db.query(
            `UPDATE ${recordId} SET settled = true, result = $result, pnl = $pnl, settleTime = $settleTime`,
            { result, pnl: pos.pnl, settleTime: pos.settleTime }
          );
          const rows = updated[0] || [];
          if (rows.length === 0) {
            // Record vanished from DB — re-insert the full settled trade
            await db.create('paper_positions', { ...pos, id: this._rawId(pos.id) });
          }
        } catch (e) {
          console.error('[Paper] Settle DB error:', e.message);
          try {
            await db.create('paper_positions', { ...pos, id: this._rawId(pos.id) });
          } catch (e2) {
            console.error('[Paper] Settle DB fallback error:', e2.message);
          }
        }

        const emoji = won ? '✅' : '❌';
        const msg = `${emoji} [${strategyName}] ${side.toUpperCase()} ${won ? 'WON' : 'LOST'} | PnL: $${pnl.toFixed(2)} | Bal: $${acct.balance.toFixed(2)}`;
        console.log(`[Paper] ${msg}`);
        await notify(msg, won ? `${strategyName} Win!` : `${strategyName} Loss`, '1');

        allSettled.push(pos);
      }

      acct.openPositions.delete(ticker);
      await this._saveState(acct);
    }

    return allSettled.length > 0 ? allSettled : null;
  }

  getOpenTickers() {
    const tickers = new Set();
    for (const [, acct] of this.accounts) {
      for (const ticker of acct.openPositions.keys()) {
        tickers.add(ticker);
      }
    }
    return Array.from(tickers);
  }

  async checkOrphans(getMarketFn) {
    const orphanTickers = this.getOpenTickers();
    if (!orphanTickers.length) return { settled: [], expired: [] };

    const results = { settled: [], expired: [] };

    for (const ticker of orphanTickers) {
      try {
        const market = await getMarketFn(ticker);
        const status = String(market?.status || '').toLowerCase();
        const result = market?.result;

        if (result) {
          console.log(`[Paper] Delayed result found for ${ticker}: "${result}"`);
          const settledPos = await this.settle(ticker, result);
          if (settledPos) results.settled.push(...settledPos);
        } else if (['expired', 'cancelled'].includes(status)) {
          console.log(`[Paper] Ticker ${ticker} marked as ${status} — force-settling as expired`);
          const expiredPos = await this._forceExpirePositions(ticker);
          if (expiredPos) results.expired.push(...expiredPos);
        }
      } catch (e) {
        console.error(`[Paper] Orphan check failed for ${ticker}:`, e.message);
      }
    }
    return results;
  }

  async _forceExpirePositions(ticker) {
    const expired = [];
    for (const [strategyName, acct] of this.accounts) {
      const positions = acct.openPositions.get(ticker);
      if (!positions || !positions.length) continue;

      for (const pos of positions) {
        pos.settled = true;
        pos.result = 'expired';
        pos.pnl = parseFloat((-pos.cost).toFixed(2));
        pos.settleTime = Date.now();

        acct.totalPnL -= pos.cost;
        acct.losses++;

        const recordId = this._recordId(pos.id);

        try {
          const updated = await db.query(
            `UPDATE ${recordId} SET settled = true, result = $result, pnl = $pnl, settleTime = $settleTime`,
            { result: 'expired', pnl: pos.pnl, settleTime: pos.settleTime }
          );
          const rows = updated[0] || [];
          if (rows.length === 0) {
            await db.create('paper_positions', { ...pos, id: this._rawId(pos.id) });
          }
        } catch (e) {
          console.error('[Paper] Force-expire DB error:', e.message);
        }

        console.log(`[Paper:${strategyName}] Force-expired position ${pos.id} for ${ticker} (lost $${pos.cost})`);
        expired.push(pos);
      }

      acct.openPositions.delete(ticker);
      await this._saveState(acct);
    }
    return expired;
  }

  async resetAll() {
    this._resetting = true;

    for (const [name, acct] of this.accounts) {
      for (const [ticker, positions] of acct.openPositions) {
        for (const pos of positions) {
          acct.balance += pos.cost;
          pos.settled = true;
          pos.result = 'cancelled';
          pos.pnl = 0;
          pos.settleTime = Date.now();

          try {
            const recordId = this._recordId(pos.id);
            await db.query(
              `UPDATE ${recordId} SET settled = true, result = $result, pnl = $pnl, settleTime = $settleTime`,
              { result: 'cancelled', pnl: 0, settleTime: pos.settleTime }
            );
          } catch (e) {}
          console.log(`[Paper:${name}] Cancelled open position ${pos.id} for ${ticker} (refunded $${pos.cost})`);
        }
      }
      acct.openPositions.clear();
    }

    for (const [name, acct] of this.accounts) {
      acct.balance = acct.initialBalance;
      acct.totalPnL = 0;
      acct.wins = 0;
      acct.losses = 0;
      acct.openPositions.clear();
      console.log(`[Paper:${name}] Reset to $${acct.initialBalance}`);
    }

    try {
      await db.query('DELETE paper_positions');
      await db.query('DELETE paper_strategy_state');
      console.log('[Paper] Cleared all DB records');
    } catch (e) {
      console.error('[Paper] Reset DB error:', e.message);
    }

    for (const [, acct] of this.accounts) {
      await this._saveState(acct);
    }

    this._resetting = false;
  }

  getStats() {
    const allOpen = [];
    let totalBalance = 0;
    let totalPnL = 0;
    let totalWins = 0;
    let totalLosses = 0;

    for (const [, acct] of this.accounts) {
      const s = acct.getStats();
      totalBalance += s.balance;
      totalPnL += s.totalPnL;
      totalWins += s.wins;
      totalLosses += s.losses;
      allOpen.push(...s.openPositions);
    }

    return {
      balance: parseFloat(totalBalance.toFixed(2)),
      totalPnL: parseFloat(totalPnL.toFixed(2)),
      wins: totalWins,
      losses: totalLosses,
      winRate: totalWins + totalLosses > 0
        ? parseFloat(((totalWins / (totalWins + totalLosses)) * 100).toFixed(1))
        : 0,
      openPositions: allOpen,
      totalTrades: totalWins + totalLosses
    };
  }

  getPerStrategyStats() {
    const result = {};
    for (const [name, acct] of this.accounts) {
      result[name] = acct.getStats();
    }
    return result;
  }

  async _saveState(acct) {
    try {
      await db.create('paper_strategy_state', {
        strategyName: acct.strategyName,
        balance: acct.balance,
        totalPnL: acct.totalPnL,
        wins: acct.wins,
        losses: acct.losses,
        timestamp: Date.now()
      });
    } catch (e) {
      console.error('[Paper] State save error:', e.message);
    }
  }
}