aaredb/backend/app/routers/dewar.py

import os
import tempfile
import time
import random
import hashlib
from fastapi import APIRouter, HTTPException, status, Depends, Response
from sqlalchemy.orm import Session, joinedload
from typing import List
import logging
from sqlalchemy.exc import SQLAlchemyError
from pydantic import ValidationError, BaseModel
from app.schemas import (
    Dewar as DewarSchema,
    DewarCreate,
    DewarUpdate,
    DewarType as DewarTypeSchema,
    DewarTypeCreate,
    DewarSerialNumber as DewarSerialNumberSchema,
    DewarSerialNumberCreate,
    Shipment as ShipmentSchema,  # Clearer name for schema
)
from app.models import (
    Dewar as DewarModel,
    Puck as PuckModel,
    Sample as SampleModel,
    DewarType as DewarTypeModel,
    DewarSerialNumber as DewarSerialNumberModel,
)
from app.dependencies import get_db
import qrcode
import io
from io import BytesIO
from PIL import Image
from reportlab.lib.pagesizes import A5, landscape
from reportlab.lib.units import cm
from reportlab.pdfgen import canvas
from app.crud import (
    get_shipments,
    get_shipment_by_id,
)  # Import CRUD functions for shipment

router = APIRouter()


def generate_unique_id(db: Session, length: int = 16) -> str:
    while True:
        base_string = f"{time.time()}{random.randint(0, 10 ** 6)}"
        hash_object = hashlib.sha256(base_string.encode())
        hash_digest = hash_object.hexdigest()
        unique_id = "".join(random.choices(hash_digest, k=length))
        existing_dewar = (
            db.query(DewarModel).filter(DewarModel.unique_id == unique_id).first()
        )
        if not existing_dewar:
            break
    return unique_id


@router.post("/", response_model=DewarSchema, status_code=status.HTTP_201_CREATED)
async def create_dewar(
    dewar: DewarCreate, db: Session = Depends(get_db)
) -> DewarSchema:
    try:
        db_dewar = DewarModel(
            dewar_name=dewar.dewar_name,
            tracking_number=dewar.tracking_number,
            status=dewar.status,
            ready_date=dewar.ready_date,
            shipping_date=dewar.shipping_date,
            arrival_date=dewar.arrival_date,
            returning_date=dewar.returning_date,
            contact_person_id=dewar.contact_person_id,
            return_address_id=dewar.return_address_id,
        )
        db.add(db_dewar)
        db.commit()
        db.refresh(db_dewar)

        for puck_data in dewar.pucks:
            puck = PuckModel(
                dewar_id=db_dewar.id,
                puck_name=puck_data.puck_name,
                puck_type=puck_data.puck_type,
                puck_location_in_dewar=puck_data.puck_location_in_dewar,
            )
            db.add(puck)
            db.commit()
            db.refresh(puck)

            for sample_data in puck_data.samples:
                logging.debug(
                    f"data_collection_parameters: "
                    f"{sample_data.data_collection_parameters}"
                )
                if sample_data.data_collection_parameters is None:
                    serialized_params = {}
                elif hasattr(sample_data.data_collection_parameters, "to_dict"):
                    serialized_params = sample_data.data_collection_parameters.to_dict()
                elif isinstance(sample_data.data_collection_parameters, BaseModel):
                    serialized_params = sample_data.data_collection_parameters.dict(
                        exclude_unset=True
                    )
                elif isinstance(sample_data.data_collection_parameters, dict):
                    serialized_params = sample_data.data_collection_parameters
                else:
                    raise ValueError(
                        "data_collection_parameters must be a dictionary,"
                        "have a to_dict method, or be None"
                    )

                sample = SampleModel(
                    puck_id=puck.id,
                    sample_name=sample_data.sample_name,
                    proteinname=sample_data.proteinname,
                    position=sample_data.position,
                    priority=sample_data.priority,
                    comments=sample_data.comments,
                    data_collection_parameters=serialized_params,
                )
                db.add(sample)
                db.commit()
                db.refresh(sample)

        return db_dewar

    except SQLAlchemyError as e:
        logging.error(f"Database error occurred: {e}")
        raise HTTPException(status_code=500, detail="Internal server error")
    except ValidationError as e:
        logging.error(f"Validation error occurred: {e}")
        raise HTTPException(status_code=400, detail="Validation error")


