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 instruments (
        symbol  TEXT    PRIMARY KEY,
        enabled INTEGER NOT NULL DEFAULT 1
    );
`);

// ── Machine-owned firm state ─────────────────────────────────────────────────
//
// Firms and account types themselves come from the master catalog (see
// lib/firm-catalog.ts) and are cached as one JSON blob in `settings`. These two
// tables hold the parts that are genuinely per-machine and must survive the
// catalog changing underneath them.
//
// Both reference firm ids as plain TEXT with NO foreign key on purpose. This
// file sets `foreign_keys = ON`, so an ON DELETE CASCADE here would wipe the
// Tradovate credentials the moment a firm was removed on the master — and
// mirroring master's deletions is routine, not exceptional. They must be
// orphan-tolerant: a firm that disappears and comes back finds its credentials
// and assignments still here.
db.exec(`
    CREATE TABLE IF NOT EXISTS firm_credentials (
        firm_id    TEXT PRIMARY KEY,
        firm_name  TEXT NOT NULL,
        username   TEXT NOT NULL,
        password   TEXT NOT NULL,
        updated_at TEXT NOT NULL
    );

    CREATE TABLE IF NOT EXISTS account_assignments (
        account_id      INTEGER PRIMARY KEY,
        firm_id         TEXT    NOT NULL,
        account_type_id TEXT,
        manual          INTEGER NOT NULL DEFAULT 0,
        resolved_at     TEXT    NOT NULL
    );
`);

// The legacy `firms`, `account_configs` and `firm_banned_symbols` tables are
// dead: firms now come from the master catalog. They are left on disk rather
// than dropped because `firms` is the only table that ever held plaintext
// Tradovate credentials, and a silent DROP would be unrecoverable if this
// change had to be rolled back. Safe to drop by hand once it has stuck.

// ── Firm credentials ─────────────────────────────────────────────────────────

export interface FirmCredentialRow {
    firm_id: string;
    firm_name: string;
    username: string;
    password: string;
    updated_at: string;
}

export function getCredentials(firmId: string): { username: string; password: string } | undefined {
    const row = db.prepare('SELECT username, password FROM firm_credentials WHERE firm_id = ?').get(firmId) as
        { username: string; password: string } | undefined;
    return row;
}

export function listCredentials(): FirmCredentialRow[] {
    return db.prepare('SELECT * FROM firm_credentials').all() as FirmCredentialRow[];
}

/** firm_name is stored alongside the id so credentials can be re-keyed by name
 *  if the master's database is ever rebuilt and the firms come back with new
 *  ids. See rekeyCredentialsByName in lib/firm-catalog.ts. */
export function setCredentials(firmId: string, firmName: string, username: string, password: string): void {
    db.prepare(`
        INSERT INTO firm_credentials (firm_id, firm_name, username, password, updated_at)
        VALUES (?, ?, ?, ?, ?)
        ON CONFLICT(firm_id) DO UPDATE SET
            firm_name = excluded.firm_name,
            username = excluded.username,
            password = excluded.password,
            updated_at = excluded.updated_at
    `).run(firmId, firmName, username, password, new Date().toISOString());
}

export function clearCredentials(firmId: string): boolean {
    return db.prepare('DELETE FROM firm_credentials WHERE firm_id = ?').run(firmId).changes > 0;
}

export function rekeyCredentials(fromFirmId: string, toFirmId: string): void {
    db.prepare('UPDATE firm_credentials SET firm_id = ? WHERE firm_id = ?').run(toFirmId, fromFirmId);
    db.prepare('UPDATE account_assignments SET firm_id = ? WHERE firm_id = ?').run(toFirmId, fromFirmId);
}

// ── Account type assignments ─────────────────────────────────────────────────

export interface AssignmentRow {
    account_id: number;
    firm_id: string;
    account_type_id: string | null;
    manual: number;   // 0 | 1
    resolved_at: string;
}

export function getAssignment(accountId: number): AssignmentRow | undefined {
    return db.prepare('SELECT * FROM account_assignments WHERE account_id = ?').get(accountId) as AssignmentRow | undefined;
}

export function listAssignments(firmId: string): AssignmentRow[] {
    return db.prepare('SELECT * FROM account_assignments WHERE firm_id = ?').all(firmId) as AssignmentRow[];
}

export function setAssignment(
    accountId: number,
    firmId: string,
    accountTypeId: string | null,
    manual: boolean,
): void {
    db.prepare(`
        INSERT INTO account_assignments (account_id, firm_id, account_type_id, manual, resolved_at)
        VALUES (?, ?, ?, ?, ?)
        ON CONFLICT(account_id) DO UPDATE SET
            firm_id = excluded.firm_id,
            account_type_id = excluded.account_type_id,
            manual = excluded.manual,
            resolved_at = excluded.resolved_at
    `).run(accountId, firmId, accountTypeId, manual ? 1 : 0, new Date().toISOString());
}

/** Writes an auto-resolved assignment, leaving any manual override alone.
 *  The `manual = 0` guard is what makes a user's pick stick across syncs. */
export function setAutoAssignment(accountId: number, firmId: string, accountTypeId: string | null): void {
    const existing = getAssignment(accountId);
    if (existing?.manual === 1) return;
    setAssignment(accountId, firmId, accountTypeId, false);
}

export function clearAssignment(accountId: number): boolean {
    return db.prepare('DELETE FROM account_assignments WHERE account_id = ?').run(accountId).changes > 0;
}

// ── Settings ─────────────────────────────────────────────────────────────────

db.exec(`
    CREATE TABLE IF NOT EXISTS settings (
        key   TEXT PRIMARY KEY,
        value TEXT NOT NULL
    );
