aaredb/backend/app/routers/puck.py

from fastapi import APIRouter, HTTPException, status, Depends
from sqlalchemy.orm import Session
from typing import List
import uuid
from app.schemas import Puck as PuckSchema, PuckCreate, PuckUpdate, SetTellPosition, PuckEvent
from app.models import Puck as PuckModel, Sample as SampleModel, PuckEvent as PuckEventModel, Slot as SlotModel, LogisticsEvent as LogisticsEventModel, Dewar as DewarModel
from app.dependencies import get_db
from datetime import datetime

router = APIRouter()


@router.get("/", response_model=List[PuckSchema])
async def get_pucks(db: Session = Depends(get_db)):
    return db.query(PuckModel).all()


@router.get("/{puck_id}", response_model=PuckSchema)
async def get_puck(puck_id: str, db: Session = Depends(get_db)):
    puck = db.query(PuckModel).filter(PuckModel.id == puck_id).first()
    if not puck:
        raise HTTPException(status_code=404, detail="Puck not found")
    return puck


@router.post("/", response_model=PuckSchema, status_code=status.HTTP_201_CREATED)
async def create_puck(puck: PuckCreate, db: Session = Depends(get_db)) -> PuckSchema:
    puck_id = f'PUCK-{uuid.uuid4().hex[:8].upper()}'
    db_puck = PuckModel(
        id=puck_id,
        puck_name=puck.puck_name,
        puck_type=puck.puck_type,
        puck_location_in_dewar=puck.puck_location_in_dewar,
        dewar_id=puck.dewar_id
    )
    db.add(db_puck)
    db.commit()
    db.refresh(db_puck)
    return db_puck


@router.put("/{puck_id}", response_model=PuckSchema)
async def update_puck(puck_id: str, updated_puck: PuckUpdate, db: Session = Depends(get_db)):
    puck = db.query(PuckModel).filter(PuckModel.id == puck_id).first()
    if not puck:
        raise HTTPException(status_code=404, detail="Puck not found")

    for key, value in updated_puck.dict(exclude_unset=True).items():
        setattr(puck, key, value)

    db.commit()
    db.refresh(puck)
    return puck


@router.delete("/{puck_id}", status_code=status.HTTP_204_NO_CONTENT)
async def delete_puck(puck_id: str, db: Session = Depends(get_db)):
    puck = db.query(PuckModel).filter(PuckModel.id == puck_id).first()
    if not puck:
        raise HTTPException(status_code=404, detail="Puck not found")

    db.delete(puck)
    db.commit()
    return

@router.put("/{puck_id}/tell_position", status_code=status.HTTP_200_OK)
async def set_tell_position(
        puck_id: int,
        request: SetTellPosition,
        db: Session = Depends(get_db)
):
    # Get the requested tell_position
    tell_position = request.tell_position

    # Define valid positions
    valid_positions = [f"{letter}{num}" for letter in "ABCDEF" for num in range(1, 6)] + ["null", None]

    # Validate tell_position
    if tell_position not in valid_positions:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=f"Invalid tell_position value. Must be one of {valid_positions}.",
        )

    # Set the correct tell_position logic
    actual_position = None if tell_position in ["null", None] else tell_position

    # Create a new PuckEvent (always a new event, even with null/None)
    new_puck_event = PuckEventModel(
        puck_id=puck_id,
        tell_position=actual_position,  # Null for disassociation, else the valid position
        event_type="tell_position_set",  # Example event type
        timestamp=datetime.utcnow(),
    )
    db.add(new_puck_event)
    db.commit()
    db.refresh(new_puck_event)

    # Send the response
    return {
        "message": "New tell position event created successfully",
        "tell_position": new_puck_event.tell_position,
        "timestamp": new_puck_event.timestamp,
    }


@router.get("/{puck_id}/last-tell-position", status_code=status.HTTP_200_OK)
async def get_last_tell_position(puck_id: str, db: Session = Depends(get_db)):
    # Query the most recent tell_position_set event for the given puck_id
    last_event = (
        db.query(PuckEventModel)
        .filter(PuckEventModel.puck_id == puck_id, PuckEventModel.event_type == "tell_position_set")
        .order_by(PuckEventModel.timestamp.desc())
        .first()
    )

    # If no event is found, return a 404 error
    if not last_event:
        raise HTTPException(
            status_code=404,
            detail=f"No 'tell_position' event found for puck with ID {puck_id}",
        )

    # Return the details of the last tell_position event
    return {
        "puck_id": puck_id,
        "tell_position": last_event.tell_position,
        "timestamp": last_event.timestamp,
    }

@router.get("/slot/{slot_id}", response_model=List[PuckSchema])
async def get_pucks_by_latest_beamline_dewar_slot(slot_id: int, db: Session = Depends(get_db)):
    """
    Retrieve all pucks for the most recent dewar associated with the given slot_id
    with a 'beamline' event type.
    """
    # Query the Slot table to ensure the slot exists
    slot = db.query(SlotModel).filter(SlotModel.id == slot_id).first()
    if not slot:
        raise HTTPException(status_code=404, detail="Slot not found")

    # Check if this slot has an associated dewar
    if not slot.dewar:
        raise HTTPException(status_code=404, detail="No dewar associated with the given slot")

    # Find the most recent dewar associated with a "beamline" event type
    recent_beamline_event = (
        db.query(LogisticsEventModel)
        .filter(
            LogisticsEventModel.slot_id == slot_id,
            LogisticsEventModel.event_type == "Beamline"  # Filter by "beamline"
        )
        .order_by(LogisticsEventModel.timestamp.desc())  # Order by most recent event
        .first()
    )

    # Confirm if we found a "beamline" event for the slot
    if not recent_beamline_event:
        raise HTTPException(
            status_code=404,
            detail="No 'beamline' event found for the given slot"
        )

    # Retrieve the dewar linked to the "beamline" event
    dewar = db.query(DewarModel).filter(DewarModel.id == recent_beamline_event.dewar_id).first()
    if not dewar:
        raise HTTPException(
            status_code=404,
            detail="No dewar associated with the most recent 'beamline' event"
        )

    # Retrieve all pucks associated with the dewar
    pucks = db.query(PuckModel).filter(PuckModel.dewar_id == dewar.id).all()

    # Return the list of pucks
    return pucks