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
import logging

router = APIRouter()

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

@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,
    where the last logistics event is of type 'beamline'.
    """
    # Step 1: Verify the slot exists
    slot = db.query(SlotModel).filter(SlotModel.id == slot_id).first()
    if not slot:
        logger.error(f"No slot found with ID={slot_id}.")
        raise HTTPException(status_code=404, detail="Slot not found")

    logger.info(f"Slot found: {slot}")

    # Step 2: Fetch the most recent 'beamline' event for the slot
    recent_beamline_event = (
        db.query(LogisticsEventModel)
        .filter(
            LogisticsEventModel.slot_id == slot_id,
            LogisticsEventModel.event_type == "beamline"
        )
        .order_by(LogisticsEventModel.timestamp.desc())
        .first()
    )

    if not recent_beamline_event:
        logger.error(f"No 'beamline' event found for slot_id={slot_id}.")
        raise HTTPException(
            status_code=404,
            detail="No 'beamline' event found for the given slot"
        )

    logger.info(f"Found beamline event: {recent_beamline_event}")

    # Step 3: Retrieve the Dewar from the recent event
    dewar = db.query(DewarModel).filter(DewarModel.id == recent_beamline_event.dewar_id).first()

    if not dewar:
        logger.error(
            f"No dewar found for the most recent beamline event. Dewar ID={recent_beamline_event.dewar_id}"
        )
        raise HTTPException(
            status_code=404,
            detail="No dewar associated with the most recent 'beamline' event"
        )

    logger.info(f"Associated dewar found: {dewar}")

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

    if not pucks:
        logger.warning(f"No pucks found for Dewar ID={dewar.id}.")
    else:
        logger.info(f"Retrieved pucks for Dewar ID={dewar.id}: {pucks}")

    return pucks