`);

// Seed defaults if missing
const seedSetting = db.prepare(`INSERT OR IGNORE INTO settings (key, value) VALUES (?, ?)`);
seedSetting.run('max_concurrent_accounts', '5');
seedSetting.run('tick_interval_seconds', '60');
seedSetting.run('master_dashboard_url', 'https://master.juicerroom.com');
seedSetting.run('instance_name', '');
seedSetting.run('trading_hours', 'full_cme');
// Solves the captchas the AutoBuyer runs into. Blank uses the key hardcoded in
// app/api/autobuyer/captcha; set this (or TWOCAPTCHA_API_KEY) to bill another
// account.
seedSetting.run('twocaptcha_api_key', '');
// Blank falls back to the proxy hardcoded in app/api/autobuyer/captcha; `none`
// solves without one.
seedSetting.run('twocaptcha_proxy', '');
// Mirrors the in-memory scheduler state so a restart can resume it.
seedSetting.run('scheduler_running', '0');
seedSetting.run('scheduler_action', 'Buy');
seedSetting.run('scheduler_symbol', 'NQ');
seedSetting.run('scheduler_stop_after_all', '0');

export function getSetting(key: string): string | null {
    const row = db.prepare('SELECT value FROM settings WHERE key = ?').get(key) as { value: string } | undefined;
    return row?.value ?? null;
}

export function setSetting(key: string, value: string): void {
    db.prepare(`INSERT INTO settings (key, value) VALUES (?, ?) ON CONFLICT(key) DO UPDATE SET value = excluded.value`).run(key, value);
}

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

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

// Enabled on a fresh install. The rest still seed, so they can be switched on
// from the Instruments page, they just start off.
const DEFAULT_ENABLED = new Set(['NQ', 'GC', 'CL']);

// 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 (?, ?)');
    for (const s of SYMBOLS) ins.run(s, DEFAULT_ENABLED.has(s) ? 1 : 0);
}

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;
}

// ── Daily P&L Cache ──────────────────────────────────────────────────────────

db.exec(`
    CREATE TABLE IF NOT EXISTS daily_pnl (
        account_id   INTEGER NOT NULL,
        account_name TEXT    NOT NULL,
        date         TEXT    NOT NULL,
        pnl          REAL    NOT NULL,
        PRIMARY KEY (account_id, date)
    );
