aaredb/backend/app/routers/puck.py

from datetime import datetime
from fastapi import APIRouter, HTTPException, status, Depends
from sqlalchemy.orm import Session, joinedload
from sqlalchemy.sql import func
from typing import List
import uuid
import re
from app.schemas import (
    Puck as PuckSchema,
    PuckCreate,
    PuckUpdate,
    PuckWithTellPosition,
    Sample,
    SetTellPositionRequest,
    DataCollectionParameters,
)
from app.models import (
    Puck as PuckModel,
    PuckEvent as PuckEventModel,
    Sample as SampleModel,
    LogisticsEvent as LogisticsEventModel,
    Dewar as DewarModel,
    SampleEvent,
    Beamtime as BeamtimeModel,
)
from app.dependencies import get_db
import logging

router = APIRouter()

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

VALID_TELL_OPTIONS = {"X06SA", "X06DA", "X10SA"}


def validate_tell(tell: str):
    if tell not in VALID_TELL_OPTIONS:
        raise ValueError(
            f"Invalid tell: {tell}. Must be one of {', '.join(VALID_TELL_OPTIONS)}"
        )


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


def resolve_slot_id(slot_identifier: str) -> int:
    """
    Convert a slot identifier (either numeric or alias) to a numeric slot ID.

    Args:
        slot_identifier (str): The slot identifier to resolve (e.g., "PXI",
         "PXII", "48").

    Returns:
        int: The numeric slot ID corresponding to the identifier.

    Raises:
        HTTPException: If the slot identifier is invalid or unrecognized.
    """
    # Map slot identifier keywords to numeric slot IDs
    slot_aliases = {
        "PXI": 47,
        "PXII": 48,
        "PXIII": 49,
        "X06SA": 47,
        "X10SA": 48,
        "X06DA": 49,
    }

    # Try to resolve the identifier
    try:
        return int(slot_identifier)  # If it's a numeric slot ID, return it directly
    except ValueError:
        # Convert alias to slot ID using the mapping
        slot_id = slot_aliases.get(slot_identifier.upper())
        if slot_id:
            return slot_id

        # Log error and raise an exception for invalid identifiers
        logger.error(f"Invalid slot identifier: {slot_identifier}")
        raise HTTPException(
            status_code=400, detail=f"Invalid slot identifier: {slot_identifier}"
        )


def get_pucks_at_beamline(slot_id: int, db: Session) -> List[PuckWithTellPosition]:
    """
    Fetch all pucks currently located at the beamline for a given slot ID.
    """
    # Subquery: Latest logistic event for each dewar
    latest_event_subquery = (
        db.query(
            LogisticsEventModel.dewar_id.label("dewar_id"),
            func.max(LogisticsEventModel.timestamp).label("latest_event_time"),
        )
        .group_by(LogisticsEventModel.dewar_id)
        .subquery(name="latest_event_subquery")
    )

    # Query dewars in the slot with the latest event "beamline"
    dewars = (
        db.query(DewarModel)
        .join(LogisticsEventModel, DewarModel.id == LogisticsEventModel.dewar_id)
        .join(
            latest_event_subquery,
            (LogisticsEventModel.dewar_id == latest_event_subquery.c.dewar_id)
            & (
                LogisticsEventModel.timestamp
                == latest_event_subquery.c.latest_event_time
            ),
        )
        .filter(
            LogisticsEventModel.slot_id == slot_id,
            LogisticsEventModel.event_type == "beamline",
        )
        .all()
    )

    if not dewars:
        logger.warning(f"No dewars found for slot ID: {slot_id}")
        return []

    # Map dewars to their details
    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: Latest event for each puck
    latest_puck_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")
    )

    # Query pucks for the selected dewars
    pucks_with_latest_events = (
        db.query(
            PuckModel,
            PuckEventModel.event_type,
            PuckEventModel.tell_position,
            PuckEventModel.timestamp,  # Useful for debugging or edge cases
            DewarModel,
        )
        .join(
            latest_puck_event_subquery,
            PuckModel.id == latest_puck_event_subquery.c.puck_id,
            isouter=True,
        )
        .join(
            PuckEventModel,
            (PuckEventModel.puck_id == latest_puck_event_subquery.c.puck_id)
            & (
                PuckEventModel.timestamp
                == latest_puck_event_subquery.c.latest_event_time
            ),
            isouter=True,
        )
        .join(DewarModel, PuckModel.dewar_id == DewarModel.id, isouter=True)
        .filter(PuckModel.dewar_id.in_(dewar_ids))
        .distinct()  # Ensure no duplicates
        .all()
    )

    # Prepare the results
    results = {}
    for (
        puck,
        event_type,
        tell_position,
        event_timestamp,
        dewar,
    ) in pucks_with_latest_events:
        dewar_name = dewar_map.get(puck.dewar_id)
        pgroup = dewar_pgroups.get(puck.dewar_id)

        # If the event is None or explicitly a "puck_removed", set `tell_position=None`
        if event_type is None or event_type == "puck_removed":
            tell_position = None

        # Always replace results since we are processing the latest event
        results[puck.id] = 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,  # Respect if `None` is explicitly set
        )

    return list(results.values())


