# Adjusting the ShipmentProcessor for better error handling and alignment

from sqlalchemy.orm import Session
from app.models import Shipment, Dewar, Puck, Sample, DataCollectionParameters
from app.schemas import ShipmentCreate, ShipmentResponse
import logging


class ShipmentProcessor:
    def __init__(self, db: Session):
        self.db = db

    def process_shipment(self, shipment: ShipmentCreate) -> ShipmentResponse:
        logger = logging.getLogger(__name__)
        try:
            new_shipment = Shipment(
                shipment_name=shipment.shipment_name,
                shipment_date=shipment.shipment_date,
                shipment_status=shipment.shipment_status,
            )
            self.db.add(new_shipment)
            self.db.commit()
            self.db.refresh(new_shipment)

            for dewar_data in shipment.dewars:
                dewar = Dewar(
                    shipment_id=new_shipment.id,
                    dewar_name=dewar_data.dewar_name,
                    tracking_number=dewar_data.tracking_number,
                    status=dewar_data.status,
                    contact_person_id=dewar_data.contact_person_id,
                    return_address_id=dewar_data.return_address_id,
                )
                self.db.add(dewar)
                self.db.commit()
                self.db.refresh(dewar)

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

                    for sample_data in puck_data.samples:
                        data_collection_params = DataCollectionParameters(
                            **sample_data.data_collection_parameters.dict(by_alias=True))
                        sample = Sample(
                            puck_id=puck.id,
                            sample_name=sample_data.sample_name,
                            position=sample_data.position,
                            data_collection_parameters=data_collection_params
                        )
                        self.db.add(sample)
                        self.db.commit()
                        self.db.refresh(sample)

            return ShipmentResponse(
                shipment_id=new_shipment.id,
                status="success",
                message="Shipment processed successfully"
            )
        except Exception as e:
            logger.error(f"Error processing shipment: {str(e)}")
            raise e