@router.post("/{dewar_id}/generate-qrcode")
async def generate_dewar_qrcode(dewar_id: int, db: Session = Depends(get_db)):
    dewar = db.query(DewarModel).filter(DewarModel.id == dewar_id).first()
    if not dewar:
        raise HTTPException(status_code=404, detail="Dewar not found")

    if not dewar.unique_id:
        dewar.unique_id = generate_unique_id(db)
        db.commit()

    qr = qrcode.QRCode(version=1, box_size=10, border=5)
    qr.add_data(dewar.unique_id)
    qr.make(fit=True)
    img = qr.make_image(fill="black", back_color="white")

    buf = io.BytesIO()
    img.save(buf)
    buf.seek(0)
    dewar.qrcode = dewar.unique_id
    dewar.qrcode_image = buf.getvalue()
    db.commit()

    return {"message": "QR Code generated", "qrcode": dewar.unique_id}


def generate_label(dewar):
    buffer = BytesIO()
    # Set page orientation to landscape
    c = canvas.Canvas(buffer, pagesize=landscape(A5))

    # Dimensions for the A5 landscape
    page_width, page_height = landscape(A5)

    # Path to the PNG logo
    file_dir = os.path.dirname(os.path.abspath(__file__))
    png_logo_path = os.path.join(file_dir, "Heidi-logo.png")

    # Open the logo with PIL to get its size
    logo_image = Image.open(png_logo_path)
    logo_aspect_ratio = logo_image.width / logo_image.height  # original aspect ratio

    # Desired logo width in the PDF (you can adjust this size)
    desired_logo_width = 4 * cm
    desired_logo_height = (
        desired_logo_width / logo_aspect_ratio
    )  # maintain aspect ratio

    # Draw header text
    c.setFont("Helvetica-Bold", 16)
    c.drawString(2 * cm, page_height - 2 * cm, "Paul Scherrer Institut")

    # Draw the Heidi logo with preserved aspect ratio
    c.drawImage(
        png_logo_path,
        page_width - desired_logo_width - 2 * cm,
        page_height - desired_logo_height - 2 * cm,
        width=desired_logo_width,
        height=desired_logo_height,
        mask="auto",
    )

    # Draw details section
    c.setFont("Helvetica", 12)

    y_position = (
        page_height - 4 * cm
    )  # Adjusted to ensure text doesn't overlap with the logo
    line_height = 0.8 * cm

    if dewar.shipment:
        c.drawString(
            2 * cm, y_position, f"Shipment Name: {dewar.shipment.shipment_name}"
        )
        y_position -= line_height

        c.drawString(2 * cm, y_position, f"Dewar Name: {dewar.dewar_name}")
        y_position -= line_height

    c.drawString(2 * cm, y_position, f"Unique ID: {dewar.unique_id}")
    y_position -= line_height

    if dewar.contact_person:
        contact_person = dewar.contact_person
        c.drawString(
            2 * cm,
            y_position,
            f"Contact: {contact_person.firstname} {contact_person.lastname}",
        )
        y_position -= line_height
        c.drawString(2 * cm, y_position, f"Email: {contact_person.email}")
        y_position -= line_height
        c.drawString(2 * cm, y_position, f"Phone: {contact_person.phone_number}")
        y_position -= line_height

    if dewar.return_address:
        return_address = dewar.return_address
        c.drawString(2 * cm, y_position, f"Return Address: {return_address.street}")
        y_position -= line_height
        c.drawString(2 * cm, y_position, f"City: {return_address.city}")
        y_position -= line_height
        c.drawString(2 * cm, y_position, f"Postal Code: {return_address.zipcode}")
        y_position -= line_height
        c.drawString(2 * cm, y_position, f"Country: {return_address.country}")
        y_position -= line_height

    c.drawString(2 * cm, y_position, "Beamtime Information: Placeholder")

    # Generate QR code
    qr = qrcode.QRCode(version=1, box_size=10, border=4)
    qr.add_data(dewar.unique_id)
    qr.make(fit=True)
    qr_img = qr.make_image(fill="black", back_color="white").convert("RGBA")

    # Save this QR code to a temporary file
    with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as temp_file:
        qr_img.save(temp_file, format="PNG")
        temp_file_path = temp_file.name

    # Add QR code to PDF
    c.drawImage(
        temp_file_path, page_width - 6 * cm, 5 * cm, width=4 * cm, height=4 * cm
    )

    # Add footer text
    c.setFont("Helvetica", 10)
    c.drawCentredString(page_width / 2, 2 * cm, "Scan for more information")

    # Draw border
    c.setLineWidth(1)
    c.rect(
        1 * cm, 1 * cm, page_width - 2 * cm, page_height - 2 * cm
    )  # Adjusted dimensions

    # Finalize the canvas
    c.showPage()
    c.save()

    buffer.seek(0)

    # Cleanup temporary file
    os.remove(temp_file_path)

    return buffer


