Chan/crypto_wyckoff/store.py

"""SQLite persistence for crypto wyckoff scan rows (per combo)."""

from __future__ import annotations

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", "combo_id", "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",
]

_CREATE_SQL = """
CREATE TABLE IF NOT EXISTS wyckoff_scan (
    trade_date TEXT NOT NULL,
    combo_id TEXT NOT NULL DEFAULT 'd_w_m',
    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, combo_id, ts_code)
)
"""


def _migrate(c: sqlite3.Connection) -> None:
    cur = c.execute(
        "SELECT name FROM sqlite_master WHERE type='table' AND name='wyckoff_scan'"
    )
    if not cur.fetchone():
        c.execute(_CREATE_SQL)
        c.execute(
            "CREATE INDEX IF NOT EXISTS idx_cw_score "
            "ON wyckoff_scan(trade_date, combo_id, overall_score DESC)"
        )
        return

    cols = {r[1] for r in c.execute("PRAGMA table_info(wyckoff_scan)")}
    if "combo_id" in cols:
        c.execute(
            "CREATE INDEX IF NOT EXISTS idx_cw_score "
            "ON wyckoff_scan(trade_date, combo_id, overall_score DESC)"
        )
        return

    # Legacy PK (trade_date, ts_code) → add combo_id via table rebuild
    c.execute("ALTER TABLE wyckoff_scan RENAME TO wyckoff_scan_old")
    c.execute(_CREATE_SQL)
    old_cols = [r[1] for r in c.execute("PRAGMA table_info(wyckoff_scan_old)")]
    shared = [col for col in _COLS if col != "combo_id" and col in old_cols]
    col_sql = ",".join(shared)
    c.execute(
        f"""
        INSERT INTO wyckoff_scan (combo_id, {col_sql})
        SELECT 'd_w_m', {col_sql} FROM wyckoff_scan_old
        """
    )
    c.execute("DROP TABLE wyckoff_scan_old")
    c.execute(
        "CREATE INDEX IF NOT EXISTS idx_cw_score "
        "ON wyckoff_scan(trade_date, combo_id, overall_score DESC)"
    )


def _conn() -> sqlite3.Connection:
    ensure_dirs()
    c = sqlite3.connect(str(SCAN_DB), timeout=60)
    c.row_factory = sqlite3.Row
    _migrate(c)
    c.commit()
    return c


def upsert_row(row: WyckoffScanRow) -> None:
    combo_id = getattr(row, "combo_id", None) or "d_w_m"
    vals = (
        row.trade_date.isoformat() if hasattr(row.trade_date, "isoformat") else str(row.trade_date),
        combo_id,
        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"{col}=excluded.{col}"
            for col in _COLS
            if col not in ("trade_date", "combo_id", "ts_code")
        )
        c.execute(
            f"""
            INSERT INTO wyckoff_scan ({col_sql}) VALUES ({placeholders})
            ON CONFLICT(trade_date, combo_id, ts_code) DO UPDATE SET {updates}
            """,
            vals,
        )
        c.commit()
    finally:
        c.close()


def latest_trade_date(combo_id: str | None = None) -> str | None:
    c = _conn()
    try:
        if combo_id:
            cur = c.execute(
                "SELECT MAX(trade_date) FROM wyckoff_scan WHERE combo_id=?",
                (combo_id,),
            )
        else:
            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, combo_id: str | None = None) -> int:
    td = trade_date or latest_trade_date(combo_id)
    if not td:
        return 0
    c = _conn()
    try:
        if combo_id:
            cur = c.execute(
                "SELECT COUNT(*) FROM wyckoff_scan WHERE trade_date=? AND combo_id=?",
                (td, combo_id),
            )
        else:
            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,
    combo_id: 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]]:
    cid = combo_id or "d_w_m"
    td = trade_date or latest_trade_date(cid)
    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=?", "combo_id=?"]
    args: list[Any] = [td, cid]
    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,
    combo_id: str | None = None,
) -> dict[str, Any] | None:
    cid = combo_id or "d_w_m"
    td = trade_date or latest_trade_date(cid)
    if not td:
        return None
    c = _conn()
    try:
        cur = c.execute(
            "SELECT * FROM wyckoff_scan WHERE trade_date=? AND combo_id=? AND ts_code=?",
            (td, cid, ts_code),
        )
        row = cur.fetchone()
        return dict(row) if row else None
    finally:
        c.close()