db/src/sleep_db/seed_sqlite.py

from __future__ import annotations

import argparse
import csv
import sqlite3
from pathlib import Path
from typing import Any

from sleep_db.config import load_settings


SCHEMA = """
CREATE TABLE IF NOT EXISTS sleep_metrics_daily (
    user_id TEXT NOT NULL,
    date TEXT NOT NULL,
    total_sleep_min INTEGER NOT NULL,
    deep_sleep_min INTEGER NOT NULL,
    light_sleep_min INTEGER NOT NULL,
    rem_sleep_min INTEGER NOT NULL,
    awake_count INTEGER NOT NULL,
    sleep_onset_latency_min INTEGER NOT NULL,
    heart_rate_avg REAL NOT NULL,
    heart_rate_baseline REAL NOT NULL,
    hrv_rmssd REAL NOT NULL,
    hrv_baseline REAL NOT NULL,
    respiratory_rate_avg REAL NOT NULL,
    spo2_min REAL NOT NULL,
    movement_count INTEGER NOT NULL,
    caffeine_after_6pm INTEGER NOT NULL,
    caffeine_after_2pm INTEGER NOT NULL DEFAULT 0,
    caffeine_mg_after_2pm REAL NOT NULL DEFAULT 0,
    alcohol_units_4h REAL NOT NULL DEFAULT 0,
    last_meal_hours_before_bed REAL NOT NULL DEFAULT 4,
    late_fluid_ml REAL NOT NULL DEFAULT 0,
    exercise_end_hour REAL NOT NULL DEFAULT 17,
    exercise_intensity REAL NOT NULL DEFAULT 0,
    training_load_score REAL NOT NULL DEFAULT 0,
    muscle_soreness_score REAL NOT NULL DEFAULT 0,
    daytime_fatigue_score REAL NOT NULL DEFAULT 0,
    stress_level REAL NOT NULL DEFAULT 0,
    breathing_rate_delta REAL NOT NULL DEFAULT 0,
    bedroom_temperature_c REAL NOT NULL DEFAULT 19,
    evening_light_lux REAL NOT NULL DEFAULT 50,
    screen_cutoff_min REAL NOT NULL DEFAULT 60,
    reaction_time_ms REAL NOT NULL DEFAULT 250,
    pvt_lapses INTEGER NOT NULL DEFAULT 0,
    mood_stress_score REAL NOT NULL DEFAULT 0,
    PRIMARY KEY (user_id, date)
);
"""


NUMERIC_COLUMNS = {
    "total_sleep_min": int,
    "deep_sleep_min": int,
    "light_sleep_min": int,
    "rem_sleep_min": int,
    "awake_count": int,
    "sleep_onset_latency_min": int,
    "heart_rate_avg": float,
    "heart_rate_baseline": float,
    "hrv_rmssd": float,
    "hrv_baseline": float,
    "respiratory_rate_avg": float,
    "spo2_min": float,
    "movement_count": int,
    "caffeine_mg_after_2pm": float,
    "alcohol_units_4h": float,
    "last_meal_hours_before_bed": float,
    "late_fluid_ml": float,
    "exercise_end_hour": float,
    "exercise_intensity": float,
    "training_load_score": float,
    "muscle_soreness_score": float,
    "daytime_fatigue_score": float,
    "stress_level": float,
    "breathing_rate_delta": float,
    "bedroom_temperature_c": float,
    "evening_light_lux": float,
    "screen_cutoff_min": float,
    "reaction_time_ms": float,
    "pvt_lapses": int,
    "mood_stress_score": float,
}

BOOLEAN_COLUMNS = {"caffeine_after_6pm", "caffeine_after_2pm"}


