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
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("/pucks/slot/{slot_id}", response_model=List[PuckSchema])
async def get_pucks_by_dewar_slot(slot_id: int, db: Session = Depends(get_db)):
    """
    Retrieve all pucks for a given slot_id by finding the associated dewar.
    """
    # Query the Slot table to get the associated dewar based on slot_id
    slot = db.query(SlotModel).filter(SlotModel.id == slot_id).first()
    if not slot:
        raise HTTPException(status_code=404, detail="Slot not found")

    # Check if the slot has a dewar associated
    if not slot.dewar:
        return []  # Return an empty list if no dewar is associated

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

    # Return the list of pucks (empty if no pucks are found)
    return pucks