from datetime import datetime from fastapi import APIRouter, HTTPException, status, Depends from sqlalchemy.orm import Session from sqlalchemy.sql import func from typing import List import uuid import re from app.schemas import ( Puck as PuckSchema, PuckCreate, PuckUpdate, PuckWithTellPosition, Sample, SetTellPosition, ) from app.models import ( Puck as PuckModel, PuckEvent as PuckEventModel, Sample as SampleModel, Slot as SlotModel, LogisticsEvent as LogisticsEventModel, Dewar as DewarModel, ) from app.dependencies import get_db import logging router = APIRouter() logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) def normalize_puck_name(name: str) -> str: """ Normalize a puck_name to remove special characters and ensure consistent formatting. """ name = str(name).strip().replace(" ", "_").upper() name = re.sub(r"[^A-Z0-9]", "", name) # Remove special characters return name @router.get("/", response_model=List[PuckSchema]) async def get_pucks(db: Session = Depends(get_db)): return db.query(PuckModel).all() @router.put("/set-tell-positions", status_code=status.HTTP_200_OK) async def set_tell_positions( pucks: List[SetTellPosition], db: Session = Depends(get_db) ): results = [] # Helper function to normalize puck names for database querying def normalize_puck_name(name: str) -> str: return str(name).strip().replace(" ", "_").upper() if not pucks: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail="Payload cannot be empty. Provide at least one puck.", ) # Retrieve all pucks in the database with their most recent # `tell_position_set` event all_pucks_with_last_event = ( db.query(PuckModel, PuckEventModel) .outerjoin(PuckEventModel, PuckEventModel.puck_id == PuckModel.id) .filter(PuckEventModel.event_type == "tell_position_set") .order_by(PuckEventModel.puck_id, PuckEventModel.timestamp.desc()) .all() ) # Dictionary mapping each puck's ID to its latest event last_events = {} for puck, last_event in all_pucks_with_last_event: if puck.id not in last_events: # Only store the latest event for each puck last_events[puck.id] = last_event # Track processed puck IDs to avoid double-processing processed_pucks = set() # Process pucks provided in the payload for puck_data in pucks: try: # Extract data from input puck_name = puck_data.puck_name new_position = ( puck_data.tell_position ) # Combined from segment + puck_in_segment normalized_name = normalize_puck_name(puck_name) # Find puck in the database puck = ( db.query(PuckModel) .filter( func.replace(func.upper(PuckModel.puck_name), "-", "") == normalized_name ) .first() ) if not puck: raise ValueError(f"Puck with name '{puck_name}' not found.") # Mark this puck as processed processed_pucks.add(puck.id) # Query the last event for this puck last_event = last_events.get(puck.id) # Rule 1: Skip if the last event's `tell_position` matches the new position if last_event and last_event.tell_position == new_position: results.append( { "puck_name": puck.puck_name, "current_position": new_position, "status": "unchanged", "message": "No change in tell_position. No event created.", } ) continue # Rule 2: Add a "puck_removed" event if the last tell_position is not None if last_event and last_event.tell_position is not None: remove_event = PuckEventModel( puck_id=puck.id, tell_position=None, event_type="puck_removed", # Event type set to "puck_removed" timestamp=datetime.utcnow(), ) db.add(remove_event) # Add a new "tell_position_set" event if new_position: new_event = PuckEventModel( puck_id=puck.id, tell_position=new_position, event_type="tell_position_set", timestamp=datetime.utcnow(), ) db.add(new_event) results.append( { "puck_name": puck.puck_name, "new_position": new_position, "previous_position": last_event.tell_position if last_event else None, "status": "updated", "message": "The tell_position was updated successfully.", } ) db.commit() except Exception as e: # Handle individual puck errors results.append({"puck_name": puck_data.puck_name, "error": str(e)}) # Process pucks not included in the payload but present in the database for puck_id, last_event in last_events.items(): # Skip pucks already processed in the previous loop if puck_id in processed_pucks: continue puck = db.query(PuckModel).filter(PuckModel.id == puck_id).first() if not puck: continue # Skip if the last event's tell_position is already null if not last_event or last_event.tell_position is None: continue try: # Add a "puck_removed" event remove_event = PuckEventModel( puck_id=puck.id, tell_position=None, event_type="puck_removed", # Event type set to "puck_removed" timestamp=datetime.utcnow(), ) db.add(remove_event) results.append( { "puck_name": puck.puck_name, "removed_position": last_event.tell_position, "status": "removed", "message": "Puck is not in payload and" " has been marked as removed from tell_position.", } ) db.commit() except Exception as e: # Handle errors for individual puck removal results.append({"puck_name": puck.puck_name, "error": str(e)}) return results @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.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_identifier}", response_model=List[PuckWithTellPosition]) async def get_pucks_by_slot(slot_identifier: str, db: Session = Depends(get_db)): """ Retrieve all pucks in a slot with their latest `tell_position`. """ # Map keywords to slot IDs slot_aliases = { "PXI": 47, "PXII": 48, "PXIII": 49, "X06SA": 47, "X10SA": 48, "X06DA": 49, } # Resolve slot ID or alias try: slot_id = int(slot_identifier) except ValueError: slot_id = slot_aliases.get(slot_identifier.upper()) if not slot_id: logger.error(f"Invalid slot alias: {slot_identifier}") raise HTTPException( status_code=400, detail=f"Invalid slot identifier: {slot_identifier}" ) logger.info(f"Resolved slot identifier: {slot_identifier} to Slot ID: {slot_id}") # Verify slot existence slot = db.query(SlotModel).filter(SlotModel.id == slot_id).first() if not slot: logger.error(f"Slot not found: {slot_identifier}") raise HTTPException( status_code=404, detail=f"Slot not found for identifier {slot_identifier}" ) # Fetch dewars in the slot dewars = ( db.query(DewarModel) .join(LogisticsEventModel, DewarModel.id == LogisticsEventModel.dewar_id) .filter( LogisticsEventModel.slot_id == slot_id, LogisticsEventModel.event_type == "beamline", ) .all() ) if not dewars: logger.warning(f"No dewars found for slot: {slot_identifier}") raise HTTPException( status_code=404, detail=f"No dewars found for slot {slot_identifier}" ) logger.info( f"Found dewars for slot {slot_identifier}: {[dewar.id for dewar in dewars]}" ) dewar_ids = [dewar.id for dewar in dewars] dewar_map = {dewar.id: dewar.dewar_name for dewar in dewars} # Subquery to fetch the latest tell_position for each puck subquery = ( db.query( PuckEventModel.puck_id, func.max(PuckEventModel.timestamp).label("latest_event_time"), ) .filter(PuckEventModel.event_type == "tell_position_set") .group_by(PuckEventModel.puck_id) .subquery() ) # Fetch pucks with their latest tell_position pucks_with_positions = ( db.query(PuckModel, PuckEventModel.tell_position) .outerjoin(subquery, subquery.c.puck_id == PuckModel.id) .outerjoin( PuckEventModel, (PuckEventModel.puck_id == PuckModel.id) & (PuckEventModel.timestamp == subquery.c.latest_event_time), ) .filter(PuckModel.dewar_id.in_(dewar_ids)) .all() ) # Log the results of the subquery and pucks fetched: logger.debug(f"Results from subquery (tell_position): {pucks_with_positions}") if not pucks_with_positions: logger.warning(f"No pucks found for slot: {slot_identifier}") raise HTTPException( status_code=404, detail=f"No pucks found for slot '{slot_identifier}'" ) # Prepare results: results = [] for puck, tell_position in pucks_with_positions: logger.debug( f"Puck ID: {puck.id}, Name: {puck.puck_name}," f" Tell Position: {tell_position}" ) dewar_name = dewar_map.get(puck.dewar_id, "Unknown") # Prepare the PuckWithTellPosition instance results.append( PuckWithTellPosition( id=puck.id, puck_name=puck.puck_name, puck_type=puck.puck_type, puck_location_in_dewar=str(puck.puck_location_in_dewar) if puck.puck_location_in_dewar else None, dewar_id=puck.dewar_id, dewar_name=dewar_name, tell_position=tell_position, # Latest tell_position from subquery ) ) logger.info(f"Final response prepared for slot {slot_identifier}: {results}") return results @router.get("/with-tell-position", response_model=List[PuckWithTellPosition]) async def get_pucks_with_tell_position(db: Session = Depends(get_db)): """ Retrieve all pucks with a `tell_position` set (not null), their associated samples, and the latest `tell_position` value (if any). """ # Query pucks with their latest `tell_position_set` event where # `tell_position` is not null pucks_with_events = ( db.query(PuckModel, PuckEventModel) .join(PuckEventModel, PuckModel.id == PuckEventModel.puck_id) .filter( PuckEventModel.tell_position.isnot(None) ) # Ensure only non-null tell_positions .order_by(PuckEventModel.timestamp.desc()) # Get the most recent event .distinct(PuckModel.id) # Ensure one row per puck (latest event is prioritized) .all() ) if not pucks_with_events: raise HTTPException( status_code=404, detail="No pucks with a `tell_position` found.", ) # Construct the response with pucks and their latest tell_position results = [] for puck, event in pucks_with_events: # Retrieve associated samples for this puck samples = db.query(SampleModel).filter(SampleModel.puck_id == puck.id).all() # Construct the response model results.append( PuckWithTellPosition( id=int(puck.id), # Explicit casting puck_name=str(puck.puck_name), puck_type=str(puck.puck_type), puck_location_in_dewar=str(puck.puck_location_in_dewar) if puck.puck_location_in_dewar else None, dewar_id=int(puck.dewar_id) if puck.dewar_id else None, samples=[ Sample( id=sample.id, sample_name=sample.sample_name, position=sample.position, ) for sample in samples ], tell_position=str(event.tell_position) if event else None, # Include tell_position ) ) return results