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,
    PuckWithTellPosition,
    Sample,
    SetTellPosition,
    DataCollectionParameters,
)
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 = re.sub(r"[^A-Z0-9]", "", name.upper())  # 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 = []

    # 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)
        .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("/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).
    Only include pucks when their latest event has a `tell_position`
     set and matches "tell_position_set".
    """
    # Step 1: Prepare a subquery to fetch the latest event timestamp for each
    # puck with a non-null tell_position
    latest_event_subquery = (
        db.query(
            PuckEventModel.puck_id,
            func.max(PuckEventModel.timestamp).label("latest_timestamp"),
        )
        .group_by(PuckEventModel.puck_id)  # Group by puck
        .subquery()
    )

    # Step 2: Query the pucks and their latest `tell_position` by joining the subquery
    pucks_with_events = (
        db.query(PuckModel, PuckEventModel, DewarModel)
        .join(PuckEventModel, PuckModel.id == PuckEventModel.puck_id)
        .join(
            latest_event_subquery,
            (PuckEventModel.puck_id == latest_event_subquery.c.puck_id)
            & (PuckEventModel.timestamp == latest_event_subquery.c.latest_timestamp),
        )
        .outerjoin(
            DewarModel, PuckModel.dewar_id == DewarModel.id
        )  # Outer join with DewarModel
        .all()
    )

    if not pucks_with_events:
        return []

    # Step 3: Construct the response with pucks and their latest tell_position
    results = []
    # Debug output for verification
    print(f"Pucks with Events and Dewars: {pucks_with_events}")
    for puck, event, dewar in pucks_with_events:
        print(f"Puck: {puck}, Event: {event}, Dewar: {dewar}")
        if event.tell_position is None:
            continue
        # Fetch 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),
                puck_name=str(puck.puck_name),
                puck_type=str(puck.puck_type),
                puck_location_in_dewar=int(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,
                dewar_name=str(dewar.dewar_name)
                if dewar and dewar.dewar_name
                else None,
                pgroup=str(dewar.pgroups)
                if dewar.pgroups
                else None,  # will be replaced later by puck pgroup
                samples=[
                    Sample(
                        id=sample.id,
                        sample_name=sample.sample_name,
                        position=sample.position,
                        puck_id=sample.puck_id,
                        data_collection_parameters=(
                            DataCollectionParameters(
                                **sample.data_collection_parameters
                            )
                            if isinstance(sample.data_collection_parameters, dict)
                            else sample.data_collection_parameters
                        ),
                    )
                    for sample in samples
                ],
                tell_position=str(event.tell_position) if event else None,
            )
        )

    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, reporting their latest event and
     `tell_position` value.
    """
    # 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}
    dewar_pgroups = {dewar.id: dewar.pgroups for dewar in dewars}

    # Subquery to fetch the latest event for each puck (any type of event)
    latest_event_subquery = (
        db.query(
            PuckEventModel.puck_id.label("puck_id"),
            func.max(PuckEventModel.timestamp).label("latest_event_time"),
        )
        .group_by(PuckEventModel.puck_id)
        .subquery(name="latest_event_subquery")  # Explicitly name the subquery
    )

    # Main query to fetch pucks and their latest events
    pucks_with_latest_events = (
        db.query(
            PuckModel,
            PuckEventModel.event_type,
            PuckEventModel.tell_position,
            DewarModel,  # Include DewarModel
        )
        .join(  # Join pucks with the latest event
            latest_event_subquery,
            PuckModel.id == latest_event_subquery.c.puck_id,
            isouter=True,
        )
        .join(  # Fetch event details from the latest event timestamp
            PuckEventModel,
            (PuckEventModel.puck_id == latest_event_subquery.c.puck_id)
            & (PuckEventModel.timestamp == latest_event_subquery.c.latest_event_time),
            isouter=True,
        )
        .join(  # Join with DewarModel to get dewar details
            DewarModel,
            PuckModel.dewar_id == DewarModel.id,
            isouter=True,
        )
        .filter(PuckModel.dewar_id.in_(dewar_ids))  # Restrict pucks to relevant dewars
        .all()
    )

    # Log the results of the query
    logger.debug(f"Results from query (latest events): {pucks_with_latest_events}")

    if not pucks_with_latest_events:
        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 the final response
    results = []
    for puck, event_type, dewar, tell_position in pucks_with_latest_events:
        logger.debug(
            f"Puck ID: {puck.id}, Name: {puck.puck_name}, Event Type: {event_type}, "
            f"Tell Position: {tell_position}"
        )

        dewar_name = dewar_map.get(puck.dewar_id, "Unknown")
        pgroup = dewar_pgroups.get(
            puck.dewar_id
        )  # will be replaced later by puck pgroup

        # For pucks with no events or whose latest event is `puck_removed`, set
        # `tell_position` to None
        if event_type is None or event_type == "puck_removed":
            tell_position = None

        # Construct the response model
        results.append(
            PuckWithTellPosition(
                id=puck.id,
                pgroup=pgroup,
                puck_name=puck.puck_name,
                puck_type=puck.puck_type,
                puck_location_in_dewar=int(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,
            )
        )

    logger.info(f"Final response for slot {slot_identifier}: {results}")
    return results