trifecta_bro/data/storage.py

"""SQLite persistence for races, runners, predictions and results.

No external DB dependency — uses the stdlib sqlite3 module.
Models are versioned so historical predictions can be compared over time.
"""
from __future__ import annotations

import json
import sqlite3
from pathlib import Path
from typing import Optional

from ..config import settings
from .models import RaceModel, RunnerModel


_SCHEMA = """
CREATE TABLE IF NOT EXISTS races (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    date TEXT NOT NULL,
    track TEXT NOT NULL,
    track_slug TEXT,
    race_number INTEGER NOT NULL,
    race_name TEXT,
    distance TEXT,
    condition TEXT,
    race_class TEXT,
    abandoned INTEGER DEFAULT 0,
    start_time TEXT,
    timezone TEXT,
    number_of_runners INTEGER,
    UNIQUE(date, track_slug, race_number)
);
CREATE TABLE IF NOT EXISTS runners (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    race_id INTEGER NOT NULL,
    number INTEGER,
    name TEXT,
    jockey TEXT,
    trainer TEXT,
    barrier INTEGER,
    weight REAL,
    form TEXT,
    data_completeness REAL,
    FOREIGN KEY(race_id) REFERENCES races(id)
);
CREATE TABLE IF NOT EXISTS predictions (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    date TEXT NOT NULL,
    track TEXT NOT NULL,
    track_slug TEXT,
    race_number INTEGER NOT NULL,
    model_version TEXT,
    generated_at TEXT,
    top3 TEXT,                 -- json list of runner numbers
    trifecta TEXT,             -- json [1st,2nd,3rd]
    alternatives TEXT,         -- json list of [a,b,c]
    confidence TEXT,
    confidence_score REAL,
    ranked_runners TEXT,       -- json full ranking w/ scores
    data_completeness REAL,
    UNIQUE(date, track_slug, race_number, model_version)
);
CREATE TABLE IF NOT EXISTS results (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    date TEXT NOT NULL,
    track TEXT NOT NULL,
    track_slug TEXT,
    race_number INTEGER NOT NULL,
    actual_trifecta TEXT,      -- json [1st,2nd,3rd] or nulls
    exact_hit INTEGER,
    box_hit INTEGER,
    top1_hit INTEGER,
    UNIQUE(date, track_slug, race_number)
);
"""


class Storage:
    def __init__(self, db_path: Optional[Path] = None) -> None:
        self.db_path = Path(db_path or settings.db_path)
        self.db_path.parent.mkdir(parents=True, exist_ok=True)
        self.conn = sqlite3.connect(str(self.db_path))
        self.conn.row_factory = sqlite3.Row
        self.conn.executescript(_SCHEMA)
        self.conn.commit()

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

    def upsert_race(self, race: RaceModel) -> int:
        cur = self.conn.cursor()
        cur.execute(
            """INSERT INTO races
               (date, track, track_slug, race_number, race_name, distance, condition,
                race_class, abandoned, start_time, timezone, number_of_runners)
               VALUES (?,?,?,?,?,?,?,?,?,?,?,?)
               ON CONFLICT(date, track_slug, race_number)
               DO UPDATE SET race_name=excluded.race_name, distance=excluded.distance,
                 condition=excluded.condition, race_class=excluded.race_class,
                 abandoned=excluded.abandoned, start_time=excluded.start_time,
                 timezone=excluded.timezone, number_of_runners=excluded.number_of_runners""",
            (race.date, race.track, race.track_slug, race.race_number, race.race_name,
             race.distance, race.condition, race.race_class, int(race.abandoned),
             race.start_time, race.timezone, race.number_of_runners),
        )
        self.conn.commit()
        cur.execute(
            "SELECT id FROM races WHERE date=? AND track_slug=? AND race_number=?",
            (race.date, race.track_slug, race.race_number),
        )
        rid = cur.fetchone()["id"]
        # refresh runners (delete+insert for simplicity)
        cur.execute("DELETE FROM runners WHERE race_id=?", (rid,))
        for r in race.runners:
            cur.execute(
                """INSERT INTO runners (race_id, number, name, jockey, trainer, barrier, weight, form, data_completeness)
                   VALUES (?,?,?,?,?,?,?,?,?)""",
                (rid, r.number, r.name, r.jockey, r.trainer, r.barrier, r.weight,
                 r.form, r.data_completeness),
            )
        self.conn.commit()
        return rid

    def save_prediction(self, record: dict) -> None:
        self.conn.execute(
            """INSERT INTO predictions
               (date, track, track_slug, race_number, model_version, generated_at,
                top3, trifecta, alternatives, confidence, confidence_score,
                ranked_runners, data_completeness)
               VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?)
               ON CONFLICT(date, track_slug, race_number, model_version)
               DO UPDATE SET generated_at=excluded.generated_at, top3=excluded.top3,
                 trifecta=excluded.trifecta, alternatives=excluded.alternatives,
                 confidence=excluded.confidence, confidence_score=excluded.confidence_score,
                 ranked_runners=excluded.ranked_runners, data_completeness=excluded.data_completeness""",
            (record["date"], record["track"], record.get("track_slug", ""),
             record["race_number"], record.get("model_version", settings.model_version),
             record.get("generated_at"), json.dumps(record.get("top3")),
             json.dumps(record.get("trifecta")), json.dumps(record.get("alternatives")),
             record.get("confidence"), record.get("confidence_score"),
             json.dumps(record.get("ranked_runners")), record.get("data_completeness")),
        )
        self.conn.commit()

    def save_result(self, date: str, track: str, track_slug: str,
                   race_number: int, actual_trifecta: Optional[list]) -> None:
        exact = box = top1 = None
        if actual_trifecta and len(actual_trifecta) >= 3 and all(isinstance(x, int) for x in actual_trifecta[:3]):
            top1 = int(actual_trifecta[0])
            # exact/box require joining to predictions; computed in evaluation
        self.conn.execute(
            """INSERT INTO results (date, track, track_slug, race_number, actual_trifecta, exact_hit, box_hit, top1_hit)
               VALUES (?,?,?,?,?,?,?,?)
               ON CONFLICT(date, track_slug, race_number)
               DO UPDATE SET actual_trifecta=excluded.actual_trifecta""",
            (date, track, track_slug, race_number, json.dumps(actual_trifecta),
             exact, box, top1),
        )
        self.conn.commit()

    def get_predictions(self, date: str, model_version: Optional[str] = None) -> list[dict]:
        if model_version:
            rows = self.conn.execute(
                "SELECT * FROM predictions WHERE date=? AND model_version=?", (date, model_version)).fetchall()
        else:
            rows = self.conn.execute("SELECT * FROM predictions WHERE date=?", (date,)).fetchall()
        out = []
        for r in rows:
            d = dict(r)
            for k in ("top3", "trifecta", "alternatives", "ranked_runners"):
                if d.get(k):
                    d[k] = json.loads(d[k])
            out.append(d)
        return out

    def __enter__(self):
        return self

    def __exit__(self, *exc):
        self.close()