# app/routers/shipment.py

from fastapi import APIRouter, HTTPException, status, Query, Depends
from sqlalchemy.orm import Session
from typing import List, Optional
import logging
from pydantic import BaseModel

from app.models import Shipment as ShipmentModel, ContactPerson as ContactPersonModel, Address as AddressModel, \
    Proposal as ProposalModel, Dewar as DewarModel
from app.schemas import ShipmentCreate, UpdateShipmentComments, Shipment as ShipmentSchema, DewarUpdate, \
    ContactPerson as ContactPersonSchema, Sample as SampleSchema, DewarCreate, PuckCreate, SampleCreate
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(id: Optional[int] = Query(None), db: Session = Depends(get_db)):
    if id:
        shipment = get_shipment_by_id(db, id)
        if not shipment:
            logging.error(f"Shipment with ID {id} not found")
            raise HTTPException(status_code=404, detail="Shipment not found")
        logging.info(f"Shipment found: {shipment}")
        return [shipment]

    shipments = get_shipments(db)
    logging.info(f"Total shipments fetched: {len(shipments)}")
    for shipment in shipments:
        logging.info(f"Shipment ID: {shipment.id}, Shipment Name: {shipment.shipment_name}")
    return shipments


@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")

    db_shipment = ShipmentModel(
        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(id: int, db: Session = Depends(get_db)):
    shipment = db.query(ShipmentModel).filter(ShipmentModel.id == 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(id: int, 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.id == 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")

        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(id: int, dewar_id: int, db: Session = Depends(get_db)):
    shipment = db.query(ShipmentModel).filter(ShipmentModel.id == 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: int, dewar_id: int, db: Session = Depends(get_db)):
    shipment = db.query(ShipmentModel).filter(ShipmentModel.id == shipment_id).first()
    if not shipment:
        raise HTTPException(status_code=404, detail="Shipment not found")

    dewar_exists = any(dw.id == dewar_id for dw in shipment.dewars)
    if not dewar_exists:
        raise HTTPException(status_code=404, detail=f"Dewar with ID {dewar_id} not found in shipment")

    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(id: int, db: Session = Depends(get_db)):
    shipment = db.query(ShipmentModel).filter(ShipmentModel.id == 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.samples)

    return samples


@router.get("/shipments/{shipment_id}/dewars/{dewar_id}/samples", response_model=List[SampleSchema])
def get_samples_in_dewar(shipment_id: int, dewar_id: int, db: Session = Depends(get_db)):
    shipment = get_shipment_by_id(db, shipment_id)
    if not shipment:
        raise HTTPException(status_code=404, detail="Shipment not found")

    dewar = next((d for d in shipment.dewars if d.id == dewar_id), None)
    if not dewar:
        raise HTTPException(status_code=404, detail="Dewar not found")

    samples = []
    for puck in dewar.pucks:
        for sample in puck.samples:
            samples.append(sample)

    return samples


@router.put("/{shipment_id}/comments", response_model=ShipmentSchema)
async def update_shipment_comments(id: int, comments_data: UpdateShipmentComments, db: Session = Depends(get_db)):
    shipment = db.query(ShipmentModel).filter(ShipmentModel.id == id).first()
    if not shipment:
        raise HTTPException(status_code=404, detail="Shipment not found")

    shipment.comments = comments_data.comments
    db.commit()
    db.refresh(shipment)
    return shipment


@router.post("/{shipment_id}/add_dewar_puck_sample", response_model=ShipmentSchema, status_code=status.HTTP_201_CREATED)
async def add_dewar_puck_sample_to_shipment(shipment_id: int, payload: DewarCreate, db: Session = Depends(get_db)):
    shipment = db.query(ShipmentModel).filter(ShipmentModel.id == shipment_id).first()
    if not shipment:
        raise HTTPException(status_code=404, detail="Shipment not found")

    for dewar_data in payload.dewars:
        dewar = Dewar(
            shipment_id=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,
        )
        db.add(dewar)
        db.commit()
        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,
            )
            db.add(puck)
            db.commit()
            db.refresh(puck)

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

    db.refresh(shipment)
    return shipment