@router.get("/{dewar_id}/download-label", response_class=Response)
async def download_dewar_label(dewar_id: int, db: Session = Depends(get_db)):
    dewar = (
        db.query(DewarModel)
        .options(
            joinedload(DewarModel.pucks).joinedload(PuckModel.samples),
            joinedload(DewarModel.contact_person),
            joinedload(DewarModel.return_address),
            joinedload(DewarModel.shipment),
        )
        .filter(DewarModel.id == dewar_id)
        .first()
    )

    if not dewar:
        raise HTTPException(status_code=404, detail="Dewar not found")
    if not dewar.unique_id:
        raise HTTPException(
            status_code=404, detail="QR Code not generated for this dewar"
        )

    buffer = generate_label(dewar)

    return Response(
        buffer.getvalue(),
        media_type="application/pdf",
        headers={
            "Content-Disposition": f"attachment; filename=dewar_label_{dewar.id}.pdf"
        },
    )


@router.get("/dewars/{dewar_id}/samples", response_model=dict)
async def get_dewar_samples(dewar_id: int, db: Session = Depends(get_db)):
    # Fetch Dewar, associated Pucks, and Samples
    dewar = db.query(DewarModel).filter(DewarModel.id == dewar_id).first()
    if not dewar:
        raise HTTPException(status_code=404, detail="Dewar not found")

    pucks = db.query(PuckModel).filter(PuckModel.dewar_id == dewar.id).all()

    data = {"dewar": {"id": dewar.id, "dewar_name": dewar.dewar_name}, "pucks": []}

    for puck in pucks:
        samples = db.query(SampleModel).filter(SampleModel.puck_id == puck.id).all()
        data["pucks"].append(
            {
                "id": puck.id,
                "name": puck.puck_name,
                "type": puck.puck_type,
                "samples": [
                    {
                        "id": sample.id,
                        "position": sample.position,
                        "dewar_name": dewar.dewar_name,
                        "sample_name": sample.sample_name,
                        "priority": sample.priority,
                        "comments": sample.comments,
                        "proteinname": sample.proteinname,
                        **(sample.data_collection_parameters or {}),
                    }
                    for sample in samples
                ],
            }
        )

    return data


@router.get("/", response_model=List[DewarSchema])
async def get_dewars(db: Session = Depends(get_db)):
    try:
        dewars = db.query(DewarModel).options(joinedload(DewarModel.pucks)).all()
        return dewars
    except SQLAlchemyError as e:
        logging.error(f"Database error occurred: {e}")
        raise HTTPException(status_code=500, detail="Internal server error")


@router.get("/dewar-types", response_model=List[DewarTypeSchema])
def get_dewar_types(db: Session = Depends(get_db)):
    return db.query(DewarTypeModel).all()