@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(
    payload: SetTellPositionRequest,  # Accept the wrapped request as a single payload
    db: Session = Depends(get_db),
):
    # Extract the tell position (slot identifier) and the list of pucks
    tell = payload.tell
    pucks = payload.pucks

    try:
        # Resolve slot ID from the provided identifier
        slot_id = resolve_slot_id(tell)
    except Exception as e:
        raise HTTPException(
            status_code=400,
            detail=f"Invalid tell or slot identifier: {tell}." f" Error: {str(e)}",
        )

    # Fetch existing pucks at the beamline slot
    pucks_at_beamline = get_pucks_at_beamline(slot_id, db)
    beamline_puck_map = {
        normalize_puck_name(puck.puck_name): puck for puck in pucks_at_beamline
    }

    # Check if the payload has any pucks
    if not pucks:  # Empty payload case
        results = []

        # Deduplicate pucks based on ID
        unique_pucks_at_beamline = {
            puck.id: puck for puck in pucks_at_beamline
        }.values()

        for puck in unique_pucks_at_beamline:
            # Fetch the most recent event for the puck
            last_event = (
                db.query(PuckEventModel)
                .filter(
                    PuckEventModel.puck_id == puck.id,
                )
                .order_by(
                    PuckEventModel.id.desc()
                )  # Order by timestamp ensures we get the latest event
                .first()
            )

            # Log the last event for the puck
            if last_event:
                logger.info(
                    f"Processing puck: {puck.puck_name}, "
                    f"Last Event -> Type: {last_event.event_type}, "
                    f"Tell Position: {last_event.tell_position}, "
                    f"Timestamp: {last_event.timestamp}"
                )
            else:
                logger.info(
                    f"Processing puck: {puck.puck_name}, No events found for this puck"
                )

            # Remove all pucks, including those without events or with None
            # tell_position
            if last_event.tell_position is not None:
                try:
                    # Add a puck_removed event
                    remove_event = PuckEventModel(
                        puck_id=puck.id,
                        tell=None,  # Nullify the `tell` for removal
                        tell_position=None,
                        event_type="puck_removed",
                        timestamp=datetime.now(),
                    )
                    db.add(remove_event)

                    # Record this removal in the response
                    results.append(
                        {
                            "puck_name": puck.puck_name,
                            "tell": tell,
                            "removed_position": last_event.tell_position
                            if last_event
                            else None,
                            "status": "removed",
                            "message": "Puck removed due to empty payload.",
                        }
                    )
                except Exception as e:
                    # Handle and log the error for this particular puck
                    results.append(
                        {
                            "puck_name": puck.puck_name,
                            "error": str(e),
                        }
                    )

        # Commit all removal events and return the results
        db.commit()
        return results

    # If the payload contains pucks, continue with the regular logic
    results = []
    processed_pucks = (
        set()
    )  # To track pucks that have been processed (unchanged or updated)

    # Existing pucks' most recent events
    last_events_map = {
        puck.id: db.query(PuckEventModel)
        .filter(
            PuckEventModel.puck_id == puck.id,
        )
        .order_by(PuckEventModel.id.desc())
        .first()
        for puck in pucks_at_beamline
    }

    # Step 1: Process each puck in the payload
    for puck_data in pucks:
        try:
            puck_name = puck_data.puck_name
            normalized_name = normalize_puck_name(puck_name)
            new_position = puck_data.tell_position

            existing_puck = beamline_puck_map.get(normalized_name)
            if not existing_puck:
                # If the puck is not found, it's a potential error
                results.append(
                    {
                        "puck_name": puck_name,
                        "error": f"Puck '{puck_name}' not found at the beamline.",
                    }
                )
                continue

            processed_pucks.add(existing_puck.id)  # Mark this puck as processed

            # Check if the tell position is unchanged
            last_event = last_events_map.get(existing_puck.id)
            if last_event and last_event.tell_position == new_position:
                results.append(
                    {
                        "puck_name": puck_name,
                        "tell": tell,
                        "current_position": new_position,
                        "status": "unchanged",
                        "message": "No change in tell_position.",
                    }
                )
                continue

            # Add a "puck_removed" event if the position is being changed (old
            # position removed)
            if last_event and last_event.tell_position is not None:
                remove_event = PuckEventModel(
                    puck_id=existing_puck.id,
                    tell=None,
                    tell_position=None,
                    event_type="puck_removed",
                    timestamp=datetime.now(),
                )
                db.add(remove_event)

            # Add a new "tell_position_set" event (new position)
            new_event = PuckEventModel(
                puck_id=existing_puck.id,
                tell=tell,
                tell_position=new_position,
                event_type="tell_position_set",
                timestamp=datetime.now(),
            )
            db.add(new_event)

            results.append(
                {
                    "puck_name": puck_name,
                    "tell": tell,
                    "new_position": new_position,
                    "previous_position": (
                        last_event.tell_position if last_event else None
                    ),
                    "status": "updated",
                    "message": "Tell position updated successfully.",
                }
            )

        except Exception as e:
            results.append(
                {
                    "puck_name": puck_name,
                    "error": str(e),
                }
            )

    # Step 2: Handle "absent" pucks for removal
    for puck in pucks_at_beamline:
        if puck.id not in processed_pucks:  # This puck was not in the payload
            last_event = last_events_map.get(
                puck.id
            )  # Fetch the last event for the puck

            # Only remove pucks that have a valid last event with a non-null
            # tell_position
            if last_event and last_event.tell_position is not None:
                try:
                    remove_event = PuckEventModel(
                        puck_id=puck.id,
                        tell=None,
                        tell_position=None,
                        event_type="puck_removed",
                        timestamp=datetime.now(),
                    )
                    db.add(remove_event)
                    results.append(
                        {
                            "puck_name": puck.puck_name,
                            "tell": tell,
                            "removed_position": last_event.tell_position,
                            "status": "removed",
                            "message": "Puck removed from tell_position.",
                        }
                    )
                except Exception as e:
                    results.append(
                        {
                            "puck_name": puck.puck_name,
                            "error": str(e),
                        }
                    )

    # Step 3: Commit all changes to the DB
    db.commit()

    return results