`);

export function saveDailyPnL(accountId: number, accountName: string, entries: { date: string; pnl: number }[]): void {
    const ins = db.prepare('INSERT OR REPLACE INTO daily_pnl (account_id, account_name, date, pnl) VALUES (?, ?, ?, ?)');
    const txn = db.transaction(() => {
        for (const e of entries) {
            ins.run(accountId, accountName, e.date, e.pnl);
        }
    });
    txn();
}

export function loadDailyPnL(accountId: number): { date: string; pnl: number }[] {
    return (db.prepare('SELECT date, pnl FROM daily_pnl WHERE account_id = ? ORDER BY date').all(accountId) as { date: string; pnl: number }[]);
}

// ── Fund Transactions Cache ───────────────────────────────────────────────────

db.exec(`
    CREATE TABLE IF NOT EXISTS fund_transactions (
        account_id   INTEGER NOT NULL,
        account_name TEXT    NOT NULL,
        date         TEXT    NOT NULL,
        amount       REAL    NOT NULL,
        PRIMARY KEY (account_id, date)
    );
`);

export function saveFundTransactions(accountId: number, accountName: string, entries: { date: string; amount: number }[]): void {
    const ins = db.prepare('INSERT OR REPLACE INTO fund_transactions (account_id, account_name, date, amount) VALUES (?, ?, ?, ?)');
    const txn = db.transaction(() => {
        for (const e of entries) ins.run(accountId, accountName, e.date, e.amount);
    });
    txn();
}

export function loadFundTransactions(accountId: number): { date: string; amount: number }[] {
    return db.prepare('SELECT date, amount FROM fund_transactions WHERE account_id = ? ORDER BY date').all(accountId) as { date: string; amount: number }[];
}

// ── Account Meta ─────────────────────────────────────────────────────────────

db.exec(`
    CREATE TABLE IF NOT EXISTS account_meta (
        account_id INTEGER NOT NULL,
        key        TEXT    NOT NULL,
        value      TEXT    NOT NULL,
        PRIMARY KEY (account_id, key)
    );
`);

export function saveAccountMeta(accountId: number, key: string, value: string): void {
    db.prepare('INSERT OR REPLACE INTO account_meta (account_id, key, value) VALUES (?, ?, ?)').run(accountId, key, value);
}

export function loadAccountMeta(accountId: number, key: string): string | null {
    const row = db.prepare('SELECT value FROM account_meta WHERE account_id = ? AND key = ?').get(accountId, key) as { value: string } | undefined;
    return row?.value ?? null;
}

// ── AutoBuyer captures ───────────────────────────────────────────────────────

db.exec(`
    CREATE TABLE IF NOT EXISTS autobuyer_captures (
        id          INTEGER PRIMARY KEY AUTOINCREMENT,
        url         TEXT    NOT NULL,
        title       TEXT    NOT NULL DEFAULT '',
        html        TEXT    NOT NULL,
        viewport    TEXT    NOT NULL DEFAULT '{}',
        elements    TEXT    NOT NULL DEFAULT '[]',
        captured_at INTEGER NOT NULL
    );
