aaredb/backend/app/schemas.py

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


class loginToken(BaseModel):
    access_token: str
    token_type: str


class loginData(BaseModel):
    username: str
    pgroups: List[int]


class DewarTypeBase(BaseModel):
    dewar_type: str


class DewarTypeCreate(DewarTypeBase):
    pass


class DewarType(DewarTypeBase):
    id: int

    class Config:
        from_attributes = True


class DewarSerialNumberBase(BaseModel):
    serial_number: str
    dewar_type_id: int


class DewarSerialNumberCreate(DewarSerialNumberBase):
    pass


class DewarSerialNumber(DewarSerialNumberBase):
    id: int
    dewar_type: DewarType

    class Config:
        from_attributes = True


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 SampleEventCreate(BaseModel):
    event_type: str


class SampleEventResponse(BaseModel):
    id: int
    sample_id: int
    event_type: str
    timestamp: datetime

    class Config:
        from_attributes = True


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


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: Optional[str] = None
    lastname: Optional[str] = None
    phone_number: Optional[str] = None
    email: Optional[EmailStr] = None


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: Optional[str] = None
    city: Optional[str] = None
    zipcode: Optional[str] = None
    country: Optional[str] = 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)
    events: List[SampleEventCreate] = []


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

    class Config:
        populate_by_name = True


class PuckEvent(BaseModel):
    id: int
    puck_id: int
    tell_position: Optional[str] = None
    event_type: str
    timestamp: datetime

    class Config:
        from_attributes = True


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
    events: List[PuckEvent] = []
    samples: List[Sample] = []  # List of samples within this puck

    class Config:
        from_attributes = True


class DewarBase(BaseModel):
    dewar_name: str
    dewar_type_id: Optional[int] = None
    dewar_serial_number_id: Optional[int] = None
    unique_id: Optional[str] = None
    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]
    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
    dewar_type_id: Optional[int] = None
    dewar_serial_number_id: Optional[int] = None
    unique_id: 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
    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


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

    class Config:
        from_attributes = True


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


class LogisticsEventCreate(BaseModel):
    dewar_qr_code: str
    location_qr_code: str
    transaction_type: str


class SlotSchema(BaseModel):
    id: int
    qr_code: str
    label: str
    qr_base: Optional[str]
    occupied: bool
    needs_refill: bool
    dewar_unique_id: Optional[str]
    dewar_name: Optional[str]
    time_until_refill: Optional[int]
    at_beamline: Optional[bool]
    retrievedTimestamp: Optional[str]
    beamlineLocation: Optional[str]
    shipment_name: Optional[str]
    contact_person: Optional[str]
    local_contact: Optional[str]

    class Config:
        from_attributes = True


class SetTellPosition(BaseModel):
    puckname: str  # The puck name is required.
    segment: Optional[str] = Field(
        None,
        pattern="^[A-F]$",  # Valid segments are A, B, C, D, E, F
        description="Segment must be one of A, B, C, D, E, or F."
        "Can be null for no tell_position.",
    )
    puck_in_segment: Optional[int] = Field(
        None,
        ge=1,
        le=5,
        description="Puck in segment must be between 1 and 5."
        "Can be null for no tell_position.",
    )

    @property
    def tell_position(self) -> Optional[str]:
        """
        Combines `segment` and `puck_in_segment` to generate the `tell_position`.
        If either value is `None`, returns `None` to indicate no `tell_position`.
        """
        if self.segment and self.puck_in_segment:
            return f"{self.segment}{self.puck_in_segment}"
        return None