market_intel/packages/db/repository.py

"""数据访问。"""

from __future__ import annotations

import json
import sqlite3
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from packages.db.schema import init_db


def _now_ms() -> int:
    return int(time.time() * 1000)


@dataclass
class OptionQuoteRow:
    ts_ms: int
    exchange: str
    underlying: str
    inst_id: str
    expiry_ymd: str
    strike: float
    side: str
    index_px: float
    ask: float | None
    bid: float | None
    ask_sz: float | None
    bid_sz: float | None
    leverage: float | None


class Repository:
    def __init__(self, db_path: str | Path) -> None:
        self.db_path = Path(db_path)
        self.conn = init_db(self.db_path)

    def close(self) -> None:
        self.conn.close()

    def insert_option_quote(self, row: OptionQuoteRow) -> int:
        cur = self.conn.execute(
            """
            INSERT INTO option_quotes (
                ts_ms, exchange, underlying, inst_id, expiry_ymd, strike, side,
                index_px, ask, bid, ask_sz, bid_sz, leverage, created_at_ms
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            (
                row.ts_ms,
                row.exchange,
                row.underlying,
                row.inst_id,
                row.expiry_ymd,
                row.strike,
                row.side,
                row.index_px,
                row.ask,
                row.bid,
                row.ask_sz,
                row.bid_sz,
                row.leverage,
                _now_ms(),
            ),
        )
        self.conn.commit()
        return int(cur.lastrowid)

    def insert_index_tick(
        self,
        *,
        ts_ms: int,
        exchange: str,
        underlying: str,
        index_px: float,
    ) -> int:
        cur = self.conn.execute(
            """
            INSERT INTO index_ticks (ts_ms, exchange, underlying, index_px, created_at_ms)
            VALUES (?, ?, ?, ?, ?)
            """,
            (ts_ms, exchange, underlying, index_px, _now_ms()),
        )
        self.conn.commit()
        return int(cur.lastrowid)

    def upsert_heartbeat(
        self,
        *,
        ok: bool,
        error: str | None = None,
        meta: dict[str, Any] | None = None,
    ) -> None:
        now = _now_ms()
        row = self.conn.execute(
            "SELECT consecutive_failures FROM collector_heartbeat WHERE id = 1"
        ).fetchone()
        fails = int(row["consecutive_failures"] if row else 0)
        if ok:
            fails = 0
            self.conn.execute(
                """
                UPDATE collector_heartbeat
                SET last_ok_ts_ms = ?, last_error = NULL, consecutive_failures = 0,
                    meta_json = COALESCE(?, meta_json)
                WHERE id = 1
                """,
                (now, json.dumps(meta, ensure_ascii=False) if meta else None),
            )
        else:
            fails += 1
            self.conn.execute(
                """
                UPDATE collector_heartbeat
                SET last_error = ?, last_error_ts_ms = ?, consecutive_failures = ?,
                    meta_json = COALESCE(?, meta_json)
                WHERE id = 1
                """,
                (
                    (error or "unknown")[:2000],
                    now,
                    fails,
                    json.dumps(meta, ensure_ascii=False) if meta else None,
                ),
            )
        self.conn.commit()

    def get_heartbeat(self) -> dict[str, Any]:
        row = self.conn.execute(
            "SELECT * FROM collector_heartbeat WHERE id = 1"
        ).fetchone()
        if not row:
            return {}
        d = dict(row)
        meta = d.get("meta_json")
        if meta:
            try:
                d["meta"] = json.loads(meta)
            except json.JSONDecodeError:
                d["meta"] = None
        else:
            d["meta"] = None
        return d

    def latest_quotes_by_side(self) -> dict[str, dict[str, Any]]:
        """返回 side -> 最新一条。"""
        out: dict[str, dict[str, Any]] = {}
        for side in ("C", "P"):
            row = self.conn.execute(
                """
                SELECT * FROM option_quotes
                WHERE side = ?
                ORDER BY ts_ms DESC, id DESC
                LIMIT 1
                """,
                (side,),
            ).fetchone()
            if row:
                out[side] = dict(row)
        return out

    def count_option_quotes(self) -> int:
        row = self.conn.execute("SELECT COUNT(*) AS n FROM option_quotes").fetchone()
        return int(row["n"] if row else 0)

    def count_index_ticks(self) -> int:
        row = self.conn.execute("SELECT COUNT(*) AS n FROM index_ticks").fetchone()
        return int(row["n"] if row else 0)

    def fetch_option_quotes(
        self,
        *,
        start_ms: int,
        end_ms: int,
        side: str = "both",
        underlying: str | None = None,
    ) -> list[dict[str, Any]]:
        """[start_ms, end_ms) 半开区间。"""
        clauses = ["ts_ms >= ?", "ts_ms < ?"]
        params: list[Any] = [int(start_ms), int(end_ms)]
        want = (side or "both").upper()
        if want in ("C", "P"):
            clauses.append("side = ?")
            params.append(want)
        if underlying:
            clauses.append("underlying = ?")
            params.append(underlying)
        sql = f"""
            SELECT ts_ms, exchange, underlying, inst_id, expiry_ymd, strike, side,
                   index_px, ask, bid, ask_sz, bid_sz, leverage
            FROM option_quotes
            WHERE {' AND '.join(clauses)}
            ORDER BY ts_ms ASC, id ASC
        """
        rows = self.conn.execute(sql, params).fetchall()
        return [dict(r) for r in rows]

    def get_settlement(self, expiry_ymd: str) -> dict[str, Any] | None:
        row = self.conn.execute(
            "SELECT * FROM expiry_settlements WHERE expiry_ymd = ?",
            (expiry_ymd,),
        ).fetchone()
        return dict(row) if row else None

    def list_settlements(self, ymds: list[str] | None = None) -> dict[str, dict[str, Any]]:
        if ymds is not None and not ymds:
            return {}
        if ymds is None:
            rows = self.conn.execute("SELECT * FROM expiry_settlements").fetchall()
        else:
            placeholders = ",".join("?" for _ in ymds)
            rows = self.conn.execute(
                f"SELECT * FROM expiry_settlements WHERE expiry_ymd IN ({placeholders})",
                list(ymds),
            ).fetchall()
        return {str(r["expiry_ymd"]): dict(r) for r in rows}

    def upsert_settlement(
        self,
        *,
        expiry_ymd: str,
        settle_ts_ms: int,
        settle_index_px: float,
        exchange: str,
        underlying: str,
    ) -> None:
        self.conn.execute(
            """
            INSERT INTO expiry_settlements (
                expiry_ymd, settle_ts_ms, settle_index_px, exchange, underlying, created_at_ms
            ) VALUES (?, ?, ?, ?, ?, ?)
            ON CONFLICT(expiry_ymd) DO UPDATE SET
                settle_ts_ms = excluded.settle_ts_ms,
                settle_index_px = excluded.settle_index_px,
                exchange = excluded.exchange,
                underlying = excluded.underlying
            """,
            (
                expiry_ymd,
                int(settle_ts_ms),
                float(settle_index_px),
                exchange,
                underlying,
                _now_ms(),
            ),
        )
        self.conn.commit()

    def nearest_index_tick(
        self,
        *,
        underlying: str,
        target_ts_ms: int,
        max_delta_ms: int,
    ) -> dict[str, Any] | None:
        row = self.conn.execute(
            """
            SELECT ts_ms, index_px, ABS(ts_ms - ?) AS delta
            FROM index_ticks
            WHERE underlying = ?
              AND ts_ms BETWEEN ? AND ?
            ORDER BY delta ASC
            LIMIT 1
            """,
            (
                int(target_ts_ms),
                underlying,
                int(target_ts_ms) - int(max_delta_ms),
                int(target_ts_ms) + int(max_delta_ms),
            ),
        ).fetchone()
        return dict(row) if row else None