@router.get(
    "/dewar-types/{type_id}/serial-numbers",
    response_model=List[DewarSerialNumberSchema],
)
def get_serial_numbers(type_id: int, db: Session = Depends(get_db)):
    return (
        db.query(DewarSerialNumberModel)
        .filter(DewarSerialNumberModel.dewar_type_id == type_id)
        .all()
    )


@router.post("/dewar-types", response_model=DewarTypeSchema)
def create_dewar_type(dewar_type: DewarTypeCreate, db: Session = Depends(get_db)):
    db_type = DewarTypeModel(**dewar_type.dict())
    db.add(db_type)
    db.commit()
    db.refresh(db_type)
    return db_type


@router.post("/dewar-serial-numbers", response_model=DewarSerialNumberSchema)
def create_dewar_serial_number(
    serial_number: DewarSerialNumberCreate, db: Session = Depends(get_db)
):
    db_serial = DewarSerialNumberModel(**serial_number.dict())
    db.add(db_serial)
    db.commit()
    db.refresh(db_serial)
    return db_serial


@router.get("/dewar-serial-numbers", response_model=List[DewarSerialNumberSchema])
def get_all_serial_numbers(db: Session = Depends(get_db)):
    try:
        serial_numbers = db.query(DewarSerialNumberModel).all()
        return serial_numbers
    except SQLAlchemyError as e:
        logging.error(f"Database error occurred: {e}")
        raise HTTPException(status_code=500, detail="Internal server error")


@router.get("/{dewar_id}", response_model=DewarSchema)
async def get_dewar(dewar_id: int, db: Session = Depends(get_db)):
    dewar = (
        db.query(DewarModel)
        .options(
            joinedload(DewarModel.pucks).joinedload(PuckModel.samples),
            joinedload(DewarModel.contact_person),
            joinedload(DewarModel.return_address),
            joinedload(DewarModel.shipment),
        )
        .filter(DewarModel.id == dewar_id)
        .first()
    )

    if not dewar:
        raise HTTPException(status_code=404, detail="Dewar not found")

    return DewarSchema.from_orm(dewar)


@router.put("/{dewar_id}", response_model=DewarSchema)
async def update_dewar(
    dewar_id: int, dewar_update: DewarUpdate, db: Session = Depends(get_db)
) -> DewarSchema:
    dewar = db.query(DewarModel).filter(DewarModel.id == dewar_id).first()

    if not dewar:
        raise HTTPException(status_code=404, detail="Dewar not found")

    for key, value in dewar_update.dict(exclude_unset=True).items():
        if hasattr(dewar, key):
            setattr(dewar, key, value)

    db.commit()
    db.refresh(dewar)
    return dewar


@router.delete("/{dewar_id}", status_code=status.HTTP_204_NO_CONTENT)
async def delete_dewar(dewar_id: int, db: Session = Depends(get_db)):
    dewar = db.query(DewarModel).filter(DewarModel.id == dewar_id).first()

    if not dewar:
        raise HTTPException(status_code=404, detail="Dewar not found")

    db.delete(dewar)
    db.commit()
    return


# New routes for shipments
@router.get("/shipments", response_model=List[ShipmentSchema])
async def get_all_shipments(db: Session = Depends(get_db)):
    try:
        shipments = get_shipments(db)
        return shipments
    except SQLAlchemyError as e:
        logging.error(f"Database error occurred: {e}")
        raise HTTPException(status_code=500, detail="Internal server error")


@router.get("/shipments/{id}", response_model=ShipmentSchema)
async def get_single_shipment(id: int, db: Session = Depends(get_db)):
    try:
        shipment = get_shipment_by_id(db, id)
        if shipment is None:
            raise HTTPException(status_code=404, detail="Shipment not found")
        return shipment
    except SQLAlchemyError as e:
        logging.error(f"Database error occurred: {e}")
        raise HTTPException(status_code=500, detail="Internal server error")