tracker/db.py

from __future__ import annotations

import json
from pathlib import Path
from typing import Iterator

from fastapi import Request
from sqlalchemy import Engine, inspect, text
from sqlmodel import Session, create_engine, select

from tracker.config import TrackerConfig, load_config
from tracker.models import TrackerUser

_engine_by_resolved_path: dict[str, Engine] = {}

USER_SCOPED_TABLES = (
    "programs",
    "workout_logs",
    "body_weight_logs",
    "run_logs",
    "planned_meals",
    "meal_log_entries",
)


def get_engine_for_config(cfg: TrackerConfig) -> Engine:
    """Return a cached engine for this database path (one engine per resolved SQLite file)."""
    resolved = str(cfg.tracker_db_full_path.resolve())
    existing = _engine_by_resolved_path.get(resolved)
    if existing is not None:
        return existing
    cfg.tracker_db_full_path.parent.mkdir(parents=True, exist_ok=True)
    eng = create_engine(f"sqlite:///{resolved}", echo=False)
    _engine_by_resolved_path[resolved] = eng
    return eng


def _default_user_email(cfg: TrackerConfig) -> str:
    if cfg.dashboard_logins:
        return cfg.dashboard_logins[0].email
    return "local@example.com"


def _ensure_default_user(engine: Engine, cfg: TrackerConfig) -> int:
    email = _default_user_email(cfg)
    login = next((item for item in cfg.dashboard_logins if item.email == email), None)
    with Session(engine) as session:
        user = session.exec(select(TrackerUser).where(TrackerUser.email == email)).first()
        if user is None:
            user = TrackerUser(
                email=email,
                label=login.label if login is not None else "Local",
                role=login.role if login is not None else None,
            )
            session.add(user)
            session.commit()
            session.refresh(user)
        elif login is not None:
            user.label = login.label
            user.role = login.role
            session.add(user)
            session.commit()
            session.refresh(user)
        if user.user_id is None:
            msg = f"Tracker user {email} has no user_id"
            raise RuntimeError(msg)
        return user.user_id


def _legacy_user_id(engine: Engine) -> int | None:
    with Session(engine) as session:
        user = session.exec(select(TrackerUser).where(TrackerUser.email == "local@example.com")).first()
        return user.user_id if user is not None else None


def _add_user_id_columns(
    engine: Engine,
    table_names: list[str],
    default_user_id: int,
) -> None:
    inspector = inspect(engine)
    legacy_user_id = _legacy_user_id(engine)
    with engine.begin() as connection:
        for table_name in USER_SCOPED_TABLES:
            if table_name not in table_names:
                continue
            column_names = {column["name"] for column in inspector.get_columns(table_name)}
            if "user_id" not in column_names:
                connection.execute(text(f"ALTER TABLE {table_name} ADD COLUMN user_id INTEGER"))
            connection.execute(
                text(f"UPDATE {table_name} SET user_id = :user_id WHERE user_id IS NULL"),
                {"user_id": default_user_id},
            )
            if legacy_user_id is not None and legacy_user_id != default_user_id:
                connection.execute(
                    text(f"UPDATE {table_name} SET user_id = :user_id WHERE user_id = :legacy_user_id"),
                    {"user_id": default_user_id, "legacy_user_id": legacy_user_id},
                )


def migrate_tracker_schema(engine: Engine, cfg: TrackerConfig) -> None:
    """Apply additive SQLite schema updates for existing tracker.db files."""
    inspector = inspect(engine)
    table_names = inspector.get_table_names()
    default_user_id = _ensure_default_user(engine, cfg)
    _add_user_id_columns(engine, table_names, default_user_id)
    if "program_exercises" in table_names:
        column_names = {column["name"] for column in inspector.get_columns("program_exercises")}
        if "starting_weight_kg" not in column_names:
            with engine.begin() as connection:
                connection.execute(
                    text("ALTER TABLE program_exercises ADD COLUMN starting_weight_kg REAL"),
                )
    if "exercises" in table_names:
        ex_cols = {column["name"] for column in inspector.get_columns("exercises")}
        with engine.begin() as connection:
            if "starting_bodyweight_fraction" not in ex_cols:
                connection.execute(
                    text("ALTER TABLE exercises ADD COLUMN starting_bodyweight_fraction REAL"),
                )
            if "starting_load_kg" not in ex_cols:
                connection.execute(
                    text("ALTER TABLE exercises ADD COLUMN starting_load_kg REAL"),
                )
            if "starting_load_reference" not in ex_cols:
                connection.execute(
                    text(
                        "ALTER TABLE exercises ADD COLUMN starting_load_reference TEXT DEFAULT ''",
                    ),
                )
    if "run_logs" in table_names:
        run_cols = {column["name"] for column in inspector.get_columns("run_logs")}
        with engine.begin() as connection:
            if "program_id" not in run_cols:
                connection.execute(text("ALTER TABLE run_logs ADD COLUMN program_id INTEGER"))
            if "day_label" not in run_cols:
                connection.execute(text("ALTER TABLE run_logs ADD COLUMN day_label TEXT"))


SEED_FILE = Path(__file__).resolve().parent.parent / "seed" / "exercises.json"


def get_session(request: Request) -> Iterator[Session]:
    with Session(request.app.state.db_engine) as session:
        yield session


def seed_exercises_if_empty(engine: Engine) -> None:
    """Bulk-insert exercises from seed/exercises.json on first run."""
    from tracker.models import Exercise

    with Session(engine) as session:
        count = session.exec(select(Exercise)).first()
        if count is not None:
            return

        if not SEED_FILE.exists():
            return

        raw = json.loads(SEED_FILE.read_text())
        for entry in raw:
            session.add(Exercise(
                name=entry["name"],
                muscle_group=entry.get("muscle_group", ""),
                equipment=entry.get("equipment", ""),
                wger_id=entry.get("wger_id"),
            ))
        session.commit()


def default_engine() -> Engine:
    """Engine for the process default config (CLI and MCP outside HTTP)."""
    return get_engine_for_config(load_config())