eth_hedge_sim/backend/app/models/db.py

from __future__ import annotations

import sqlite3
import threading
import time
from pathlib import Path
from typing import Any

from ..config import Settings, get_settings

_SCHEMA = """
CREATE TABLE IF NOT EXISTS settings (
  key TEXT PRIMARY KEY,
  value TEXT NOT NULL
);

CREATE TABLE IF NOT EXISTS ledger_meta (
  id INTEGER PRIMARY KEY CHECK (id = 1),
  equity REAL NOT NULL,
  available REAL NOT NULL,
  reserved REAL NOT NULL DEFAULT 0,
  updated_at_ms INTEGER NOT NULL
);

CREATE TABLE IF NOT EXISTS groups (
  group_id TEXT PRIMARY KEY,
  status TEXT NOT NULL,
  bias TEXT,
  option_side TEXT,
  perp_side TEXT,
  option_inst_id TEXT,
  perp_inst_id TEXT,
  strike REAL,
  expiry_ymd TEXT,
  entry_index_px REAL,
  initial_premium REAL DEFAULT 0,
  open_at_ms INTEGER,
  close_at_ms INTEGER,
  close_reason TEXT,
  realized_pnl REAL DEFAULT 0,
  fees REAL DEFAULT 0,
  slip_cost REAL DEFAULT 0,
  note TEXT
);

CREATE TABLE IF NOT EXISTS fills (
  id INTEGER PRIMARY KEY AUTOINCREMENT,
  group_id TEXT NOT NULL,
  leg TEXT NOT NULL,
  action TEXT NOT NULL,
  side TEXT NOT NULL,
  inst_id TEXT NOT NULL,
  qty_eth REAL NOT NULL,
  qty_contracts REAL,
  base_px REAL,
  fill_px REAL NOT NULL,
  fee REAL NOT NULL,
  slip REAL NOT NULL,
  notional REAL NOT NULL,
  ts_ms INTEGER NOT NULL,
  FOREIGN KEY(group_id) REFERENCES groups(group_id)
);

CREATE TABLE IF NOT EXISTS positions (
  id INTEGER PRIMARY KEY CHECK (id = 1),
  group_id TEXT,
  perp_side TEXT,
  perp_qty_eth REAL DEFAULT 0,
  perp_entry_px REAL,
  option_inst_id TEXT,
  option_side TEXT,
  option_qty_eth REAL DEFAULT 0,
  option_qty_contracts REAL DEFAULT 0,
  option_entry_px REAL,
  entry_index_px REAL,
  initial_premium REAL DEFAULT 0,
  status TEXT
);

CREATE TABLE IF NOT EXISTS ledger_entries (
  id INTEGER PRIMARY KEY AUTOINCREMENT,
  group_id TEXT,
  kind TEXT NOT NULL,
  amount REAL NOT NULL,
  balance_after REAL NOT NULL,
  note TEXT,
  ts_ms INTEGER NOT NULL
);

CREATE TABLE IF NOT EXISTS strategy_state (
  id INTEGER PRIMARY KEY CHECK (id = 1),
  running INTEGER NOT NULL DEFAULT 0,
  phase TEXT NOT NULL DEFAULT 'idle',
  rounds_done INTEGER NOT NULL DEFAULT 0,
  window_key TEXT,
  rest_until_ms INTEGER,
  last_error TEXT,
  updated_at_ms INTEGER NOT NULL
);
"""


def default_db_path(settings: Settings | None = None) -> Path:
    s = settings or get_settings()
    if s.db_path:
        return Path(s.db_path)
    root = Path(__file__).resolve().parents[2]  # backend/
    return root / "data" / "hedge.db"


class Database:
    def __init__(self, path: Path | None = None) -> None:
        self.path = path or default_db_path()
        self.path.parent.mkdir(parents=True, exist_ok=True)
        self._lock = threading.RLock()
        self._conn = sqlite3.connect(str(self.path), check_same_thread=False)
        self._conn.row_factory = sqlite3.Row
        self._conn.execute("PRAGMA journal_mode=WAL;")
        self._conn.executescript(_SCHEMA)
        self._conn.commit()
        self._ensure_seed()

    def close(self) -> None:
        with self._lock:
            self._conn.close()

    def _ensure_seed(self) -> None:
        s = get_settings()
        now = int(time.time() * 1000)
        with self._lock:
            row = self._conn.execute("SELECT id FROM ledger_meta WHERE id=1").fetchone()
            if row is None:
                self._conn.execute(
                    "INSERT INTO ledger_meta(id, equity, available, reserved, updated_at_ms) VALUES (1,?,?,0,?)",
                    (s.initial_equity, s.initial_equity, now),
                )
            pos = self._conn.execute("SELECT id FROM positions WHERE id=1").fetchone()
            if pos is None:
                self._conn.execute(
                    "INSERT INTO positions(id, status) VALUES (1, 'flat')"
                )
            st = self._conn.execute("SELECT id FROM strategy_state WHERE id=1").fetchone()
            if st is None:
                self._conn.execute(
                    "INSERT INTO strategy_state(id, running, phase, rounds_done, updated_at_ms) VALUES (1,0,'idle',0,?)",
                    (now,),
                )
            defaults = {
                "fee_rate": str(s.fee_rate),
                "initial_equity": str(s.initial_equity),
                "exit_move_points": str(s.exit_move_points),
                "rest_seconds": str(s.rest_seconds),
                "max_rounds": str(s.max_rounds),
                "perp_qty_eth": str(s.perp_qty_eth),
                "option_qty_eth": str(s.option_qty_eth),
            }
            for k, v in defaults.items():
                exists = self._conn.execute(
                    "SELECT key FROM settings WHERE key=?", (k,)
                ).fetchone()
                if exists is None:
                    self._conn.execute(
                        "INSERT INTO settings(key, value) VALUES (?,?)", (k, v)
                    )
            self._conn.commit()

    def execute(self, sql: str, params: tuple[Any, ...] | list[Any] = ()) -> sqlite3.Cursor:
        with self._lock:
            cur = self._conn.execute(sql, params)
            self._conn.commit()
            return cur

    def executemany(self, sql: str, seq: list[tuple[Any, ...]]) -> None:
        with self._lock:
            self._conn.executemany(sql, seq)
            self._conn.commit()

    def fetchone(self, sql: str, params: tuple[Any, ...] | list[Any] = ()) -> sqlite3.Row | None:
        with self._lock:
            return self._conn.execute(sql, params).fetchone()

    def fetchall(self, sql: str, params: tuple[Any, ...] | list[Any] = ()) -> list[sqlite3.Row]:
        with self._lock:
            return list(self._conn.execute(sql, params).fetchall())

    def get_setting(self, key: str, default: str | None = None) -> str | None:
        row = self.fetchone("SELECT value FROM settings WHERE key=?", (key,))
        if row is None:
            return default
        return str(row["value"])

    def set_setting(self, key: str, value: str) -> None:
        self.execute(
            "INSERT INTO settings(key, value) VALUES(?,?) ON CONFLICT(key) DO UPDATE SET value=excluded.value",
            (key, value),
        )


_db: Database | None = None


def get_db() -> Database:
    global _db
    if _db is None:
        _db = Database()
    return _db


def set_db(db: Database | None) -> None:
    global _db
    _db = db