database/resolution_repository.py

"""SQLite-backed helpers for five-session prediction resolution."""

from __future__ import annotations

from datetime import date, datetime, timezone
from typing import Any, Callable

from jobs.resolver import resolve_prediction_rows


def fetch_unresolved_predictions(conn) -> list[dict[str, Any]]:
    with conn.cursor() as cur:
        cur.execute("""
            SELECT id, prediction_date, symbol, target_threshold, target_horizon_days
            FROM predictions WHERE resolved_at IS NULL ORDER BY prediction_date, symbol
        """)
        return cur.fetchall()


def fetch_future_ohlcv(conn, symbol: str, prediction_date: date, forward_days: int) -> list[dict[str, Any]]:
    with conn.cursor() as cur:
        cur.execute("""
            SELECT o.trading_date AS timestamp, o.open, o.close
            FROM ohlcv o JOIN market_tickers t ON t.id = o.ticker_id
            WHERE t.symbol = ? AND t.market = 'NSE' AND o.trading_date > ?
            ORDER BY o.trading_date ASC LIMIT ?
        """, (symbol, prediction_date, forward_days))
        return cur.fetchall()


def resolve_pending_predictions(
    conn,
    *,
    on_progress: Callable[[int, int, int], None] | None = None,
    commit_interval: int = 25,
) -> int:
    resolved_count = 0
    failed_count = 0
    processed_count = 0

    def checkpoint() -> None:
        if processed_count % commit_interval == 0:
            conn.commit()
            if on_progress is not None:
                on_progress(processed_count, resolved_count, failed_count)

    for prediction in fetch_unresolved_predictions(conn):
        rows = fetch_future_ohlcv(conn, prediction["symbol"], prediction["prediction_date"], prediction["target_horizon_days"])
        try:
            resolution = resolve_prediction_rows(
                prediction_date=prediction["prediction_date"], symbol=prediction["symbol"],
                future_rows=[{"date": row["timestamp"], "open": row["open"], "close": row["close"]} for row in rows],
                target_return=prediction["target_threshold"], forward_days=prediction["target_horizon_days"],
            )
        except ValueError:
            failed_count += 1
        else:
            if resolution is not None:
                with conn.cursor() as cur:
                    cur.execute("""
                        UPDATE predictions SET entry_date = ?, entry_open = ?, max_close_5d = ?,
                            evaluation_end_date = ?, actual_return = ?, actual_label = ?, resolved_at = ?
                        WHERE id = ? AND resolved_at IS NULL
                    """, (resolution.entry_date, resolution.entry_open, resolution.max_close,
                          resolution.evaluation_dates[-1], resolution.actual_return, resolution.actual_label,
                          datetime.now(timezone.utc), prediction["id"]))
                    resolved_count += cur.rowcount
        processed_count += 1
        checkpoint()
    conn.commit()
    if on_progress is not None:
        on_progress(processed_count, resolved_count, failed_count)
    return resolved_count