autofirmer-expanded/lib/db.ts

import Database from 'better-sqlite3';
import path from 'path';

const db = new Database(path.join(process.cwd(), 'autotrader.sqlite'));

db.pragma('foreign_keys = ON');

db.exec(`
    CREATE TABLE IF NOT EXISTS firms (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        name TEXT NOT NULL UNIQUE,
        username TEXT NOT NULL,
        password TEXT NOT NULL
    );

    CREATE TABLE IF NOT EXISTS account_configs (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        firm_id INTEGER NOT NULL REFERENCES firms(id) ON DELETE CASCADE,
        prefix TEXT NOT NULL,
        profit_target REAL NOT NULL DEFAULT 3000,
        consistency REAL NOT NULL DEFAULT 0.5,
        min_day_pnl REAL NOT NULL DEFAULT -999,
        min_trading_days INTEGER NOT NULL DEFAULT 5
    );

    CREATE TABLE IF NOT EXISTS firm_fees (
        firm_id INTEGER NOT NULL REFERENCES firms(id) ON DELETE CASCADE,
        symbol  TEXT    NOT NULL,
        allin_fee     REAL NOT NULL DEFAULT 0,
        roundtrip_fee REAL NOT NULL DEFAULT 0,
        PRIMARY KEY (firm_id, symbol)
    );

    CREATE TABLE IF NOT EXISTS instruments (
        symbol  TEXT    PRIMARY KEY,
        enabled INTEGER NOT NULL DEFAULT 1
    );
`);

// Migration: add account_size column if it doesn't exist yet
try {
    db.exec('ALTER TABLE account_configs ADD COLUMN account_size REAL NOT NULL DEFAULT 50000');
} catch {
    // Column already exists
}

// Migration: add max_loss column (0 = no max loss limit)
try {
    db.exec('ALTER TABLE account_configs ADD COLUMN max_loss REAL NOT NULL DEFAULT 0');
} catch {
    // Column already exists
}

// Seed default firms if empty
const firmCount = (db.prepare('SELECT COUNT(*) as count FROM firms').get() as { count: number }).count;
if (firmCount === 0) {
    const insertFirm = db.prepare('INSERT INTO firms (name, username, password) VALUES (?, ?, ?)');
    const insertAccount = db.prepare(
        'INSERT INTO account_configs (firm_id, prefix, profit_target, consistency, min_day_pnl, min_trading_days, account_size) VALUES (?, ?, ?, ?, ?, ?, ?)'
    );

    const alpha = insertFirm.run('Alpha', 'brandonsenoli72786', '-Z2kPm7nBg');
    insertAccount.run(alpha.lastInsertRowid, 'AFSTDEV',  9000, 0.51, -999, 5,  150000);
    insertAccount.run(alpha.lastInsertRowid, 'AFSTDQA',  4500, 0.40, -999, 7,  150000);
    insertAccount.run(alpha.lastInsertRowid, 'AFZEROEV', 3000, 0.50, -999, 5,  100000);
    insertAccount.run(alpha.lastInsertRowid, 'AFZEROQA', 3000, 0.50, -999, 5,  100000);
    insertAccount.run(alpha.lastInsertRowid, 'AF',       3000, 0.50, -999, 5,  100000);

    const tpt = insertFirm.run('TakeProfitTrader', 'BRANDONLI1', 'W4592F5512U2817tv=');
    insertAccount.run(tpt.lastInsertRowid, 'TAKEPROFIT', 9000, 0.50, -999, 5, 150000);

    console.log('[db] Seeded default firms.');
}


// ── Interfaces ─────────────────────────────────────────────────────────────

export interface AccountConfigRow {
    id: number;
    firm_id: number;
    prefix: string;
    profit_target: number;
    consistency: number;
    min_day_pnl: number;
    min_trading_days: number;
    account_size: number;
    max_loss: number;
}

export interface FirmRow {
    id: number;
    name: string;
    username: string;
    password: string;
}

export interface FirmWithAccounts extends FirmRow {
    accounts: AccountConfigRow[];
}

export interface FirmFee {
    firmId: number;
    symbol: string;
    allinFee: number;
    roundtripFee: number;
}

// ── Firms ───────────────────────────────────────────────────────────────────

export function getFirms(): FirmWithAccounts[] {
    const firms = db.prepare('SELECT * FROM firms ORDER BY id').all() as FirmRow[];
    const getAccounts = db.prepare('SELECT * FROM account_configs WHERE firm_id = ? ORDER BY id');
    return firms.map((firm) => ({
        ...firm,
        accounts: getAccounts.all(firm.id) as AccountConfigRow[],
    }));
}

