aaredb/backend/app/routers/shipment.py

from fastapi import APIRouter, HTTPException, status, Query, Depends
from sqlalchemy.orm import Session
from typing import List, Optional
import uuid
import json
from datetime import date

from app.models import Shipment as ShipmentModel, ContactPerson as ContactPersonModel, Address as AddressModel, Proposal as ProposalModel, Dewar as DewarModel
from app.schemas import ShipmentCreate, Shipment as ShipmentSchema, DewarUpdate, ContactPerson as ContactPersonSchema
from app.schemas import Sample as SampleSchema
from app.database import get_db
from app.crud import get_shipments, get_shipment_by_id

router = APIRouter()

def default_serializer(obj):
    if isinstance(obj, date):
        return obj.isoformat()
    raise TypeError(f"Type {type(obj)} not serializable")

@router.get("", response_model=List[ShipmentSchema])
async def fetch_shipments(shipment_id: Optional[str] = Query(None), db: Session = Depends(get_db)):
    if shipment_id:
        shipment = get_shipment_by_id(db, shipment_id)
        if not shipment:
            raise HTTPException(status_code=404, detail="Shipment not found")
        return [shipment]
    return get_shipments(db)

@router.post("", response_model=ShipmentSchema, status_code=status.HTTP_201_CREATED)
async def create_shipment(shipment: ShipmentCreate, db: Session = Depends(get_db)):
    contact_person = db.query(ContactPersonModel).filter(ContactPersonModel.id == shipment.contact_person_id).first()
    return_address = db.query(AddressModel).filter(AddressModel.id == shipment.return_address_id).first()
    proposal = db.query(ProposalModel).filter(ProposalModel.id == shipment.proposal_id).first()

    if not (contact_person or return_address or proposal):
        raise HTTPException(status_code=404, detail="Associated entity not found")

    shipment_id = f'SHIP-{uuid.uuid4().hex[:8].upper()}'
    db_shipment = ShipmentModel(
        shipment_id=shipment_id,
        shipment_name=shipment.shipment_name,
        shipment_date=shipment.shipment_date,
        shipment_status=shipment.shipment_status,
        comments=shipment.comments,
        contact_person_id=contact_person.id,
        return_address_id=return_address.id,
        proposal_id=proposal.id,
    )

    # Handling dewars association
    if shipment.dewars:
        dewar_ids = [dewar.dewar_id for dewar in shipment.dewars]
        dewars = db.query(DewarModel).filter(DewarModel.id.in_(dewar_ids)).all()
        if len(dewars) != len(shipment.dewars):
            raise HTTPException(status_code=404, detail="One or more dewars not found")
        db_shipment.dewars.extend(dewars)

    db.add(db_shipment)
    db.commit()
    db.refresh(db_shipment)

    return db_shipment

@router.delete("/{shipment_id}", status_code=status.HTTP_204_NO_CONTENT)
async def delete_shipment(shipment_id: str, db: Session = Depends(get_db)):
    shipment = db.query(ShipmentModel).filter(ShipmentModel.shipment_id == shipment_id).first()
    if not shipment:
        raise HTTPException(status_code=404, detail="Shipment not found")
    db.delete(shipment)
    db.commit()
    return

