from utils.db_connectors.base_connector import BaseConnector
from utils.db_objects.direct_delivery import Asset, AssetLocation, \
    AssetLocationDetail


class AssetsConnector(BaseConnector):
    def remove_assets_if_exists(self, upc, asset_type_id=None):
        asset_ids = self.get_asset_id(upc, asset_type_id)
        if asset_ids:
            for asset in asset_ids:
                self.delete_asset_records(asset.asset_id)

    def get_asset_id(self, upc, asset_type_id=None):
        session = self.Session()
        filter_by_params = {'upc': upc}
        if asset_type_id:
            filter_by_params['asset_type_id'] = str(asset_type_id)
        results = session.query(Asset).filter_by(**filter_by_params)
        session.close()
        if results is None:
            return None
        else:
            return results

    def delete_asset_records(self, asset_id):
        session = self.Session()
        location_rows = session.query(AssetLocation).filter_by(
            asset_id=asset_id)
        for row in location_rows:
            location_id = row.asset_location_id
            session.query(AssetLocationDetail).filter_by(
                asset_location_id=location_id).delete()
            session.delete(row)
        session.query(Asset).filter_by(asset_id=asset_id).delete()
        session.commit()
        session.close()