aaredb/backend/app/routers/puck.py

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,
)
from app.models import (
    Puck as PuckModel,
    PuckEvent as PuckEventModel,
    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[dict],  # Accept a list of puck definitions from the client
    db: Session = Depends(get_db),
):
    """
    Set the tell positions for multiple pucks.

    Args:
        pucks (List[dict]): A list of puck definitions, where each puck contains:
            - `puck_name` (str): The cleaned name of the puck.
            - `segment` (str): The segment in the dewar (e.g., "A-F").
            - `puck_in_segment` (int): The position within the segment (1-5).

    Returns:
        List[dict]: A list of responses, one for each successfully processed puck.
    """
    results = []

    for puck_data in pucks:
        try:
            # Extract data from input
            puck_name = puck_data.get("puckname")
            segment = puck_data.get("segment")
            puck_in_segment = puck_data.get("puck_in_segment")

            # 1. Validate `segment` and `puck_in_segment`
            if not segment or segment not in "ABCDEF":
                raise ValueError(
                    f"Invalid segment '{segment}'. Must be A, B, C, D, E, or F."
                )
            if not (1 <= puck_in_segment <= 5):
                raise ValueError(
                    f"Invalid puck_in_segment "
                    f"'{puck_in_segment}'. Must be in range 1-5."
                )

            # Generate tell_position
            tell_position = f"{segment}{puck_in_segment}"

            # 2. Find the puck by its cleaned name
            normalized_name = normalize_puck_name(puck_name)

            # Use SQLAlchemy to match normalized names
            puck = (
                db.query(PuckModel)
                .filter(
                    func.replace(func.upper(PuckModel.puck_name), "-", "")
                    == normalized_name
                )
                .first()
            )

            if not puck:
                raise ValueError(
                    f"Puck with cleaned name '{puck_name}' not found in the database."
                )

            # 3. Find the dewar associated with the puck
            dewar = db.query(DewarModel).filter(DewarModel.id == puck.dewar_id).first()
            if not dewar:
                raise ValueError(f"Dewar associated with puck '{puck_name}' not found.")

            # 4. Find the most recent logistics event for the dewar (type 'beamline')
            logistics_event = (
                db.query(LogisticsEventModel)
                .filter(
                    LogisticsEventModel.dewar_id == dewar.id,
                    LogisticsEventModel.event_type == "beamline",
                )
                .order_by(LogisticsEventModel.timestamp.desc())
                .first()
            )
            if not logistics_event:
                raise ValueError(
                    f"No recent 'beamline' logistics event found for dewar "
                    f"'{dewar.dewar_name}'."
                )

            # 5. Retrieve the slot from the logistics event
            slot = (
                db.query(SlotModel)
                .filter(SlotModel.id == logistics_event.slot_id)
                .first()
            )
            if not slot:
                raise ValueError(
                    f"No slot associated with the most recent 'beamline' "
                    f"logistics event "
                    f"for dewar '{dewar.dewar_name}'."
                )

            # 6. Create a PuckEvent to set the 'tell_position'
            new_puck_event = PuckEventModel(
                puck_id=puck.id,
                tell_position=tell_position,
                event_type="tell_position_set",
                timestamp=datetime.utcnow(),
            )
            db.add(new_puck_event)
            db.commit()
            db.refresh(new_puck_event)

            # Add success result to the results list
            results.append(
                {
                    "puck_id": puck.id,
                    "puck_name": puck.puck_name,
                    "dewar_id": dewar.id,
                    "dewar_name": dewar.dewar_name,
                    "slot_id": slot.id,
                    "slot_label": slot.label,
                    "tell_position": tell_position,
                    "event_timestamp": new_puck_event.timestamp,
                }
            )
        except Exception as e:
            # Handle errors for individual pucks and continue processing others
            results.append(
                {
                    "puck_name": 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[dict])
async def get_pucks_by_slot(slot_identifier: str, db: Session = Depends(get_db)):
    """
    Retrieve all pucks associated with all dewars linked to the given slot
    (by ID or keyword) via 'beamline' events.

    - Accepts slot keywords like PXI, PXII, PXIII.
    - Retrieves all dewars (and their names) associated with the slot.
    """
    # Map keywords to slot IDs
    slot_aliases = {
        "PXI": 47,
        "PXII": 48,
        "PXIII": 49,
        "X06SA": 47,
        "X10SA": 48,
        "X06DA": 49,
    }

    # Check if the slot identifier is an alias or ID
    try:
        slot_id = int(slot_identifier)  # If the user provided a numeric ID
        alias = next(
            (k for k, v in slot_aliases.items() if v == slot_id), slot_identifier
        )
    except ValueError:
        slot_id = slot_aliases.get(slot_identifier.upper())  # Try mapping alias
        alias = slot_identifier.upper()  # Keep alias as-is for error messages
        if not slot_id:
            raise HTTPException(
                status_code=400,
                detail="Invalid slot identifier."
                "Must be an ID or one of the following:"
                "PXI, PXII, PXIII, X06SA, X10SA, X06DA.",
            )

    # Verify that the slot exists
    slot = db.query(SlotModel).filter(SlotModel.id == slot_id).first()
    if not slot:
        raise HTTPException(
            status_code=404, detail=f"Slot not found for identifier '{alias}'."
        )

    logger.info(f"Slot found: ID={slot.id}, Label={slot.label}")

    # Retrieve all beamline events associated with the slot
    beamline_events = (
        db.query(LogisticsEventModel)
        .filter(
            LogisticsEventModel.slot_id == slot_id,
            LogisticsEventModel.event_type == "beamline",
        )
        .order_by(LogisticsEventModel.timestamp.desc())
        .all()
    )

    if not beamline_events:
        logger.warning(f"No dewars associated to this beamline '{alias}'.")
        raise HTTPException(
            status_code=404, detail=f"No dewars found for the given beamline '{alias}'."
        )

    logger.info(f"Found {len(beamline_events)} beamline events for slot_id={slot_id}.")

    # Use the beamline events to find all associated dewars
    dewar_ids = {event.dewar_id for event in beamline_events if event.dewar_id}
    dewars = db.query(DewarModel).filter(DewarModel.id.in_(dewar_ids)).all()

    if not dewars:
        logger.warning(f"No dewars found for beamline '{alias}'.")
        raise HTTPException(
            status_code=404, detail=f"No dewars found for beamline '{alias}'."
        )

    logger.info(f"Found {len(dewars)} dewars for beamline '{alias}'.")

    # Create a mapping of dewar_id to dewar_name
    dewar_mapping = {dewar.id: dewar.dewar_name for dewar in dewars}

    # Retrieve all pucks associated with the dewars
    puck_list = (
        db.query(PuckModel)
        .filter(PuckModel.dewar_id.in_([dewar.id for dewar in dewars]))
        .all()
    )

    if not puck_list:
        logger.warning(f"No pucks found for dewars associated with beamline '{alias}'.")
        raise HTTPException(
            status_code=404,
            detail=f"No pucks found for dewars associated with beamline '{alias}'.",
        )

    logger.info(f"Found {len(puck_list)} pucks for beamline '{alias}'.")

    # Add the dewar_name to the output for each puck
    puck_output = [
        {
            "id": puck.id,
            "puck_name": puck.puck_name,
            "puck_type": puck.puck_type,
            "dewar_id": puck.dewar_id,
            "dewar_name": dewar_mapping.get(puck.dewar_id),  # Link dewar_name
        }
        for puck in puck_list
    ]

    # Return the list of pucks with their associated dewar names
    return puck_output