from sqlalchemy.orm import Session from db.utils import db_transaction from db.models import Warehouse, User, StagedFileProduct, StorageBlock, ProductBlock, File, Card from uuid import uuid4 as uuid from typing import List, TypedDict class ProductAttributes(TypedDict): product_id: str card_number: str class StorageService: def __init__(self, db: Session): self.db = db def get_or_create_user(self, username: str) -> User: user = self.db.query(User).filter(User.username == username).first() if user is None: user = User( id = str(uuid()), username = username ) with db_transaction(self.db): self.db.add(user) return user def get_or_create_warehouse(self) -> Warehouse: warehouse = self.db.query(Warehouse).first() user = self.get_or_create_user('admin') if warehouse is None: warehouse = Warehouse( id = str(uuid()), user_id = user.id ) with db_transaction(self.db): self.db.add(warehouse) return warehouse def get_staged_product(self, file_id: str) -> List[StagedFileProduct]: staged_product = self.db.query(StagedFileProduct).filter(StagedFileProduct.file_id == file_id).all() return staged_product def get_storage_block_name(self, warehouse: Warehouse, file_id: str) -> str: # Get file type from id current_file = self.db.query(File).filter(File.id == file_id).first() if not current_file: raise ValueError(f"No file found with id {file_id}") # Determine storage block type storage_block_type = 'rare' if 'rare' in current_file.type else 'common' prefix = storage_block_type[0] # Get most recent storage block with same type and warehouse id latest_block = ( self.db.query(StorageBlock) .filter( StorageBlock.warehouse_id == warehouse.id, StorageBlock.type == storage_block_type ) .order_by(StorageBlock.date_created.desc()) .first() ) # If no existing block, start with number 1 if not latest_block: return f"{prefix}1" # Start with the next number after the latest block number = int(latest_block.name[1:]) while True: number += 1 new_name = f"{prefix}{number}" # Check if the new name already exists exists = ( self.db.query(StorageBlock) .filter( StorageBlock.warehouse_id == warehouse.id, StorageBlock.name == new_name ) .first() ) if not exists: return new_name def create_storage_block(self, warehouse: Warehouse, file_id: str) -> StorageBlock: current_file = self.db.query(File).filter(File.id == file_id).first() if not current_file: raise ValueError(f"No file found with id {file_id}") storage_block_type = 'rare' if 'rare' in current_file.type else 'common' storage_block = StorageBlock( id = str(uuid()), warehouse_id = warehouse.id, name = self.get_storage_block_name(warehouse, file_id), type = storage_block_type ) with db_transaction(self.db): self.db.add(storage_block) return storage_block def add_staged_product_to_product_block( self, staged_product: StagedFileProduct, storage_block: StorageBlock, product_attributes: ProductAttributes, block_index: int ) -> ProductBlock: """Create a new ProductBlock for a single unit of a staged product.""" product_block = ProductBlock( id=str(uuid()), product_id=staged_product.product_id, block_id=storage_block.id, block_index=block_index ) with db_transaction(self.db): self.db.add(product_block) return product_block def get_staged_product_attributes_for_storage( self, staged_product: StagedFileProduct ) -> List[ProductAttributes]: """Get attributes for each unit of a staged product.""" result = ( self.db.query( StagedFileProduct.product_id, StagedFileProduct.quantity, Card.number ) .join(Card, Card.product_id == StagedFileProduct.product_id) .filter(StagedFileProduct.id == staged_product.id) .first() ) if not result: return [] return [ ProductAttributes( product_id=result.product_id, card_number=result.number ) for i in range(result.quantity) ] def store_staged_products_for_file(self, file_id: str) -> StorageBlock: """Store all staged products for a file in a new storage block.""" warehouse = self.get_or_create_warehouse() storage_block = self.create_storage_block(warehouse, file_id) staged_products = self.get_staged_product(file_id) # Collect all product attributes first all_product_attributes = [] for staged_product in staged_products: product_attributes_list = self.get_staged_product_attributes_for_storage(staged_product) for attrs in product_attributes_list: all_product_attributes.append((staged_product, attrs)) # Sort by card number as integer to determine block indices sorted_attributes = sorted(all_product_attributes, key=lambda x: int(x[1]['card_number'])) # Add products with correct block indices for block_index, (staged_product, product_attributes) in enumerate(sorted_attributes, 1): self.add_staged_product_to_product_block( staged_product=staged_product, storage_block=storage_block, product_attributes=product_attributes, block_index=block_index ) return storage_block