csci-backend/app/db/repository/vehicle.py

from sqlalchemy.orm import Session
from schemas.vehicle import CreateVehicle, UpdateVehicle
from db.models.vehicle import Vehicle
from db.base import User
from db.repository.user import get_car_driver


def create_new_vehicle(vehicle: CreateVehicle, db: Session):
    vehicle_object = Vehicle(
        **vehicle.model_dump(),
        Fuel=0,
        Status="Inactive",
        CurrentLocation=[],
        MaintenanceNotes=[]
    )
    db.add(vehicle_object)
    db.commit()
    db.refresh(vehicle_object)
    return vehicle_object


def assign_vehicle_driver(vehicle_id: int, driver_id: int, db: Session):
    driver = (
        db.query(User).filter((User.Id == driver_id) & (User.Role == "Driver")).first()
    )
    print(driver)
    vehicle = db.query(Vehicle).filter(Vehicle.Id == vehicle_id).first()
    if not vehicle:
        return "novehicle"
    if not driver:
        return "nodriver"
    driver.AssignedVehicle = vehicle_id
    db.add(driver)
    db.commit()
    return vehicle


def list_vehicles(db: Session):
    vehicles = db.query(Vehicle).all()
    # assign AssignedDriver to all vehicles based on their drivers
    for vehicle in vehicles:
        driver = get_car_driver(vehicle.Id, db)
        if driver:
            vehicle.AssignedDriver = driver.Id
        else:
            vehicle.AssignedDriver = None
    return vehicles


def get_vehicle_by_id(vehicle_id: int, db: Session):
    vehicle = db.query(Vehicle).filter(Vehicle.Id == vehicle_id).first()
    driver = get_car_driver(vehicle.Id, db)
    if driver:
        vehicle.AssignedDriver = driver.Id
    else:
        vehicle.AssignedDriver = None
    return vehicle


def replace_vehicle_data(id: int, vehicle: UpdateVehicle, db: Session):
    vehicle_db = db.query(Vehicle).filter(Vehicle.Id == id)
    vehicle_object = vehicle_db.first()
    if not vehicle_object:
        return "vehiclenotfound"
    vehicle_object.CurrentLocation = vehicle.CurrentLocation
    vehicle_object.Fuel = vehicle.Fuel
    vehicle_object.LicensePlate = vehicle.LicensePlate
    vehicle_object.MaintenanceNotes = vehicle.MaintenanceNotes
    vehicle_object.Mileage = vehicle.Mileage
    vehicle_object.Model = vehicle.Model
    vehicle_object.Status = vehicle.Status
    res = assign_vehicle_driver(id, vehicle.AssignedDriver, db)
    if res == "nodriver":
        return "driverNotFound"
    vehicle_object.Capacity = vehicle.Capacity
    vehicle_object.Year = vehicle.Year
    print(vehicle_object)
    db.add(vehicle_object)
    db.commit()
    vehicle_object.AssignedDriver = vehicle.AssignedDriver
    return vehicle_object


def delete_vehicle_data(id: int, db: Session):
    vehicle_db = db.query(Vehicle).filter(Vehicle.Id == id)
    vehicle_object = vehicle_db.first()
    if not vehicle_object:
        return "vehiclenotfound"
    db.delete(vehicle_object)
    db.commit()
    return vehicle_object