Source code for trader.persistence.snapshots

"""Repository for account/position snapshots (requirements §10).

The only repository wired in Slice 1. Later slices add repositories for
trades, decisions, watchlist events, and backtest runs against the same
schema.
"""

from collections.abc import Sequence
from datetime import UTC, datetime

from sqlalchemy import select
from sqlalchemy.orm import Session, selectinload, sessionmaker

from trader.domain import Account, Position
from trader.persistence.models import AccountSnapshot, PositionSnapshot

__all__ = ["SnapshotRepository"]




class SnapshotRepository:
    """Reads and writes portfolio snapshots."""

    def __init__(self, session_factory: sessionmaker[Session]) -> None:
        self._session_factory = session_factory



    def save_snapshot(
        self,
        account: Account,
        positions: Sequence[Position],
        captured_at: datetime | None = None,
    ) -> int:
        """Persist an account snapshot and its positions. Returns the new id."""
        snapshot = AccountSnapshot(
            captured_at=captured_at or datetime.now(UTC),
            account_id=account.account_id,
            cash=account.cash,
            equity=account.equity,
            buying_power=account.buying_power,
            portfolio_value=account.portfolio_value,
            is_paper=account.is_paper,
            positions=[
                PositionSnapshot(
                    ticker=position.symbol,
                    quantity=position.quantity,
                    avg_entry_price=position.avg_entry_price,
                    current_price=position.current_price,
                    market_value=position.market_value,
                    unrealized_pl=position.unrealized_pl,
                )
                for position in positions
            ],
        )
        with self._session_factory() as session:
            session.add(snapshot)
            session.commit()
            return snapshot.id




    def latest_snapshot(self) -> AccountSnapshot | None:
        """The most recently captured snapshot, with positions loaded."""
        statement = (
            select(AccountSnapshot)
            .options(selectinload(AccountSnapshot.positions))
            .order_by(AccountSnapshot.captured_at.desc(), AccountSnapshot.id.desc())
            .limit(1)
        )
        with self._session_factory() as session:
            return session.scalars(statement).first()




    def snapshots_in_range(
        self, start: datetime | None = None, end: datetime | None = None
    ) -> list[AccountSnapshot]:
        """Snapshots captured within `[start, end]` (both bounds inclusive),
        ordered oldest to newest — feeds `trader.performance.benchmark.
        account_return` (issue #19), which compares the FIRST snapshot in the
        list against the LAST.

        Both bounds inclusive, unlike the news archive's half-open window:
        that window exists to stop a decision seeing a headline it is being
        scored for anticipating, which has no analogue here — there is no
        lookahead risk in including a snapshot captured exactly at `end`.

        `None` for either bound means "no floor"/"no ceiling" — omitting both
        returns every snapshot ever captured.

        Positions are deliberately NOT eager-loaded here (contrast
        `latest_snapshot`): every caller of this method reads only
        `captured_at`/`portfolio_value`, so `selectinload` would cost an extra
        query per snapshot for data nothing uses.
        """
        statement = select(AccountSnapshot).order_by(
            AccountSnapshot.captured_at.asc(), AccountSnapshot.id.asc()
        )
        if start is not None:
            statement = statement.where(AccountSnapshot.captured_at >= start)
        if end is not None:
            statement = statement.where(AccountSnapshot.captured_at <= end)
        with self._session_factory() as session:
            return list(session.scalars(statement))




    def snapshots_since(self, start: datetime | None = None) -> Sequence[AccountSnapshot]:
        """Snapshots in chronological order, oldest first (§12.3 equity curve).

        `id` is the tiebreak for equal `captured_at` values, ascending to match
        `latest_snapshot`'s descending tiebreak — both resolve a tie in favour
        of the row inserted later. Positions are deliberately not eager-loaded
        here: callers of this method (the equity curve) only read scalar
        account-level columns, so pulling `position_snapshots` along for every
        row would cost a join no caller needs.
        """
        statement = select(AccountSnapshot).order_by(
            AccountSnapshot.captured_at.asc(), AccountSnapshot.id.asc()
        )
        if start is not None:
            statement = statement.where(AccountSnapshot.captured_at >= start)
        with self._session_factory() as session:
            return session.scalars(statement).all()