from fastapi import APIRouter, Depends, HTTPException, Query from sqlalchemy.orm import Session from sqlalchemy import and_ from sqlalchemy.exc import IntegrityError from datetime import date, timedelta from typing import Optional from ..database import get_db from ..dependencies import get_current_user from ..models import DailyShift, Staff, PumpReading from ..schemas import RosterEntry, RosterOut router = APIRouter() def _shift_to_out(shift: DailyShift) -> dict: # Collect unique pumper names from pump_readings (preserving assignment order) seen: set = set() pumper_names: list = [] for r in (shift.pump_readings or []): if r.staff_id and r.staff_id not in seen and r.pumper: seen.add(r.staff_id) pumper_names.append(r.pumper.full_name) return { "id": shift.id, "roster_date": shift.record_date, "shift_type": shift.shift_type, "staff_id": shift.staff_id, "staff_name": shift.staff_member.full_name if shift.staff_member else None, "pumper_names": pumper_names, "notes": shift.notes, } @router.get("/roster", response_model=list[RosterOut]) def get_roster_month( month: int = Query(..., ge=1, le=12), year: int = Query(..., ge=2020), db: Session = Depends(get_db), _=Depends(get_current_user), ): """Return all shift assignments for a given month/year.""" from calendar import monthrange last_day = monthrange(year, month)[1] start = date(year, month, 1) end = date(year, month, last_day) rows = ( db.query(DailyShift) .filter(and_(DailyShift.record_date >= start, DailyShift.record_date <= end)) .order_by(DailyShift.record_date, DailyShift.shift_type) .all() ) return [_shift_to_out(r) for r in rows] @router.get("/roster/week", response_model=list[RosterOut]) def get_roster_week( date_str: Optional[str] = Query(None, alias="date"), db: Session = Depends(get_db), _=Depends(get_current_user), ): """Return 7-day view starting from given date (or today).""" start = date.fromisoformat(date_str) if date_str else date.today() end = start + timedelta(days=6) rows = ( db.query(DailyShift) .filter(and_(DailyShift.record_date >= start, DailyShift.record_date <= end)) .order_by(DailyShift.record_date, DailyShift.shift_type) .all() ) return [_shift_to_out(r) for r in rows] @router.post("/roster", response_model=RosterOut, status_code=201) def assign_roster( body: RosterEntry, db: Session = Depends(get_db), current_user=Depends(get_current_user), ): """Assign a staff member to a shift slot (creates a DailyShift record).""" staff = db.query(Staff).filter(Staff.id == body.staff_id).first() if not staff: raise HTTPException(status_code=404, detail="Staff not found") existing = ( db.query(DailyShift) .filter( DailyShift.record_date == body.roster_date, DailyShift.shift_type == body.shift_type, ) .first() ) from .audit import log_audit if existing: # Update staff assignment on existing shift existing.staff_id = body.staff_id if body.notes is not None: existing.notes = body.notes log_audit(db, current_user.id, 'UPDATE', 'daily_shifts', existing.id, f"Roster: assigned manager to {body.shift_type} on {body.roster_date}") db.commit() db.refresh(existing) return _shift_to_out(existing) try: from ..models import ShiftTemplate tmpl = db.query(ShiftTemplate).filter(ShiftTemplate.code == body.shift_type.upper()).first() shift = DailyShift( record_date=body.roster_date, shift_type=body.shift_type, shift_template_id=tmpl.id if tmpl else None, shift_start_time=tmpl.start_time if tmpl else None, shift_end_time=tmpl.end_time if tmpl else None, staff_id=body.staff_id, notes=body.notes, ) db.add(shift) db.flush() log_audit(db, current_user.id, 'CREATE', 'daily_shifts', shift.id, f"Roster: created {body.shift_type} shift on {body.roster_date}") db.commit() db.refresh(shift) return _shift_to_out(shift) except IntegrityError: # Duplicate key — another request created the shift first; update it instead db.rollback() existing = ( db.query(DailyShift) .filter( DailyShift.record_date == body.roster_date, DailyShift.shift_type == body.shift_type, ) .first() ) if not existing: raise HTTPException(status_code=500, detail="Failed to create or find shift") existing.staff_id = body.staff_id if body.notes is not None: existing.notes = body.notes log_audit(db, current_user.id, 'UPDATE', 'daily_shifts', existing.id, f"Roster: assigned manager to {body.shift_type} on {body.roster_date}") db.commit() db.refresh(existing) return _shift_to_out(existing) @router.put("/roster/{shift_id}", response_model=RosterOut) def update_roster( shift_id: int, body: RosterEntry, db: Session = Depends(get_db), current_user=Depends(get_current_user), ): shift = db.query(DailyShift).filter(DailyShift.id == shift_id).first() if not shift: raise HTTPException(status_code=404, detail="Roster entry not found") staff = db.query(Staff).filter(Staff.id == body.staff_id).first() if not staff: raise HTTPException(status_code=404, detail="Staff not found") shift.staff_id = body.staff_id if body.notes is not None: shift.notes = body.notes from .audit import log_audit log_audit(db, current_user.id, 'UPDATE', 'daily_shifts', shift_id, f"Roster: updated manager for {shift.shift_type} on {shift.record_date}") db.commit() db.refresh(shift) return _shift_to_out(shift) @router.delete("/roster/{shift_id}", status_code=204) def remove_roster_staff( shift_id: int, db: Session = Depends(get_db), current_user=Depends(get_current_user), ): """Remove the staff assignment from a shift slot (nullifies staff_id).""" shift = db.query(DailyShift).filter(DailyShift.id == shift_id).first() if not shift: raise HTTPException(status_code=404, detail="Roster entry not found") if shift.is_locked: raise HTTPException(status_code=400, detail="Cannot modify a locked shift") shift.staff_id = None from .audit import log_audit log_audit(db, current_user.id, 'UPDATE', 'daily_shifts', shift_id, f"Roster: removed manager from {shift.shift_type} on {shift.record_date}") db.commit()