crypto_monitor/lib/hedge_plan/hedge_plan_db.py

"""对冲计划 SQLite 表."""
from __future__ import annotations

import sqlite3
from typing import Any, Optional


def init_hedge_plan_tables(conn: sqlite3.Connection) -> None:
    conn.execute(
        """
        CREATE TABLE IF NOT EXISTS hedge_plans (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            plan_type TEXT NOT NULL,
            status TEXT NOT NULL,
            underlying TEXT NOT NULL,
            direction TEXT,
            entry_mark REAL,
            tp REAL,
            sl REAL,
            target_price REAL,
            sizing_mode_at_open TEXT,
            perp_size REAL,
            margin REAL,
            leverage REAL,
            premium_total REAL,
            realized_pnl_perp REAL,
            realized_pnl_options REAL,
            realized_pnl_total REAL,
            stats_bucket TEXT,
            close_reason TEXT,
            wechat_start_sent INTEGER DEFAULT 0,
            wechat_end_sent INTEGER DEFAULT 0,
            note TEXT,
            preview_json TEXT,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            opened_at TIMESTAMP,
            closed_at TIMESTAMP
        )
        """
    )
    conn.execute(
        """
        CREATE TABLE IF NOT EXISTS hedge_plan_legs (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            plan_id INTEGER NOT NULL,
            leg_role TEXT NOT NULL,
            symbol TEXT,
            inst_id TEXT,
            opt_type TEXT,
            strike REAL,
            side TEXT,
            size REAL,
            avg_open REAL,
            premium REAL,
            status TEXT,
            linked_monitor_id INTEGER,
            options_trade_id INTEGER,
            exchange_ord_id TEXT,
            realized_pnl REAL,
            close_reason TEXT,
            opened_at TIMESTAMP,
            closed_at TIMESTAMP,
            FOREIGN KEY(plan_id) REFERENCES hedge_plans(id)
        )
        """
    )
    conn.execute(
        "CREATE INDEX IF NOT EXISTS idx_hedge_plans_status ON hedge_plans(status)"
    )
    conn.execute(
        "CREATE INDEX IF NOT EXISTS idx_hedge_plan_legs_plan ON hedge_plan_legs(plan_id)"
    )


def count_active_plans(conn: sqlite3.Connection, plan_type: Optional[str] = None) -> int:
    if plan_type:
        row = conn.execute(
            "SELECT COUNT(1) AS c FROM hedge_plans WHERE status IN ('opening','active','partial') AND plan_type=?",
            (plan_type,),
        ).fetchone()
    else:
        row = conn.execute(
            "SELECT COUNT(1) AS c FROM hedge_plans WHERE status IN ('opening','active','partial')"
        ).fetchone()
    return int((row["c"] if row else 0) or 0)


def insert_plan(conn: sqlite3.Connection, row: dict[str, Any]) -> int:
    cols = list(row.keys())
    placeholders = ",".join(["?"] * len(cols))
    conn.execute(
        f"INSERT INTO hedge_plans ({','.join(cols)}) VALUES ({placeholders})",
        [row[c] for c in cols],
    )
    return int(conn.execute("SELECT last_insert_rowid()").fetchone()[0])


def insert_leg(conn: sqlite3.Connection, row: dict[str, Any]) -> int:
    cols = list(row.keys())
    placeholders = ",".join(["?"] * len(cols))
    conn.execute(
        f"INSERT INTO hedge_plan_legs ({','.join(cols)}) VALUES ({placeholders})",
        [row[c] for c in cols],
    )
    return int(conn.execute("SELECT last_insert_rowid()").fetchone()[0])


def update_plan(conn: sqlite3.Connection, plan_id: int, **fields: Any) -> None:
    if not fields:
        return
    sets = ", ".join(f"{k}=?" for k in fields)
    conn.execute(f"UPDATE hedge_plans SET {sets} WHERE id=?", [*fields.values(), plan_id])


def list_plans(
    conn: sqlite3.Connection,
    *,
    status: Optional[str] = None,
    plan_type: Optional[str] = None,
    underlying: Optional[str] = None,
    limit: int = 50,
) -> list[dict[str, Any]]:
    wheres: list[str] = []
    args: list[Any] = []
    if status:
        wheres.append("status=?")
        args.append(status)
    if plan_type:
        wheres.append("plan_type=?")
        args.append(plan_type)
    if underlying:
        wheres.append("underlying=?")
        args.append(underlying)
    where = (" WHERE " + " AND ".join(wheres)) if wheres else ""
    rows = conn.execute(
        f"SELECT * FROM hedge_plans{where} ORDER BY id DESC LIMIT ?",
        [*args, int(limit)],
    ).fetchall()
    return [dict(r) for r in rows]


def get_plan(conn: sqlite3.Connection, plan_id: int) -> Optional[dict[str, Any]]:
    row = conn.execute("SELECT * FROM hedge_plans WHERE id=?", (plan_id,)).fetchone()
    return dict(row) if row else None


def get_plan_legs(conn: sqlite3.Connection, plan_id: int) -> list[dict[str, Any]]:
    rows = conn.execute(
        "SELECT * FROM hedge_plan_legs WHERE plan_id=? ORDER BY id", (plan_id,)
    ).fetchall()
    return [dict(r) for r in rows]


def stats_summary(conn: sqlite3.Connection) -> dict[str, Any]:
    rows = conn.execute(
        """
        SELECT plan_type, close_reason, COUNT(1) AS n,
               COALESCE(SUM(realized_pnl_total), 0) AS pnl
        FROM hedge_plans
        WHERE status='closed'
        GROUP BY plan_type, close_reason
        """
    ).fetchall()
    closed = conn.execute(
        "SELECT COUNT(1) AS c, COALESCE(SUM(realized_pnl_total),0) AS pnl FROM hedge_plans WHERE status='closed'"
    ).fetchone()
    active = count_active_plans(conn)
    return {
        "active": active,
        "closed_count": int((closed["c"] if closed else 0) or 0),
        "closed_pnl_total": float((closed["pnl"] if closed else 0) or 0),
        "by_reason": [dict(r) for r in rows],
    }