export function getFirmById(id: number): FirmWithAccounts | undefined {
    const firm = db.prepare('SELECT * FROM firms WHERE id = ?').get(id) as FirmRow | undefined;
    if (!firm) return undefined;
    const accounts = db.prepare('SELECT * FROM account_configs WHERE firm_id = ? ORDER BY id').all(id) as AccountConfigRow[];
    return { ...firm, accounts };
}

export function createFirm(name: string, username: string, password: string): FirmRow {
    const stmt = db.prepare('INSERT INTO firms (name, username, password) VALUES (?, ?, ?)');
    const result = stmt.run(name, username, password);
    return db.prepare('SELECT * FROM firms WHERE id = ?').get(result.lastInsertRowid) as FirmRow;
}

export function deleteFirm(id: number): boolean {
    const result = db.prepare('DELETE FROM firms WHERE id = ?').run(id);
    return result.changes > 0;
}

// ── Account Configs ─────────────────────────────────────────────────────────

export function createAccountConfig(firmId: number, data: {
    prefix: string;
    profitTarget: number;
    consistency: number;
    minDayPnL: number;
    minTradingDays: number;
    accountSize: number;
    maxLoss: number;
}): AccountConfigRow {
    const stmt = db.prepare(
        'INSERT INTO account_configs (firm_id, prefix, profit_target, consistency, min_day_pnl, min_trading_days, account_size, max_loss) VALUES (?, ?, ?, ?, ?, ?, ?, ?)'
    );
    const result = stmt.run(firmId, data.prefix, data.profitTarget, data.consistency, data.minDayPnL, data.minTradingDays, data.accountSize, data.maxLoss);
    return db.prepare('SELECT * FROM account_configs WHERE id = ?').get(result.lastInsertRowid) as AccountConfigRow;
}

export function deleteAccountConfig(id: number): boolean {
    const result = db.prepare('DELETE FROM account_configs WHERE id = ?').run(id);
    return result.changes > 0;
}

export function updateAccountConfig(id: number, data: {
    prefix: string;
    profitTarget: number;
    consistency: number;
    minDayPnL: number;
    minTradingDays: number;
    accountSize: number;
    maxLoss: number;
}): boolean {
    const result = db.prepare(`
        UPDATE account_configs
        SET prefix = ?, profit_target = ?, consistency = ?, min_day_pnl = ?, min_trading_days = ?, account_size = ?, max_loss = ?
        WHERE id = ?
    `).run(data.prefix, data.profitTarget, data.consistency, data.minDayPnL, data.minTradingDays, data.accountSize, data.maxLoss, id);
    return result.changes > 0;
}

// ── Firm Fees ────────────────────────────────────────────────────────────────

export function getFirmFees(firmId: number): FirmFee[] {
    return (db.prepare('SELECT firm_id, symbol, allin_fee, roundtrip_fee FROM firm_fees WHERE firm_id = ? ORDER BY symbol').all(firmId) as {
        firm_id: number;
        symbol: string;
        allin_fee: number;
        roundtrip_fee: number;
    }[]).map((r) => ({
        firmId: r.firm_id,
        symbol: r.symbol,
        allinFee: r.allin_fee,
        roundtripFee: r.roundtrip_fee,
    }));
}

export function upsertFirmFee(firmId: number, symbol: string, allinFee: number, roundtripFee: number): void {
    db.prepare(`
        INSERT INTO firm_fees (firm_id, symbol, allin_fee, roundtrip_fee)
        VALUES (?, ?, ?, ?)
        ON CONFLICT(firm_id, symbol) DO UPDATE SET
            allin_fee     = excluded.allin_fee,
            roundtrip_fee = excluded.roundtrip_fee
    `).run(firmId, symbol, allinFee, roundtripFee);
}

// ── Instruments ──────────────────────────────────────────────────────────────

const SYMBOLS = ['NQ','MNQ','ES','MES','YM','MYM','RTY','M2K','GC','MGC','SI','CL','MCL','NG','ZB','ZN','ZF','6E','6J','6B'];

// Seed instruments table if empty
const instrCount = (db.prepare('SELECT COUNT(*) as count FROM instruments').get() as { count: number }).count;
if (instrCount === 0) {
    const ins = db.prepare('INSERT INTO instruments (symbol, enabled) VALUES (?, 1)');
    for (const s of SYMBOLS) ins.run(s);
}

export interface InstrumentRow {
    symbol: string;
    enabled: boolean;
}

export function getInstruments(): InstrumentRow[] {
    return (db.prepare('SELECT symbol, enabled FROM instruments ORDER BY symbol').all() as { symbol: string; enabled: number }[])
        .map((r) => ({ symbol: r.symbol, enabled: r.enabled === 1 }));
}

export function setInstrumentEnabled(symbol: string, enabled: boolean): boolean {
    const result = db.prepare('UPDATE instruments SET enabled = ? WHERE symbol = ?').run(enabled ? 1 : 0, symbol);
    return result.changes > 0;
}