"""Checkpoint queries."""

from __future__ import annotations

from uuid import UUID

from sqlmodel import Session, select

from .tables import Checkpoint


def get_latest_checkpoint(session: Session, timeline_id: UUID, parent_id: UUID | None = None) -> Checkpoint | None:
    stmt = select(Checkpoint).where(Checkpoint.timeline_id == timeline_id)
    if parent_id is not None:
        stmt = stmt.where(Checkpoint.parent_id == parent_id)
    stmt = stmt.order_by(Checkpoint.created_at.desc())
    return session.exec(stmt).first()


def get_root_checkpoint(session: Session, timeline_id: UUID) -> Checkpoint | None:
    stmt = select(Checkpoint).where(
        Checkpoint.timeline_id == timeline_id,
        Checkpoint.parent_id == None,
    )
    return session.exec(stmt).first()


def list_checkpoints(session: Session, timeline_id: UUID) -> list[Checkpoint]:
    stmt = (
        select(Checkpoint)
        .where(Checkpoint.timeline_id == timeline_id)
        .order_by(Checkpoint.created_at)
    )
    return list(session.exec(stmt).all())


def list_scenarios(session: Session) -> list[Checkpoint]:
    stmt = (
        select(Checkpoint)
        .where(Checkpoint.is_scenario == True)
        .order_by(Checkpoint.created_at.desc())
    )
    return list(session.exec(stmt).all())