aaredb/backend/app/routers/puck.py

from datetime import datetime
from fastapi import APIRouter, HTTPException, status, Depends
from sqlalchemy.orm import Session, load_only
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,
    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],  # Accept a validated Pydantic model
    db: Session = Depends(get_db),
):
    """
    Set tell positions for multiple pucks with updated rules.

    Args:
        pucks (List[SetTellPosition]): A list including puck_name,
         segment, and puck_in_segment.

    Rules:
        1. If a puck already has a tell_position matching the payload,
         it is ignored (no new timestamps).
        2. If a puck is assigned a different position,
         set the new position while nullifying the previous.
        3. Pucks that have a last `tell_position_set` event with
         a non-null `tell_position` but
           are not in the payload will get
            a `"puck_removed"` event nullifying their position.
        4. If the last event for a puck is already null
         or it has no events, nothing is added or updated.

    Returns:
        List[dict]: Status information for processed and ignored pucks.
    """
    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.",
        )

    # Normalize payload puck names
    input_puck_names = {normalize_puck_name(p.puck_name): p for p in pucks}

    # Retrieve all pucks in the database
    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:
            last_events[puck.id] = last_event

    # 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

            # Normalize puck name
            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.")

            # 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: If the last tell_position is not None, nullify it with a
            # puck_removed event
            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",  # Set event_type to "puck_removed"
                    timestamp=datetime.utcnow(),
                )
                db.add(remove_event)

            # Add a new event with the updated tell_position
            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.",
                    }
                )

            # Commit changes after processing each puck
            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 DB
    for puck_id, last_event in last_events.items():
        puck = db.query(PuckModel).filter(PuckModel.id == puck_id).first()
        normalized_name = normalize_puck_name(puck.puck_name)

        # Check if the puck is missing from the payload
        if (
            normalized_name not in input_puck_names
            and last_event
            and last_event.tell_position is not None
        ):
            try:
                # Add a puck_removed event
                remove_event = PuckEventModel(
                    puck_id=puck.id,
                    tell_position=None,
                    event_type="puck_removed",  # Set event_type to "puck_removed"
                    timestamp=datetime.utcnow(),
                )
                db.add(remove_event)

                # Append to results
                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[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


class SampleModel:
    pass


@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)
    and their associated samples.
    """
    # Query pucks with events where `tell_position` is not null
    pucks = (
        db.query(PuckModel)
        .join(PuckEventModel, PuckModel.id == PuckEventModel.puck_id)
        .filter(PuckEventModel.tell_position.isnot(None))
        .options(
            load_only(
                PuckModel.id,
                PuckModel.puck_name,
                PuckModel.puck_type,
                PuckModel.puck_location_in_dewar,
                PuckModel.dewar_id,
            )
        )
        .all()
    )

    if not pucks:
        raise HTTPException(
            status_code=404,
            detail="No pucks with a `tell_position` found.",
        )

    results = [
        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),
            samples=[
                Sample(
                    id=sample.id,
                    sample_name=sample.sample_name,
                    position=sample.position,
                )
                for sample in db.query(SampleModel)
                .filter(SampleModel.puck_id == puck.id)
                .all()
            ],
        )
        for puck in pucks
    ]

    return results