aaredb/backend/app/schemas.py

from typing import List, Optional
from pydantic import BaseModel, EmailStr, constr, Field
from datetime import date


class DataCollectionParameters(BaseModel):
    priority: Optional[int] = None
    comments: Optional[str] = None
    directory: Optional[str] = None
    proteinname: Optional[str] = None
    oscillation: Optional[float] = None
    aperture: Optional[str] = None
    exposure: Optional[float] = None
    totalrange: Optional[int] = None
    transmission: Optional[int] = None
    dose: Optional[float] = None
    targetresolution: Optional[float] = None
    datacollectiontype: Optional[str] = None
    processingpipeline: Optional[str] = None
    spacegroupnumber: Optional[int] = None
    cellparameters: Optional[str] = None
    rescutkey: Optional[str] = None
    rescutvalue: Optional[float] = None
    userresolution: Optional[float] = None
    pdbid: Optional[str] = None
    autoprocfull: Optional[bool] = None
    procfull: Optional[bool] = None
    adpenabled: Optional[bool] = None
    noano: Optional[bool] = None
    ffcscampaign: Optional[bool] = None
    trustedhigh: Optional[float] = None
    autoprocextraparams: Optional[str] = None
    chiphiangles: Optional[float] = None

    class Config:
        from_attributes = True


class Results(BaseModel):
    # Define attributes for Results here
    pass


# Contact Person schemas
class ContactPersonBase(BaseModel):
    firstname: str
    lastname: str
    phone_number: str
    email: EmailStr


class ContactPersonCreate(ContactPersonBase):
    pass


class ContactPerson(ContactPersonBase):
    id: int

    class Config:
        from_attributes = True


class ContactPersonUpdate(BaseModel):
    firstname: str | None = None
    lastname: str | None = None
    phone_number: str | None = None
    email: EmailStr | None = None


# Address schemas
class AddressCreate(BaseModel):
    street: str
    city: str
    zipcode: str
    country: str


class Address(AddressCreate):
    id: int

    class Config:
        from_attributes = True


class AddressUpdate(BaseModel):
    street: str | None = None
    city: str | None = None
    zipcode: str | None = None
    country: str | None = None


class Sample(BaseModel):
    id: int
    sample_name: str
    position: int  # Position within the puck
    puck_id: int
    crystalname: Optional[str] = Field(None)
    positioninpuck: Optional[int] = Field(None)


class SampleCreate(BaseModel):
    sample_name: str = Field(..., alias="crystalname")
    position: int = Field(..., alias="positioninpuck")
    data_collection_parameters: Optional[DataCollectionParameters] = None
    results: Optional[Results] = None

    class Config:
        populate_by_name = True



# Puck schemas
class PuckBase(BaseModel):
    puck_name: str
    puck_type: str
    puck_location_in_dewar: int


class PuckCreate(BaseModel):
    puck_name: str
    puck_type: str
    puck_location_in_dewar: int
    samples: List[SampleCreate] = []


class PuckUpdate(BaseModel):
    puck_name: Optional[str] = None
    puck_type: Optional[str] = None
    puck_location_in_dewar: Optional[int] = None
    dewar_id: Optional[int] = None


class Puck(BaseModel):
    id: int
    puck_name: str
    puck_type: str
    puck_location_in_dewar: int
    dewar_id: int
    samples: List[Sample] = []  # List of samples within this puck

    class Config:
        from_attributes = True


# Dewar schemas
class DewarBase(BaseModel):
    dewar_name: str
    tracking_number: str
    number_of_pucks: int
    number_of_samples: int
    status: str
    ready_date: Optional[date]
    shipping_date: Optional[date]
    arrival_date: Optional[date]
    returning_date: Optional[date]
    qrcode: str
    contact_person_id: Optional[int]
    return_address_id: Optional[int]
    pucks: List[PuckCreate] = []


class DewarCreate(DewarBase):
    pass


class Dewar(DewarBase):
    id: int
    shipment_id: Optional[int]
    contact_person: Optional[ContactPerson]
    return_address: Optional[Address]
    pucks: List[Puck] = []  # List of pucks within this dewar

    class Config:
        from_attributes = True


class DewarUpdate(BaseModel):
    dewar_name: Optional[str] = None
    tracking_number: Optional[str] = None
    status: Optional[str] = None
    ready_date: Optional[date] = None
    shipping_date: Optional[date] = None
    arrival_date: Optional[date] = None
    returning_date: Optional[date] = None
    qrcode: Optional[str] = None
    contact_person_id: Optional[int] = None
    address_id: Optional[int] = None

class DewarSchema(BaseModel):
    id: int
    dewar_name: str
    tracking_number: str
    status: str
    contact_person_id: int
    return_address_id: int

    class Config:
        from_attributes = True


# Proposal schemas
class Proposal(BaseModel):
    id: int
    number: str

    class Config:
        from_attributes = True


# Shipment schemas
class Shipment(BaseModel):
    id: int
    shipment_name: str
    shipment_date: date
    shipment_status: str
    comments: Optional[str]
    contact_person: Optional[ContactPerson]
    return_address: Optional[Address]
    proposal: Optional[Proposal]
    dewars: List[Dewar] = []

    class Config:
        from_attributes = True


class ShipmentCreate(BaseModel):
    shipment_name: str
    shipment_date: date
    shipment_status: str
    comments: Optional[constr(max_length=200)]
    contact_person_id: int
    return_address_id: int
    proposal_id: int
    dewars: List[DewarCreate] = []

    class Config:
        from_attributes = True


class UpdateShipmentComments(BaseModel):
    comments: str