@router.get("/with-tell-position", response_model=List[PuckWithTellPosition])
async def get_pucks_with_tell_position(
    tell: str,  # Specify tell as a query parameter
    db: Session = Depends(get_db),
):
    """
    Retrieve all pucks with a valid `tell_position` set (non-null),
    their associated samples, and the latest `tell_position` value (if any),
    filtered by a specific `tell`.
    """
    # Validate the incoming `tell` value
    try:
        validate_tell(tell)  # Ensure `tell` is valid using predefined valid options
    except ValueError as error:
        raise HTTPException(
            status_code=400, detail=str(error)
        )  # Raise error for invalid tells

    # Step 1: Prepare a subquery to fetch the latest event timestamp for each puck.
    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: Main query - fetch pucks with latest `tell_position` event details
    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
        )  # Optional, include related dewar info
        .filter(
            PuckEventModel.tell_position.isnot(None)
        )  # Only include non-null `tell_position`
        .filter(
            PuckEventModel.event_type == "tell_position_set"
        )  # Only include relevant event types
        .filter(PuckEventModel.tell == tell)  # Filter by the specific `tell` variable
        .all()
    )

    # Return an empty list if no relevant pucks are found
    if not pucks_with_events:
        return []

    # Step 3: Construct the response with pucks and their valid tell_position
    results = []
    for puck, event, dewar in pucks_with_events:
        # 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 and dewar.pgroups
                else None,  # Replace later by puck pgroup if needed
                samples=[
                    Sample(
                        id=sample.id,
                        sample_name=sample.sample_name,
                        mount_count=sample.mount_count,
                        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),
            )
        )

    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.
    Args:
        slot_identifier (str): The slot identifier (e.g., "PXI", "48").
        db (Session): Database session dependency.

    Returns:
        List[PuckWithTellPosition]: List of pucks in the specified slot.
    """
    # Resolve the slot identifier to a numeric slot ID using the function
    try:
        slot_id = resolve_slot_id(slot_identifier)
    except HTTPException as e:
        logger.error(
            f"Failed to resolve slot identifier: {slot_identifier}. Error: {e.detail}"
        )
        raise e

    logger.info(f"Resolved slot identifier '{slot_identifier}' to Slot ID: {slot_id}")

    # Fetch the pucks at the beamline for the resolved slot ID
    pucks = get_pucks_at_beamline(slot_id, db)

    if not pucks:
        logger.warning(
            f"No pucks found for the slot '{slot_identifier}' (ID: {slot_id})"
        )
        raise HTTPException(
            status_code=404, detail=f"No pucks found for slot '{slot_identifier}'"
        )

    logger.info(
        f"Found {len(pucks)} pucks for slot '{slot_identifier}' (ID: {slot_id})"
    )

    return pucks


@router.patch("/puck/{puck_id}/assign-beamtime", operation_id="assignPuckToBeamtime")
async def assign_beamtime_to_puck(
    puck_id: int,
    beamtime_id: int,  # expects ?beamtime_id=123 in the query
    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")

    beamtime = (
        db.query(BeamtimeModel).filter(BeamtimeModel.id == beamtime_id).first()
        if beamtime_id
        else None
    )

    if beamtime_id == 0:
        puck.beamtimes = []
    else:
        puck.beamtimes = [
            beamtime
        ]  # or use .append(beamtime) if you want to support multiple

    db.commit()
    db.refresh(puck)
    # Update samples as well
    for sample in puck.samples:
        has_sample_event = (
            db.query(SampleEvent).filter(SampleEvent.sample_id == sample.id).count() > 0
        )
        if not has_sample_event:
            if beamtime_id == 0:
                sample.beamtimes = []
            else:
                sample.beamtimes = [beamtime]
    db.commit()
    return {"status": "success", "puck_id": puck.id, "beamtime_id": beamtime_id}


@router.get(
    "/by-beamtime/{beamtime_id}",
    response_model=List[PuckSchema],
    operation_id="get_pucks_by_beamtime",
)
async def get_pucks_by_beamtime(beamtime_id: int, db: Session = Depends(get_db)):
    beamtime = (
        db.query(BeamtimeModel)
        .options(joinedload(BeamtimeModel.pucks))  # eager load pucks
        .filter(BeamtimeModel.id == beamtime_id)
        .first()
    )
    if not beamtime:
        raise HTTPException(status_code=404, detail="Beamtime not found")
    return beamtime.pucks