aaredb/backend/app/routers/sample.py

from fastapi import APIRouter, HTTPException, Depends, UploadFile, File, Form
from fastapi.encoders import jsonable_encoder
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,
    SampleEventCreate,
    Sample,
    Image,
    ImageCreate,
    SampleResult,
    ExperimentParametersCreate,
    ExperimentParametersRead,
    ImageInfo,
    ResultResponse,
    ResultCreate,
    Results as ProcessingResults,
    Datasets,
)
from app.models import (
    Puck as PuckModel,
    Sample as SampleModel,
    SampleEvent as SampleEventModel,
    Image as ImageModel,
    Dewar as DewarModel,
    ExperimentParameters as ExperimentParametersModel,
    # ExperimentParameters,
    Results as ResultsModel,
    Jobs as JobModel,
    JobStatus,
)
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(
    active_pgroup: str, db: Session = Depends(get_db)
):
    logging.info(
        "Fetching all pucks with " "samples and events for active_pgroup: %s",
        active_pgroup,
    )

    pucks = (
        db.query(PuckModel)
        .join(PuckModel.samples)  # Join samples related to the puck
        .join(PuckModel.dewar)  # Join the dewar from the puck
        .join(SampleModel.events)  # Join sample events
        .filter(DewarModel.pgroups == active_pgroup)  # Filter by the dewar's group
        .options(
            joinedload(PuckModel.samples).joinedload(SampleModel.events),
            joinedload(PuckModel.dewar),
        )
        .distinct()  # Avoid duplicate puck rows if there are multiple events/samples
        .all()
    )

    if not pucks:
        raise HTTPException(
            status_code=404,
            detail="No pucks found with" " sample events for the active pgroup",
        )

    # Extract samples from each puck if needed
    filtered_samples = []
    for puck in pucks:
        if puck.dewar and getattr(puck.dewar, "pgroups", None) == active_pgroup:
            for sample in puck.samples:
                filtered_samples.append(sample)

    # Depending on what your endpoint expects,
    # you may choose to return pucks or samples.
    # For now, we're returning the list of pucks.
    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`


@router.post(
    "/{sample_id}/upload-images",
    response_model=Image,
    operation_id="upload_sample_image",
)
async def upload_sample_image(
    sample_id: int,
    uploaded_file: UploadFile = File(...),
    comment: str = Form(None),
    db: Session = Depends(get_db),
):
    logging.info(f"Received file: {uploaded_file.filename}")

    # Validate Sample
    sample = db.query(SampleModel).filter(SampleModel.id == sample_id).first()
    if not sample:
        raise HTTPException(status_code=404, detail="Sample not found")

    # Retrieve the most recent sample event for the sample
    sample_event = (
        db.query(SampleEventModel)
        .filter(SampleEventModel.sample_id == sample_id)
        .order_by(SampleEventModel.timestamp.desc())  # Sort by most recent event
        .first()
    )

    if not sample_event:
        logging.debug(f"No events found for sample with id: {sample_id}")
        raise HTTPException(
            status_code=404, detail="No events found for the specified sample"
        )

    # Log the found sample event for debugging
    logging.debug(
        f"Most recent event found for sample_id {sample_id}: "
        f"event_id={sample_event.id}, "
        f"type={sample_event.event_type}, "
        f"timestamp={sample_event.timestamp}"
    )

    # Extract event type and timestamp for directory structure
    event_type = sample_event.event_type
    event_timestamp = sample_event.timestamp.strftime("%Y-%m-%d_%H-%M-%S")

    # Define Directory Structure
    pgroup = sample.puck.dewar.pgroups  # adjust to sample or puck pgroups as needed
    today = datetime.now().strftime("%Y-%m-%d")
    dewar_name = (
        sample.puck.dewar.dewar_name
        if sample.puck and sample.puck.dewar
        else "default_dewar"
    )
    puck_name = sample.puck.puck_name if sample.puck else "default_puck"
    position = sample.position if sample.position else "default_position"

    # Add 'run/event' specific details to the folder structure
    base_dir = Path(
        f"images/{pgroup}/{today}/{dewar_name}/{puck_name}/"
        f"{position}/{event_type}_{event_timestamp}"
    )
    base_dir.mkdir(parents=True, exist_ok=True)

    # Validate MIME type and Save the File
    if not uploaded_file.content_type.startswith("image/"):
        raise HTTPException(
            status_code=400,
            detail=f"Invalid file type: {uploaded_file.filename}."
            f"Only images are accepted.",
        )

    file_path = base_dir / uploaded_file.filename
    logging.debug(f"Saving file {uploaded_file.filename} to {file_path}")

    try:
        with file_path.open("wb") as buffer:
            shutil.copyfileobj(uploaded_file.file, buffer)
        logging.info(f"File saved: {file_path}")
    except Exception as e:
        logging.error(f"Error saving file {uploaded_file.filename}: {str(e)}")
        raise HTTPException(
            status_code=500,
            detail=f"Could not save file {uploaded_file.filename}. "
            f"Ensure the server has correct permissions.",
        )

    # Create the payload from the Pydantic schema
    image_payload = ImageCreate(
        pgroup=pgroup,
        comment=comment,
        filepath=str(file_path),
        status="active",
        sample_id=sample_id,
        sample_event_id=int(sample_event.id),  # Link to the most recent sample event
    ).dict()

    # Convert the payload to your mapped SQLAlchemy model instance
    new_image = ImageModel(**image_payload)
    db.add(new_image)
    db.commit()
    db.refresh(new_image)

    logging.info(
        f"Uploaded 1 file for sample {sample_id} and event {sample_event.id} and "
        f"added record {new_image.id} to the database."
    )

    # Returning the mapped SQLAlchemy object, which will be converted to the
    # Pydantic response model.
    return new_image


@router.get(
    "/results", response_model=List[SampleResult], operation_id="get_sample_results"
)
async def get_sample_results(active_pgroup: str, db: Session = Depends(get_db)):
    # Query samples for the active pgroup using joins.
    samples = (
        db.query(SampleModel)
        .join(SampleModel.puck)
        .join(PuckModel.dewar)
        .filter(DewarModel.pgroups == active_pgroup)
        .all()
    )
    if not samples:
        raise HTTPException(
            status_code=404, detail="No samples found for the active pgroup"
        )

    results = []
    for sample in samples:
        # Query images associated with the sample, including the related event_type
        images = (
            db.query(ImageModel)
            .options(joinedload(ImageModel.sample_event))
            .filter(ImageModel.sample_id == sample.id)
            .all()
        )

        # Query experiment parameters (which include beamline parameters) for the
        # sample.
        experiment_parameters = (
            db.query(ExperimentParametersModel)
            .filter(ExperimentParametersModel.sample_id == sample.id)
            .all()
        )
        print("Experiment Parameters for sample", sample.id, experiment_parameters)

        results.append(
            SampleResult(
                sample_id=sample.id,
                sample_name=sample.sample_name,
                puck_name=sample.puck.puck_name if sample.puck else None,
                dewar_name=sample.puck.dewar.dewar_name
                if (sample.puck and sample.puck.dewar)
                else None,
                images=[
                    ImageInfo(
                        id=img.id,
                        filepath=img.filepath,
                        event_type=img.sample_event.event_type
                        if img.sample_event
                        else "Unknown",
                        comment=img.comment,
                    )
                    for img in images
                ],
                experiment_runs=[
                    ExperimentParametersRead(
                        id=ex.id,
                        type=ex.type,
                        run_number=ex.run_number,
                        beamline_parameters=ex.beamline_parameters,
                        sample_id=ex.sample_id,
                    )
                    for ex in experiment_parameters
                ],
            )
        )

    return results


@router.post(
    "/samples/{sample_id}/experiment_parameters",
    response_model=ExperimentParametersRead,
    operation_id="create_experiment_parameters_for_sample",
)
def create_experiment_parameters_for_sample(
    sample_id: int,
    exp_params: ExperimentParametersCreate,
    db: Session = Depends(get_db),
):
    # Calculate the new run_number for the given sample.
    # This assumes that the run_number is computed as one plus the maximum
    # current value.
    last_exp = (
        db.query(ExperimentParametersModel)
        .filter(ExperimentParametersModel.sample_id == sample_id)
        .order_by(ExperimentParametersModel.run_number.desc())
        .first()
    )
    new_run_number = last_exp.run_number + 1 if last_exp else 1

    # Create a new ExperimentParameters record. The beamline_parameters are
    # stored as JSON.
    new_exp = ExperimentParametersModel(
        run_number=new_run_number,
        type=exp_params.type,
        beamline_parameters=exp_params.beamline_parameters.dict()
        if exp_params.beamline_parameters
        else None,
        sample_id=sample_id,
    )
    db.add(new_exp)
    db.commit()
    db.refresh(new_exp)

    # Create a "Collecting" sample event associated with the new experiment parameters
    new_event = SampleEventModel(
        sample_id=sample_id,
        event_type="Collecting",  # The event type
        timestamp=datetime.now(),  # Use current timestamp
    )
    db.add(new_event)
    db.commit()

    return new_exp


@router.patch(
    "/update-dataset/{sample_id}/{run_id}",
    response_model=ExperimentParametersRead,
    operation_id="update_dataset_for_experiment_run",
)
def update_experiment_run_dataset(
    sample_id: int,
    run_id: int,
    dataset: Datasets,
    db: Session = Depends(get_db),
):
    # Find the run for this sample and run_id
    exp = (
        db.query(ExperimentParametersModel)
        .filter(
            ExperimentParametersModel.sample_id == sample_id,
            ExperimentParametersModel.id == run_id,
        )
        .first()
    )
    if not exp:
        raise HTTPException(
            status_code=404,
            detail="ExperimentParameters (run) not found for this sample",
        )

    exp.dataset = jsonable_encoder(dataset)
    db.commit()
    db.refresh(exp)

    # Only create a job if status is "written" and job does not exist yet
    if dataset.status == "written":
        job_exists = (
            db.query(JobModel)
            .filter(JobModel.sample_id == sample_id, JobModel.run_id == run_id)
            .first()
        )
        if not job_exists:
            new_job = JobModel(
                sample_id=sample_id,
                run_id=run_id,
                experiment_parameters=exp,  # adjust this line as appropriate
                status=JobStatus.TODO,
            )
            db.add(new_job)
            db.commit()
            db.refresh(new_job)

    return exp


@router.post(
    "/processing-results", response_model=ResultResponse, operation_id="create_result"
)
def create_result(payload: ResultCreate, db: Session = Depends(get_db)):
    # Check experiment existence
    experiment = (
        db.query(ExperimentParametersModel)
        .filter(ExperimentParametersModel.id == payload.run_id)
        .first()
    )
    if not experiment:
        raise HTTPException(
            status_code=404, detail="Experiment parameters (run) not found"
        )

    result_entry = ResultsModel(
        sample_id=payload.sample_id,
        run_id=payload.run_id,
        result=payload.result.model_dump(),  # Serialize entire result to JSON
    )

    db.add(result_entry)
    db.commit()
    db.refresh(result_entry)

    return ResultResponse(
        id=result_entry.id,
        sample_id=result_entry.sample_id,
        run_id=result_entry.run_id,
        result=payload.result,  # return original payload directly
    )


@router.get(
    "/processing-results/{sample_id}/{run_id}",
    response_model=List[ResultResponse],
    operation_id="get_results_for_run_and_sample",
)
async def get_results_for_run_and_sample(
    sample_id: int, run_id: int, db: Session = Depends(get_db)
):
    results = (
        db.query(ResultsModel)
        .filter(ResultsModel.sample_id == sample_id, ResultsModel.run_id == run_id)
        .all()
    )

    if not results:
        raise HTTPException(status_code=404, detail="Results not found.")

    formatted_results = [
        ResultResponse(
            id=result.id,
            sample_id=result.sample_id,
            run_id=result.run_id,
            result=ProcessingResults(**result.result),
        )
        for result in results
    ]

    return formatted_results