`);

export interface CaptureRow {
    id: number;
    url: string;
    title: string;
    html: string;
    viewport: string;   // JSON ViewportMetrics
    elements: string;   // JSON ElementPosition[]
    captured_at: number;
}

/** Number of captures kept on disk — the extension overwrites constantly, so only
 *  a short tail is useful and an unbounded table would grow by megabytes a minute. */
const CAPTURE_HISTORY = 5;

export function saveCapture(c: Omit<CaptureRow, 'id'>): void {
    db.prepare(
        'INSERT INTO autobuyer_captures (url, title, html, viewport, elements, captured_at) VALUES (?, ?, ?, ?, ?, ?)'
    ).run(c.url, c.title, c.html, c.viewport, c.elements, c.captured_at);

    db.prepare(`
        DELETE FROM autobuyer_captures
        WHERE id NOT IN (SELECT id FROM autobuyer_captures ORDER BY id DESC LIMIT ?)
    `).run(CAPTURE_HISTORY);
}

export function getLatestCapture(): CaptureRow | undefined {
    return db.prepare('SELECT * FROM autobuyer_captures ORDER BY id DESC LIMIT 1').get() as CaptureRow | undefined;
}

export function clearCaptures(): void {
    db.prepare('DELETE FROM autobuyer_captures').run();
}

// ── AutoBuyer locate/click requests ──────────────────────────────────────────
//
// The Python clicker can't see inside Chrome, and the extension can't move the
// mouse. So they meet here: Python enqueues "where is <selector>?", the extension
// focuses the tab, measures the element and writes back desktop coordinates.

db.exec(`
    CREATE TABLE IF NOT EXISTS autobuyer_locate (
        id          INTEGER PRIMARY KEY AUTOINCREMENT,
        selector    TEXT    NOT NULL,
        match_index INTEGER NOT NULL DEFAULT 0,
        url_pattern TEXT    NOT NULL DEFAULT '',
        status      TEXT    NOT NULL DEFAULT 'pending',
        result      TEXT,
        error       TEXT,
        created_at  INTEGER NOT NULL,
        resolved_at INTEGER
    );
`);

export interface LocateRow {
    id: number;
    selector: string;
    match_index: number;
    url_pattern: string;
    open_url: string;
    navigate_url: string;
    options: string;      // JSON, per-request extras (scroll parameters, ...)
    status: 'pending' | 'claimed' | 'done' | 'error';
    result: string | null;
    error: string | null;
    created_at: number;
    resolved_at: number | null;
}

const LOCATE_HISTORY = 20;

export function createLocateRequest(selector: string, matchIndex: number, urlPattern: string, openUrl = '', navigateUrl = '', options = '{}'): LocateRow {
    const res = db.prepare(
        'INSERT INTO autobuyer_locate (selector, match_index, url_pattern, open_url, navigate_url, options, created_at) VALUES (?, ?, ?, ?, ?, ?, ?)'
    ).run(selector, matchIndex, urlPattern, openUrl, navigateUrl, options, Date.now());

    db.prepare(`
        DELETE FROM autobuyer_locate
        WHERE id NOT IN (SELECT id FROM autobuyer_locate ORDER BY id DESC LIMIT ?)
    `).run(LOCATE_HISTORY);

    return db.prepare('SELECT * FROM autobuyer_locate WHERE id = ?').get(res.lastInsertRowid) as LocateRow;
}

export function getLocateRequest(id: number): LocateRow | undefined {
    return db.prepare('SELECT * FROM autobuyer_locate WHERE id = ?').get(id) as LocateRow | undefined;
}

export function countPendingLocate(): number {
    const row = db.prepare("SELECT COUNT(*) AS n FROM autobuyer_locate WHERE status = 'pending'").get() as { n: number };
    return row.n;
}

/** Hand the oldest pending request to the extension, marking it claimed so a
 *  slow round-trip doesn't cause the same click to be served twice. */
export function claimLocateRequest(): LocateRow | undefined {
    const row = db.prepare("SELECT * FROM autobuyer_locate WHERE status = 'pending' ORDER BY id LIMIT 1").get() as LocateRow | undefined;
    if (!row) return undefined;
    db.prepare("UPDATE autobuyer_locate SET status = 'claimed' WHERE id = ?").run(row.id);
    return { ...row, status: 'claimed' };
}

export function resolveLocateRequest(id: number, result: unknown | null, error: string | null): void {
    db.prepare('UPDATE autobuyer_locate SET status = ?, result = ?, error = ?, resolved_at = ? WHERE id = ?')
        .run(error ? 'error' : 'done', result ? JSON.stringify(result) : null, error, Date.now(), id);
}

// ── AutoBuyer automation runs ────────────────────────────────────────────────
//
// The dashboard queues a run; the Python runner claims it and works through the
// steps, reporting progress back so the page can show what is happening.

db.exec(`
    CREATE TABLE IF NOT EXISTS autobuyer_runs (
        id            INTEGER PRIMARY KEY AUTOINCREMENT,
        automation_id TEXT    NOT NULL,
        status        TEXT    NOT NULL DEFAULT 'queued',
        step_index    INTEGER NOT NULL DEFAULT 0,
        total_steps   INTEGER NOT NULL DEFAULT 0,
        log           TEXT    NOT NULL DEFAULT '[]',
        inputs        TEXT    NOT NULL DEFAULT '{}',
        error         TEXT,
        created_at    INTEGER NOT NULL,
        updated_at    INTEGER NOT NULL
    );
