giga_tcg/services/storage.py

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