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
  }

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

  async init() {
    try {
      // Load saved per-strategy states
      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`);
      }

      // Load open positions
      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) {
    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: `pt_${Date.now()}_${Math.random().toString(36).slice(2, 8)}`,
      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}`);

    return trade;
  }

  async settle(ticker, rawResult) {
    // Normalize result to lowercase — Kalshi may return "Yes"/"No"/"yes"/"no"
    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) {
        // Normalize stored side too
        const side = String(pos.side || '').toLowerCase();
        const won = side === result;

        // Payout: if won, each contract pays 100¢. Cost is in dollars, price in cents.
        // Contracts bought = cost / (price / 100). Payout = contracts * $1.
        // Simplified: payout = cost * (100 / price)
        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++;

        try {
          await db.query(
            `UPDATE paper_positions SET settled = true, result = $result, pnl = $pnl, settleTime = $settleTime WHERE id = $id`,
            { id: pos.id, result, pnl: pos.pnl, settleTime: pos.settleTime }
          );
        } catch (e) {
          console.error('[Paper] Settle DB error:', e.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`);

        allSettled.push(pos);
      }

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

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

  /**
   * Attempt to settle any orphaned open positions by checking market status.
   * Called on startup to clean up positions from previous runs.
   */
  async settleOrphans(getMarketFn) {
    const orphanTickers = new Set();
    for (const [, acct] of this.accounts) {
      for (const ticker of acct.openPositions.keys()) {
        orphanTickers.add(ticker);
      }
    }

    if (!orphanTickers.size) return;
    console.log(`[Paper] Checking ${orphanTickers.size} ticker(s) for orphaned positions...`);

    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] Orphan ticker ${ticker} has result "${result}" — settling now`);
          await this.settle(ticker, result);
        } else if (['closed', 'settled', 'expired', 'finalized'].includes(status)) {
          // Market closed but no result yet — mark as expired loss
          console.log(`[Paper] Orphan ticker ${ticker} status="${status}" with no result — force-settling as expired`);
          await this._forceExpirePositions(ticker);
        } else {
          console.log(`[Paper] Orphan ticker ${ticker} still active (status="${status}") — keeping open`);
        }
      } catch (e) {
        console.error(`[Paper] Orphan check failed for ${ticker}:`, e.message);
      }
    }
  }

  async _forceExpirePositions(ticker) {
    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++;

        try {
          await db.query(
            `UPDATE paper_positions SET settled = true, result = $result, pnl = $pnl, settleTime = $settleTime WHERE id = $id`,
            { id: pos.id, result: 'expired', pnl: pos.pnl, settleTime: pos.settleTime }
          );
        } 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})`);
      }

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

  /**
   * Reset all strategy stats and clear trade history.
   */
  async resetAll() {
    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);
    }

    // Save fresh state for each account
    for (const [, acct] of this.accounts) {
      await this._saveState(acct);
    }
  }

  /**
   * Get combined stats (backward compat) — aggregates all strategies.
   */
  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
    };
  }

  /**
   * Get per-strategy stats for the paper dashboard.
   */
  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);
    }
  }
}