benchpress/backend/app/services/schedule.py

from __future__ import annotations

from datetime import date, timedelta

from sqlalchemy.orm import Session

from app.models import AccessoryState, Workout, WorkoutExercise, WorkoutSet
from app.services.program import load_program
from app.services.weights import bench_weight


WEEKDAY_NAMES = ["monday", "tuesday", "wednesday", "thursday", "friday", "saturday", "sunday"]


def next_dates_for_weekdays(start: date, weekdays: list[int], count: int) -> list[date]:
    """Return `count` dates on selected weekdays starting from start (inclusive if matches)."""
    selected = set(weekdays)
    dates: list[date] = []
    cursor = start
    # safety limit
    for _ in range(400):
        if cursor.weekday() in selected:
            dates.append(cursor)
            if len(dates) >= count:
                break
        cursor += timedelta(days=1)
    if len(dates) < count:
        raise ValueError("Не удалось построить расписание")
    return dates


def clear_cycle_workouts(db: Session, cycle_number: int) -> None:
    workouts = db.query(Workout).filter(Workout.cycle_number == cycle_number).all()
    for w in workouts:
        db.delete(w)
    db.query(AccessoryState).filter(AccessoryState.cycle_number == cycle_number).delete()
    db.flush()


def generate_schedule(
    db: Session,
    *,
    one_rm_kg: float,
    training_weekdays: list[int],
    start_date: date,
    cycle_number: int,
) -> list[Workout]:
    program = load_program()
    total_days = sum(len(w["days"]) for w in program["weeks"])
    dates = next_dates_for_weekdays(start_date, training_weekdays, total_days)

    # Map chronological day_index within week order from template (day 0,1,2 of each week)
    date_idx = 0
    created: list[Workout] = []

    for week in program["weeks"]:
        for day_index, day in enumerate(week["days"]):
            scheduled = dates[date_idx]
            date_idx += 1
            workout = Workout(
                cycle_number=cycle_number,
                week_number=week["week"],
                day_index=day_index,
                scheduled_date=scheduled,
                status="planned",
            )
            db.add(workout)
            db.flush()

            for order_index, ex in enumerate(day["exercises"]):
                kind = ex["kind"]
                accessory_state = None
                needs_weight = False
                prescribed = None
                locked = False

                if kind == "accessory":
                    accessory_state = (
                        db.query(AccessoryState)
                        .filter(
                            AccessoryState.cycle_number == cycle_number,
                            AccessoryState.exercise_name == ex["name"],
                        )
                        .first()
                    )
                    if accessory_state and accessory_state.locked and accessory_state.weight_kg is not None:
                        prescribed = accessory_state.weight_kg
                        locked = True
                    elif accessory_state and accessory_state.pending_weight_kg is not None:
                        prescribed = accessory_state.pending_weight_kg
                        needs_weight = False
                    elif accessory_state and accessory_state.weight_kg is not None and not accessory_state.locked:
                        prescribed = accessory_state.weight_kg
                        needs_weight = False
                    else:
                        needs_weight = True

                we = WorkoutExercise(
                    workout_id=workout.id,
                    order_index=order_index,
                    name=ex["name"],
                    kind=kind,
                    prescribed_weight_kg=prescribed,
                    weight_locked=locked,
                    needs_weight_input=needs_weight,
                )
                db.add(we)
                db.flush()

                for set_index, s in enumerate(ex["sets"]):
                    if kind == "bench":
                        tw = bench_weight(one_rm_kg, s["percent"])
                        ws = WorkoutSet(
                            exercise_id=we.id,
                            set_index=set_index,
                            target_percent=s["percent"],
                            target_reps=s["reps"],
                            target_weight_kg=tw,
                            is_failure=False,
                        )
                    elif kind == "failure":
                        ws = WorkoutSet(
                            exercise_id=we.id,
                            set_index=set_index,
                            target_reps=None,
                            target_weight_kg=None,
                            is_failure=True,
                        )
                    else:
                        reps = s.get("reps") if isinstance(s, dict) else None
                        ws = WorkoutSet(
                            exercise_id=we.id,
                            set_index=set_index,
                            target_reps=reps,
                            target_weight_kg=prescribed,
                            is_failure=False,
                        )
                    db.add(ws)

            created.append(workout)

    db.flush()
    return created


def apply_accessory_weight_to_future(
    db: Session,
    *,
    cycle_number: int,
    exercise_name: str,
    weight_kg: float,
    from_date: date,
    lock: bool,
) -> None:
    state = (
        db.query(AccessoryState)
        .filter(
            AccessoryState.cycle_number == cycle_number,
            AccessoryState.exercise_name == exercise_name,
        )
        .first()
    )
    if not state:
        state = AccessoryState(cycle_number=cycle_number, exercise_name=exercise_name)
        db.add(state)

    if lock:
        state.weight_kg = weight_kg
        state.locked = True
        state.pending_weight_kg = None
    else:
        state.pending_weight_kg = weight_kg
        state.weight_kg = weight_kg
        state.locked = False

    # Update future workouts' exercises of same name
    future = (
        db.query(WorkoutExercise)
        .join(Workout)
        .filter(
            Workout.cycle_number == cycle_number,
            Workout.scheduled_date >= from_date,
            Workout.status != "completed",
            WorkoutExercise.name == exercise_name,
            WorkoutExercise.kind == "accessory",
        )
        .all()
    )
    for we in future:
        if we.sets and all(s.completed for s in we.sets):
            continue
        we.prescribed_weight_kg = weight_kg
        we.weight_locked = lock
        we.needs_weight_input = False
        for s in we.sets:
            if not s.completed:
                s.target_weight_kg = weight_kg