tvsignals-to-tg/app/state.py

from __future__ import annotations

import sqlite3
import time
from pathlib import Path


class ScannerStore:
    """Persist last processed bar per (strategy, symbol) and sent Heryon nonces."""

    def __init__(self, path: str | Path) -> None:
        self._path = Path(path)
        self._path.parent.mkdir(parents=True, exist_ok=True)
        self._conn = sqlite3.connect(self._path, check_same_thread=False)
        self._conn.execute("PRAGMA journal_mode=WAL")
        self._conn.execute(
            """
            CREATE TABLE IF NOT EXISTS processed (
                strategy TEXT NOT NULL,
                symbol TEXT NOT NULL,
                bar_open_ts INTEGER NOT NULL,
                PRIMARY KEY (strategy, symbol)
            )
            """
        )
        self._conn.execute(
            """
            CREATE TABLE IF NOT EXISTS sent_nonces (
                nonce TEXT PRIMARY KEY,
                sent_at INTEGER NOT NULL
            )
            """
        )
        self._conn.execute(
            """
            CREATE TABLE IF NOT EXISTS last_open (
                strategy TEXT NOT NULL,
                symbol TEXT NOT NULL,
                side TEXT NOT NULL,
                entry REAL NOT NULL,
                bar_open_ts INTEGER NOT NULL DEFAULT 0,
                PRIMARY KEY (strategy, symbol)
            )
            """
        )
        self._conn.execute(
            """
            CREATE TABLE IF NOT EXISTS last_tg_message (
                thread_id INTEGER NOT NULL,
                ticker TEXT NOT NULL,
                message_id INTEGER NOT NULL,
                PRIMARY KEY (thread_id, ticker)
            )
            """
        )
        last_open_cols = {
            row[1] for row in self._conn.execute("PRAGMA table_info(last_open)")
        }
        if "bar_open_ts" not in last_open_cols:
            self._conn.execute(
                "ALTER TABLE last_open ADD COLUMN bar_open_ts INTEGER NOT NULL DEFAULT 0"
            )
        self._conn.commit()

    def get_last_bar(self, strategy: str, symbol: str) -> int | None:
        row = self._conn.execute(
            "SELECT bar_open_ts FROM processed WHERE strategy = ? AND symbol = ?",
            (strategy, symbol),
        ).fetchone()
        return int(row[0]) if row else None

    def set_last_bar(self, strategy: str, symbol: str, bar_open_ts: int) -> None:
        self._conn.execute(
            """
            INSERT INTO processed (strategy, symbol, bar_open_ts)
            VALUES (?, ?, ?)
            ON CONFLICT(strategy, symbol) DO UPDATE SET bar_open_ts = excluded.bar_open_ts
            """,
            (strategy, symbol, bar_open_ts),
        )
        self._conn.commit()

    def nonce_sent(self, nonce: str) -> bool:
        row = self._conn.execute(
            "SELECT 1 FROM sent_nonces WHERE nonce = ?", (nonce,)
        ).fetchone()
        return row is not None

    def mark_nonce(self, nonce: str) -> None:
        self._conn.execute(
            "INSERT OR IGNORE INTO sent_nonces (nonce, sent_at) VALUES (?, ?)",
            (nonce, int(time.time())),
        )
        self._conn.commit()

    def get_last_open(self, strategy: str, symbol: str) -> tuple[str, float, int] | None:
        row = self._conn.execute(
            "SELECT side, entry, bar_open_ts FROM last_open WHERE strategy = ? AND symbol = ?",
            (strategy, symbol),
        ).fetchone()
        if row is None:
            return None
        return str(row[0]), float(row[1]), int(row[2] or 0)

    def set_last_open(
        self,
        strategy: str,
        symbol: str,
        side: str,
        entry: float,
        bar_open_ts: int,
    ) -> None:
        self._conn.execute(
            """
            INSERT INTO last_open (strategy, symbol, side, entry, bar_open_ts)
            VALUES (?, ?, ?, ?, ?)
            ON CONFLICT(strategy, symbol) DO UPDATE SET
                side = excluded.side,
                entry = excluded.entry,
                bar_open_ts = excluded.bar_open_ts
            """,
            (strategy, symbol, side, entry, bar_open_ts),
        )
        self._conn.commit()

    def get_last_tg_message(self, thread_id: int, ticker: str) -> int | None:
        row = self._conn.execute(
            "SELECT message_id FROM last_tg_message WHERE thread_id = ? AND ticker = ?",
            (thread_id, ticker.strip().upper()),
        ).fetchone()
        return int(row[0]) if row else None

    def set_last_tg_message(self, thread_id: int, ticker: str, message_id: int) -> None:
        self._conn.execute(
            """
            INSERT INTO last_tg_message (thread_id, ticker, message_id)
            VALUES (?, ?, ?)
            ON CONFLICT(thread_id, ticker) DO UPDATE SET message_id = excluded.message_id
            """,
            (thread_id, ticker.strip().upper(), message_id),
        )
        self._conn.commit()

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