aaredb/backend/app/routers/sample.py

from fastapi import APIRouter, HTTPException, Depends, UploadFile, File
from sqlalchemy.orm import Session
from pathlib import Path
from typing import List
from datetime import datetime
import shutil
from app.schemas import (
    Puck as PuckSchema,
    Sample as SampleSchema,
    SampleEventResponse,
    SampleEventCreate,
    Sample,
)
from app.models import (
    Puck as PuckModel,
    Sample as SampleModel,
    SampleEvent as SampleEventModel,
)
from app.dependencies import get_db
import logging
from sqlalchemy.orm import joinedload


router = APIRouter()


@router.get("/{puck_id}/samples", response_model=List[SampleSchema])
async def get_samples_with_events(puck_id: int, 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")

    samples = db.query(SampleModel).filter(SampleModel.puck_id == puck_id).all()

    for sample in samples:
        sample.events = (
            db.query(SampleEventModel)
            .filter(SampleEventModel.sample_id == sample.id)
            .all()
        )

    return samples


@router.get("/pucks-samples", response_model=List[PuckSchema])
async def get_all_pucks_with_samples_and_events(db: Session = Depends(get_db)):
    logging.info("Fetching all pucks with samples and events")

    pucks = (
        db.query(PuckModel)
        .options(
            joinedload(PuckModel.samples).joinedload(
                SampleModel.events
            ),  # Correct nested relationship
            joinedload(PuckModel.events),  # If Puck has its own events relationship
        )
        .all()
    )

    if not pucks:
        raise HTTPException(status_code=404, detail="No pucks found in the database")

    return pucks


# Route to post a new sample event
@router.post("/samples/{sample_id}/events", response_model=Sample)
async def create_sample_event(
    sample_id: int, event: SampleEventCreate, db: Session = Depends(get_db)
):
    # Ensure the sample exists
    sample = db.query(SampleModel).filter(SampleModel.id == sample_id).first()
    if not sample:
        raise HTTPException(status_code=404, detail="Sample not found")

    # Create the event
    sample_event = SampleEventModel(
        sample_id=sample_id,
        event_type=event.event_type,
        timestamp=datetime.now(),  # Use the current timestamp
    )
    db.add(sample_event)
    db.commit()
    db.refresh(sample_event)

    # Load events for the sample to be serialized in the response
    sample.events = (
        db.query(SampleEventModel).filter(SampleEventModel.sample_id == sample_id).all()
    )

    return sample  # Return the sample, now including `mount_count`


# Route to fetch the last (most recent) sample event
@router.get("/samples/{sample_id}/events/last", response_model=SampleEventResponse)
async def get_last_sample_event(sample_id: int, db: Session = Depends(get_db)):
    # Ensure the sample exists
    sample = db.query(SampleModel).filter(SampleModel.id == sample_id).first()
    if not sample:
        raise HTTPException(status_code=404, detail="Sample not found")

    # Get the most recent event for the sample
    last_event = (
        db.query(SampleEventModel)
        .filter(SampleEventModel.sample_id == sample_id)
        .order_by(SampleEventModel.timestamp.desc())
        .first()
    )

    if not last_event:
        raise HTTPException(status_code=404, detail="No events found for the sample")

    return SampleEventResponse(
        id=last_event.id,
        sample_id=last_event.sample_id,
        event_type=last_event.event_type,
        timestamp=last_event.timestamp,
    )  # Response will automatically use the SampleEventResponse schema


@router.post("/samples/{sample_id}/upload-images")
async def upload_sample_images(
    sample_id: int,
    uploaded_files: List[UploadFile] = File(...),  # Accept multiple files
    db: Session = Depends(get_db),
):
    """
    Uploads images for a sample and stores them in a directory structure:
    images/user/date/dewar_name/puck_name/position/.

    Args:
        sample_id (int): ID of the sample.
        uploaded_files (List[UploadFile]): List of image files to be uploaded.
        db (Session): SQLAlchemy database session.
    """
    # Fetch sample details from the database
    sample = db.query(SampleModel).filter(SampleModel.id == sample_id).first()
    if not sample:
        raise HTTPException(status_code=404, detail="Sample not found")

    # Retrieve associated dewar_name, puck_name and position
    puck = sample.puck
    if not puck:
        raise HTTPException(
            status_code=404, detail=f"No puck associated with sample ID {sample_id}"
        )

    dewar_name = puck.dewar.dewar_name if puck.dewar else None
    if not dewar_name:
        raise HTTPException(
            status_code=404, detail=f"No dewar associated with puck ID {puck.id}"
        )

    puck_name = puck.puck_name
    position = sample.position

    # Retrieve username (hardcoded for now—can be fetched dynamically if needed)
    username = "e16371"

    # Today's date in the format YYYY-MM-DD
    today = datetime.now().strftime("%Y-%m-%d")

    # Generate the directory path based on the structure
    base_dir = (
        Path("images") / username / today / dewar_name / puck_name / str(position)
    )

    # Create directories if they don't exist
    base_dir.mkdir(parents=True, exist_ok=True)

    # Save each uploaded image to the directory
    for file in uploaded_files:
        # Validate file content type
        if not file.content_type.startswith("image/"):
            raise HTTPException(
                status_code=400,
                detail=f"Invalid file type: {file.filename}. Must be an image.",
            )

        # Create a file path for storing the uploaded file
        file_path = base_dir / file.filename

        try:
            # Save the file
            with file_path.open("wb") as buffer:
                shutil.copyfileobj(file.file, buffer)
        except Exception as e:
            raise HTTPException(
                status_code=500,
                detail=f"Error saving file {file.filename}: {str(e)}",
            )

    return {
        "message": f"{len(uploaded_files)} images uploaded successfully.",
        "path": str(base_dir),  # Return the base directory for reference
    }