@router.put("/{shipment_id}", response_model=ShipmentSchema)
async def update_shipment(shipment_id: str, updated_shipment: ShipmentCreate, db: Session = Depends(get_db)):
    print("Received payload:", json.dumps(updated_shipment.dict(), indent=2, default=default_serializer))

    shipment = db.query(ShipmentModel).filter(ShipmentModel.shipment_id == shipment_id).first()
    if not shipment:
        raise HTTPException(status_code=404, detail="Shipment not found")

    # Validate relationships by IDs
    contact_person = db.query(ContactPersonModel).filter(ContactPersonModel.id == updated_shipment.contact_person_id).first()
    return_address = db.query(AddressModel).filter(AddressModel.id == updated_shipment.return_address_id).first()
    if not contact_person:
        raise HTTPException(status_code=404, detail="Contact person not found")
    if not return_address:
        raise HTTPException(status_code=404, detail="Return address not found")

    # Update shipment details
    shipment.shipment_name = updated_shipment.shipment_name
    shipment.shipment_date = updated_shipment.shipment_date
    shipment.shipment_status = updated_shipment.shipment_status
    shipment.comments = updated_shipment.comments
    shipment.contact_person_id = updated_shipment.contact_person_id
    shipment.return_address_id = updated_shipment.return_address_id

    # Process and update dewars' details
    for dewar_data in updated_shipment.dewars:
        dewar = db.query(DewarModel).filter(DewarModel.id == dewar_data.dewar_id).first()
        if not dewar:
            raise HTTPException(status_code=404, detail=f"Dewar with ID {dewar_data.dewar_id} not found")

        # Dynamically update the dewar fields based on provided input
        update_fields = dewar_data.dict(exclude_unset=True)
        for key, value in update_fields.items():
            if key == 'contact_person_id':
                contact_person = db.query(ContactPersonModel).filter(ContactPersonModel.id == value).first()
                if not contact_person:
                    raise HTTPException(status_code=404, detail=f"Contact person with ID {value} for Dewar {dewar_data.dewar_id} not found")
            if key == 'return_address_id':
                address = db.query(AddressModel).filter(AddressModel.id == value).first()
                if not address:
                    raise HTTPException(status_code=404, detail=f"Address with ID {value} for Dewar {dewar_data.dewar_id} not found")

        for key, value in update_fields.items():
            if key != 'dewar_id':
                setattr(dewar, key, value)

    db.commit()
    db.refresh(shipment)
    return shipment

@router.post("/{shipment_id}/add_dewar", response_model=ShipmentSchema)
async def add_dewar_to_shipment(shipment_id: str, dewar_id: str, db: Session = Depends(get_db)):
    shipment = db.query(ShipmentModel).filter(ShipmentModel.shipment_id == shipment_id).first()
    if not shipment:
        raise HTTPException(status_code=404, detail="Shipment not found")
    dewar = db.query(DewarModel).filter(DewarModel.id == dewar_id).first()
    if not dewar:
        raise HTTPException(status_code=404, detail="Dewar not found")

    if dewar not in shipment.dewars:
        shipment.dewars.append(dewar)
        db.commit()
        db.refresh(shipment)
    return shipment

@router.delete("/{shipment_id}/remove_dewar/{dewar_id}", response_model=ShipmentSchema)
async def remove_dewar_from_shipment(shipment_id: str, dewar_id: str, db: Session = Depends(get_db)):
    shipment = db.query(ShipmentModel).filter(ShipmentModel.shipment_id == shipment_id).first()
    if not shipment:
        raise HTTPException(status_code=404, detail="Shipment not found")

    shipment.dewars = [dw for dw in shipment.dewars if dw.id != dewar_id]
    db.commit()
    db.refresh(shipment)
    return shipment

@router.get("/contact_persons", response_model=List[ContactPersonSchema])
async def get_shipment_contact_persons(db: Session = Depends(get_db)):
    contact_persons = db.query(ContactPersonModel).all()
    return contact_persons

@router.get("/{shipment_id}/samples", response_model=List[SampleSchema])
def get_samples_in_shipment(shipment_id: str, db: Session = Depends(get_db)):
    shipment = db.query(ShipmentModel).filter(ShipmentModel.shipment_id == shipment_id).first()
    if shipment is None:
        raise HTTPException(status_code=404, detail="Shipment not found")

    samples = []
    for dewar in shipment.dewars:
        for puck in dewar.pucks:
            samples.extend(puck.positions)

    return samples

@router.get("/{shipment_id}/dewars/{dewar_id}/samples", response_model=List[SampleSchema])
def get_samples_in_dewar(shipment_id: str, dewar_id: str, db: Session = Depends(get_db)):
    shipment = db.query(ShipmentModel).filter(ShipmentModel.shipment_id == shipment_id).first()
    if shipment is None:
        raise HTTPException(status_code=404, detail="Shipment not found")

    dewar = db.query(DewarModel).filter(DewarModel.id == dewar_id, DewarModel.shipment_id == shipment_id).first()
    if dewar is None:
        raise HTTPException(status_code=404, detail="Dewar not found in shipment")

    samples = []
    for puck in dewar.pucks:
        samples.extend(puck.positions)

    return samples