def seed_sqlite(csv_path: Path, sqlite_path: Path) -> int:
    sqlite_path.parent.mkdir(parents=True, exist_ok=True)
    rows = _read_rows(csv_path)
    with sqlite3.connect(sqlite_path) as conn:
        conn.execute(SCHEMA)
        _ensure_extra_columns(conn)
        conn.executemany(
            """
            INSERT OR REPLACE INTO sleep_metrics_daily (
                user_id, date, total_sleep_min, deep_sleep_min, light_sleep_min,
                rem_sleep_min, awake_count, sleep_onset_latency_min,
                heart_rate_avg, heart_rate_baseline, hrv_rmssd, hrv_baseline,
                respiratory_rate_avg, spo2_min, movement_count, caffeine_after_6pm,
                caffeine_after_2pm, caffeine_mg_after_2pm, alcohol_units_4h,
                last_meal_hours_before_bed, late_fluid_ml, exercise_end_hour,
                exercise_intensity, training_load_score, muscle_soreness_score,
                daytime_fatigue_score, stress_level, breathing_rate_delta,
                bedroom_temperature_c, evening_light_lux, screen_cutoff_min,
                reaction_time_ms, pvt_lapses, mood_stress_score
            ) VALUES (
                :user_id, :date, :total_sleep_min, :deep_sleep_min, :light_sleep_min,
                :rem_sleep_min, :awake_count, :sleep_onset_latency_min,
                :heart_rate_avg, :heart_rate_baseline, :hrv_rmssd, :hrv_baseline,
                :respiratory_rate_avg, :spo2_min, :movement_count, :caffeine_after_6pm,
                :caffeine_after_2pm, :caffeine_mg_after_2pm, :alcohol_units_4h,
                :last_meal_hours_before_bed, :late_fluid_ml, :exercise_end_hour,
                :exercise_intensity, :training_load_score, :muscle_soreness_score,
                :daytime_fatigue_score, :stress_level, :breathing_rate_delta,
                :bedroom_temperature_c, :evening_light_lux, :screen_cutoff_min,
                :reaction_time_ms, :pvt_lapses, :mood_stress_score
            )
            """,
            rows,
        )
    return len(rows)


def get_sleep_metrics(user_id: str, date: str, sqlite_path: Path | None = None) -> dict[str, Any]:
    settings = load_settings()
    db_path = sqlite_path or settings.sqlite_path
    with sqlite3.connect(db_path) as conn:
        conn.row_factory = sqlite3.Row
        row = conn.execute(
            "SELECT * FROM sleep_metrics_daily WHERE user_id = ? AND date = ?",
            (user_id, date),
        ).fetchone()
    if row is None:
        raise LookupError(f"No sleep metrics for user_id={user_id!r} date={date!r}")
    result = dict(row)
    for column in BOOLEAN_COLUMNS:
        result[column] = bool(result[column])
    return result


def _read_rows(csv_path: Path) -> list[dict[str, Any]]:
    with csv_path.open("r", encoding="utf-8", newline="") as handle:
        rows = []
        for row in csv.DictReader(handle):
            parsed = dict(row)
            for column, caster in NUMERIC_COLUMNS.items():
                parsed[column] = caster(parsed[column])
            for column in BOOLEAN_COLUMNS:
                parsed[column] = 1 if parsed[column].lower() == "true" else 0
            rows.append(parsed)
    return rows


def _ensure_extra_columns(conn: sqlite3.Connection) -> None:
    existing = {row[1] for row in conn.execute("PRAGMA table_info(sleep_metrics_daily)").fetchall()}
    column_defs = {
        "caffeine_after_2pm": "INTEGER NOT NULL DEFAULT 0",
        "caffeine_mg_after_2pm": "REAL NOT NULL DEFAULT 0",
        "alcohol_units_4h": "REAL NOT NULL DEFAULT 0",
        "last_meal_hours_before_bed": "REAL NOT NULL DEFAULT 4",
        "late_fluid_ml": "REAL NOT NULL DEFAULT 0",
        "exercise_end_hour": "REAL NOT NULL DEFAULT 17",
        "exercise_intensity": "REAL NOT NULL DEFAULT 0",
        "training_load_score": "REAL NOT NULL DEFAULT 0",
        "muscle_soreness_score": "REAL NOT NULL DEFAULT 0",
        "daytime_fatigue_score": "REAL NOT NULL DEFAULT 0",
        "stress_level": "REAL NOT NULL DEFAULT 0",
        "breathing_rate_delta": "REAL NOT NULL DEFAULT 0",
        "bedroom_temperature_c": "REAL NOT NULL DEFAULT 19",
        "evening_light_lux": "REAL NOT NULL DEFAULT 50",
        "screen_cutoff_min": "REAL NOT NULL DEFAULT 60",
        "reaction_time_ms": "REAL NOT NULL DEFAULT 250",
        "pvt_lapses": "INTEGER NOT NULL DEFAULT 0",
        "mood_stress_score": "REAL NOT NULL DEFAULT 0",
    }
    for column, definition in column_defs.items():
        if column not in existing:
            conn.execute(f"ALTER TABLE sleep_metrics_daily ADD COLUMN {column} {definition}")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--csv", type=Path)
    parser.add_argument("--sqlite", type=Path)
    args = parser.parse_args()

    settings = load_settings()
    csv_path = args.csv or settings.seed_metrics_csv
    sqlite_path = args.sqlite or settings.sqlite_path
    count = seed_sqlite(csv_path, sqlite_path)
    print(f"Seeded {count} rows into {sqlite_path}")


if __name__ == "__main__":
    main()