sqlite_storage.py

from __future__ import annotations

import sqlite3
from pathlib import Path
from typing import Dict

import pandas as pd

from data_loader import DATA_DIR, prepare_frames


SCHEMA_PATH = DATA_DIR / "schema.sql"
DEFAULT_DB_PATH = DATA_DIR / "market_data.db"


def ensure_database(
    tickers_df: pd.DataFrame | None = None,
    prices_df: pd.DataFrame | None = None,
    *,
    rebuild: bool = False,
    db_path: Path | None = None,
) -> sqlite3.Connection:
    """
    Guarantee that the SQLite database exists and contains the current snapshot.
    """
    db_path = Path(db_path or DEFAULT_DB_PATH)
    if rebuild and db_path.exists():
        db_path.unlink()

    needs_setup = not db_path.exists()
    conn = sqlite3.connect(db_path)

    if not needs_setup:
        return conn

    if tickers_df is None or prices_df is None:
        tickers_df, prices_df = prepare_frames()

    with open(SCHEMA_PATH, "r", encoding="utf-8") as fh:
        conn.executescript(fh.read())

    tickers_df.to_sql("tickers", conn, if_exists="append", index=False)

    price_records = prices_df.copy()
    if "ticker" in price_records.columns:
        price_records = price_records.drop(columns=["ticker"])
    price_records["timestamp"] = price_records["timestamp"].dt.strftime(
        "%Y-%m-%d %H:%M:%S"
    )
    price_records.to_sql("prices", conn, if_exists="append", index=False)
    return conn


def query_tsla_range(
    conn: sqlite3.Connection,
    *,
    start_ts: str,
    end_ts: str,
) -> pd.DataFrame:
    sql = """
        SELECT
            p.timestamp,
            t.symbol,
            p.open,
            p.high,
            p.low,
            p.close,
            p.volume
        FROM prices p
        JOIN tickers t ON t.ticker_id = p.ticker_id
        WHERE
            t.symbol = ?
            AND p.timestamp BETWEEN ? AND ?
        ORDER BY p.timestamp;
    """
    return pd.read_sql_query(sql, conn, params=("TSLA", start_ts, end_ts))


def query_average_daily_volume(conn: sqlite3.Connection) -> pd.DataFrame:
    sql = """
        SELECT
            t.symbol,
            DATE(p.timestamp) AS trade_date,
            AVG(p.volume) AS avg_daily_volume
        FROM prices p
        JOIN tickers t ON t.ticker_id = p.ticker_id
        GROUP BY t.symbol, DATE(p.timestamp)
        ORDER BY trade_date, t.symbol;
    """
    return pd.read_sql_query(sql, conn)


def query_top_returns(conn: sqlite3.Connection, *, limit: int = 3) -> pd.DataFrame:
    sql = f"""
        WITH first_trade AS (
            SELECT
                p.ticker_id,
                p.open AS first_open
            FROM prices p
            JOIN (
                SELECT ticker_id, MIN(timestamp) AS first_ts
                FROM prices
                GROUP BY ticker_id
            ) bounds
                ON bounds.ticker_id = p.ticker_id
                AND bounds.first_ts = p.timestamp
        ),
        last_trade AS (
            SELECT
                p.ticker_id,
                p.close AS last_close
            FROM prices p
            JOIN (
                SELECT ticker_id, MAX(timestamp) AS last_ts
                FROM prices
                GROUP BY ticker_id
            ) bounds
                ON bounds.ticker_id = p.ticker_id
                AND bounds.last_ts = p.timestamp
        )
        SELECT
            t.symbol,
            ROUND((last_trade.last_close - first_trade.first_open) / first_trade.first_open, 6)
                AS total_return
        FROM first_trade
        JOIN last_trade USING (ticker_id)
        JOIN tickers t ON t.ticker_id = first_trade.ticker_id
        ORDER BY total_return DESC
        LIMIT {limit};
    """
    return pd.read_sql_query(sql, conn)


def query_first_last_close(conn: sqlite3.Connection) -> pd.DataFrame:
    sql = """
        WITH daily_bounds AS (
            SELECT
                ticker_id,
                DATE(timestamp) AS trade_date,
                MIN(timestamp) AS first_ts,
                MAX(timestamp) AS last_ts
            FROM prices
            GROUP BY ticker_id, DATE(timestamp)
        )
        SELECT
            t.symbol,
            d.trade_date,
            first_trade.close AS first_close,
            last_trade.close AS last_close
        FROM daily_bounds d
        JOIN prices first_trade
            ON first_trade.ticker_id = d.ticker_id
            AND first_trade.timestamp = d.first_ts
        JOIN prices last_trade
            ON last_trade.ticker_id = d.ticker_id
            AND last_trade.timestamp = d.last_ts
        JOIN tickers t ON t.ticker_id = d.ticker_id
        ORDER BY d.trade_date, t.symbol;
    """
    return pd.read_sql_query(sql, conn)


def describe_queries(conn: sqlite3.Connection) -> Dict[str, pd.DataFrame]:
    return {
        "sqlite_tsla_range": query_tsla_range(
            conn,
            start_ts="2025-11-17 00:00:00",
            end_ts="2025-11-18 23:59:59",
        ),
        "sqlite_avg_daily_volume": query_average_daily_volume(conn),
        "sqlite_top_returns": query_top_returns(conn),
        "sqlite_daily_first_last": query_first_last_close(conn),
    }