`);

// Migration: free-form per-request options, so a new kind of request doesn't
// need a new column each time.
try {
    db.exec("ALTER TABLE autobuyer_locate ADD COLUMN options TEXT NOT NULL DEFAULT '{}'");
} catch {
    // Column already exists
}

// Migration: the page to send the tab to before locating.
try {
    db.exec("ALTER TABLE autobuyer_locate ADD COLUMN navigate_url TEXT NOT NULL DEFAULT ''");
} catch {
    // Column already exists
}

// Migration: the page to open when no tab matches the pattern.
try {
    db.exec("ALTER TABLE autobuyer_locate ADD COLUMN open_url TEXT NOT NULL DEFAULT ''");
} catch {
    // Column already exists
}

// Migration: per-run inputs the user set on the dashboard.
try {
    db.exec("ALTER TABLE autobuyer_runs ADD COLUMN inputs TEXT NOT NULL DEFAULT '{}'");
} catch {
    // Column already exists
}

export type RunStatus = 'queued' | 'running' | 'done' | 'error' | 'cancelled';

export interface RunRow {
    id: number;
    automation_id: string;
    status: RunStatus;
    step_index: number;
    total_steps: number;
    log: string;      // JSON { at: number; step: string; ok: boolean; detail?: string }[]
    inputs: string;   // JSON Record<string, number>
    error: string | null;
    created_at: number;
    updated_at: number;
}

const RUN_HISTORY = 30;

export function createRun(automationId: string, totalSteps: number, inputs: Record<string, number> = {}): RunRow {
    const now = Date.now();
    const res = db.prepare(
        'INSERT INTO autobuyer_runs (automation_id, total_steps, inputs, created_at, updated_at) VALUES (?, ?, ?, ?, ?)'
    ).run(automationId, totalSteps, JSON.stringify(inputs), now, now);

    db.prepare(`
        DELETE FROM autobuyer_runs
        WHERE id NOT IN (SELECT id FROM autobuyer_runs ORDER BY id DESC LIMIT ?)
    `).run(RUN_HISTORY);

    return db.prepare('SELECT * FROM autobuyer_runs WHERE id = ?').get(res.lastInsertRowid) as RunRow;
}

export function getRun(id: number): RunRow | undefined {
    return db.prepare('SELECT * FROM autobuyer_runs WHERE id = ?').get(id) as RunRow | undefined;
}

export function getRecentRuns(limit = 5): RunRow[] {
    return db.prepare('SELECT * FROM autobuyer_runs ORDER BY id DESC LIMIT ?').all(limit) as RunRow[];
}

export function countActiveRuns(): number {
    const row = db.prepare("SELECT COUNT(*) AS n FROM autobuyer_runs WHERE status IN ('queued','running')").get() as { n: number };
    return row.n;
}

/** Grace period before a 'running' run with nobody behind it is written off.
 *  Long enough to cover the gap between the runner claiming a run and its next
 *  heartbeat reporting busy. */
const ORPHAN_GRACE_MS = 30_000;

/** Fail runs that were left mid-flight when their runner went away.
 *
 *  Only one run executes at a time, so a stale 'running' row blocks every future
 *  run — kill the daemon during a run (or restart it, which is the same thing)
 *  and nothing would ever be claimed again. The heartbeat says whether anything
 *  is actually executing: a runner that is gone, or up and idle, is not driving
 *  this run no matter what its status column claims.
 */
export function reapOrphanedRuns(): number {
    const hb = getRunnerHeartbeat();
    const alive = hb !== null && Date.now() - hb.at < RUNNER_TIMEOUT_MS;
    if (alive && hb!.busy) return 0;

    const now = Date.now();
    return db.prepare(`
        UPDATE autobuyer_runs
        SET status = 'error',
            error = 'the runner stopped while this run was in progress',
            updated_at = ?
        WHERE status = 'running' AND updated_at < ?
    `).run(now, now - ORPHAN_GRACE_MS).changes;
}

/** The runner takes the oldest queued run. Only one runs at a time — two
 *  processes driving the same physical mouse would interleave clicks. */
export function claimRun(): RunRow | undefined {
    reapOrphanedRuns();

    const running = db.prepare("SELECT 1 FROM autobuyer_runs WHERE status = 'running'").get();
    if (running) return undefined;

    const row = db.prepare("SELECT * FROM autobuyer_runs WHERE status = 'queued' ORDER BY id LIMIT 1").get() as RunRow | undefined;
    if (!row) return undefined;

    db.prepare("UPDATE autobuyer_runs SET status = 'running', updated_at = ? WHERE id = ?").run(Date.now(), row.id);
    return { ...row, status: 'running' };
}

export function appendRunLog(id: number, entry: unknown, stepIndex: number): void {
    const row = getRun(id);
    if (!row) return;
    let log: unknown[];
    try { log = JSON.parse(row.log); } catch { log = []; }
    log.push(entry);
    db.prepare('UPDATE autobuyer_runs SET log = ?, step_index = ?, updated_at = ? WHERE id = ?')
        .run(JSON.stringify(log), stepIndex, Date.now(), id);
}

export function finishRun(id: number, status: RunStatus, error: string | null): void {
    db.prepare('UPDATE autobuyer_runs SET status = ?, error = ?, updated_at = ? WHERE id = ?')
        .run(status, error, Date.now(), id);
}

/** Cancelling a queued run stops it starting; cancelling a running one is seen
 *  by the runner between steps. */
export function cancelRun(id: number): boolean {
    const row = getRun(id);
    if (!row || row.status === 'done' || row.status === 'error' || row.status === 'cancelled') return false;
    finishRun(id, 'cancelled', null);
    return true;
}

// ── Runner heartbeat ─────────────────────────────────────────────────────────
//
// The runner is a desktop process, so the server can't tell whether it's alive.
// It reports in every couple of seconds; if the beats stop, the dashboard greys
// out the buttons rather than queueing runs nobody will execute.

/** Three missed beats. Long enough to ride out a slow tick, short enough that a
 *  killed runner shows as offline before you press anything. */
export const RUNNER_TIMEOUT_MS = 7000;

/** The runner version this server's step vocabulary requires. A running process
 *  doesn't reload when the source changes, so an older one silently fails on
 *  steps it predates — the dashboard warns instead. */
export const RUNNER_EXPECTED_VERSION = '0.18.0';

export interface RunnerHeartbeat {
    at: number;
    host?: string;
    pid?: number;
    dryRun?: boolean;
    busy?: boolean;
    version?: string;
}

export function setRunnerHeartbeat(info: Omit<RunnerHeartbeat, 'at'>): void {
    setSetting('runner_heartbeat', JSON.stringify({ ...info, at: Date.now() }));
}

export function getRunnerHeartbeat(): RunnerHeartbeat | null {
    const raw = getSetting('runner_heartbeat');
    if (!raw) return null;
    try { return JSON.parse(raw) as RunnerHeartbeat; } catch { return null; }
}