Chan/crypto_wyckoff/store.py

"""SQLite persistence for crypto wyckoff scan rows."""

from __future__ import annotations

import json
import sqlite3
from datetime import datetime
from typing import Any

from crypto_wyckoff.domain_models import WyckoffScanRow
from crypto_wyckoff.io import SCAN_DB, ensure_dirs

_COLS = [
    "trade_date", "ts_code", "name", "industry", "engine_version",
    "m_cycle", "cycle_confidence", "trend_score",
    "w_cycle", "w_phase", "w_current_event", "w_recent_events_json",
    "phase_confidence", "structure_score",
    "d_current_event", "d_recent_events_json", "event_confidence", "entry_score",
    "entry", "stop", "target1", "target2", "rr",
    "alignment", "stars", "decision_signal", "signal_confidence",
    "overall_confidence", "overall_score", "risk", "reasons_json",
    "feature_snapshot_json", "markers_json", "scanned_at",
]


def _conn() -> sqlite3.Connection:
    ensure_dirs()
    c = sqlite3.connect(str(SCAN_DB), timeout=60)
    c.row_factory = sqlite3.Row
    c.execute(
        """
        CREATE TABLE IF NOT EXISTS wyckoff_scan (
            trade_date TEXT NOT NULL,
            ts_code TEXT NOT NULL,
            name TEXT DEFAULT '',
            industry TEXT DEFAULT '',
            engine_version TEXT,
            m_cycle TEXT, cycle_confidence REAL, trend_score REAL,
            w_cycle TEXT, w_phase TEXT, w_current_event TEXT, w_recent_events_json TEXT,
            phase_confidence REAL, structure_score REAL,
            d_current_event TEXT, d_recent_events_json TEXT, event_confidence REAL, entry_score REAL,
            entry REAL, stop REAL, target1 REAL, target2 REAL, rr REAL,
            alignment REAL, stars INTEGER, decision_signal TEXT, signal_confidence REAL,
            overall_confidence REAL, overall_score REAL, risk TEXT, reasons_json TEXT,
            feature_snapshot_json TEXT, markers_json TEXT, scanned_at TEXT,
            PRIMARY KEY (trade_date, ts_code)
        )
        """
    )
    c.execute(
        "CREATE INDEX IF NOT EXISTS idx_cw_score ON wyckoff_scan(trade_date, overall_score DESC)"
    )
    return c


def upsert_row(row: WyckoffScanRow) -> None:
    vals = (
        row.trade_date.isoformat() if hasattr(row.trade_date, "isoformat") else str(row.trade_date),
        row.ts_code, row.name, row.industry, row.engine_version,
        row.m_cycle, row.cycle_confidence, row.trend_score,
        row.w_cycle, row.w_phase, row.w_current_event, row.w_recent_events_json,
        row.phase_confidence, row.structure_score,
        row.d_current_event, row.d_recent_events_json, row.event_confidence, row.entry_score,
        row.entry, row.stop, row.target1, row.target2, row.rr,
        row.alignment, row.stars, row.decision_signal, row.signal_confidence,
        row.overall_confidence, row.overall_score, row.risk, row.reasons_json,
        row.feature_snapshot_json, row.markers_json,
        row.scanned_at.isoformat() if isinstance(row.scanned_at, datetime) else str(row.scanned_at),
    )
    c = _conn()
    try:
        placeholders = ",".join("?" * len(_COLS))
        col_sql = ",".join(_COLS)
        updates = ",".join(f"{c}=excluded.{c}" for c in _COLS if c not in ("trade_date", "ts_code"))
        c.execute(
            f"""
            INSERT INTO wyckoff_scan ({col_sql}) VALUES ({placeholders})
            ON CONFLICT(trade_date, ts_code) DO UPDATE SET {updates}
            """,
            vals,
        )
        c.commit()
    finally:
        c.close()


def latest_trade_date() -> str | None:
    c = _conn()
    try:
        cur = c.execute("SELECT MAX(trade_date) FROM wyckoff_scan")
        row = cur.fetchone()
        return row[0] if row and row[0] else None
    finally:
        c.close()


def count_for_date(trade_date: str | None = None) -> int:
    td = trade_date or latest_trade_date()
    if not td:
        return 0
    c = _conn()
    try:
        cur = c.execute("SELECT COUNT(*) FROM wyckoff_scan WHERE trade_date=?", (td,))
        return int(cur.fetchone()[0])
    finally:
        c.close()


def query_scan(
    *,
    trade_date: str | None = None,
    m_cycle: str | None = None,
    w_phase: str | None = None,
    d_event: str | None = None,
    decision_signal: str | None = None,
    min_overall_score: float | None = None,
    min_alignment: float | None = None,
    sort: str = "overall_score",
    limit: int = 100,
    offset: int = 0,
) -> list[dict[str, Any]]:
    td = trade_date or latest_trade_date()
    if not td:
        return []
    sort_col = sort if sort in {
        "overall_score", "alignment", "entry_score", "trend_score", "structure_score", "stars"
    } else "overall_score"
    clauses = ["trade_date=?"]
    args: list[Any] = [td]
    if m_cycle:
        clauses.append("m_cycle=?")
        args.append(m_cycle)
    if w_phase:
        clauses.append("w_phase=?")
        args.append(w_phase)
    if d_event:
        clauses.append("d_current_event=?")
        args.append(d_event)
    if decision_signal:
        clauses.append("decision_signal=?")
        args.append(decision_signal)
    if min_overall_score is not None:
        clauses.append("overall_score>=?")
        args.append(min_overall_score)
    if min_alignment is not None:
        clauses.append("alignment>=?")
        args.append(min_alignment)
    where = " AND ".join(clauses)
    args.extend([limit, offset])
    c = _conn()
    try:
        cur = c.execute(
            f"SELECT * FROM wyckoff_scan WHERE {where} ORDER BY {sort_col} DESC LIMIT ? OFFSET ?",
            args,
        )
        return [dict(r) for r in cur.fetchall()]
    finally:
        c.close()


def get_symbol(ts_code: str, trade_date: str | None = None) -> dict[str, Any] | None:
    td = trade_date or latest_trade_date()
    if not td:
        return None
    c = _conn()
    try:
        cur = c.execute(
            "SELECT * FROM wyckoff_scan WHERE trade_date=? AND ts_code=?",
            (td, ts_code),
        )
        row = cur.fetchone()
        return dict(row) if row else None
    finally:
        c.close()