from sqlalchemy.orm import Session
from db.models import Product, Inventory
from schemas.inventory import UpdateInventoryResponse
from db.utils import db_transaction

class InventoryService:
    def __init__(self, db: Session):
        self.db = db
    
    def add_inventory(self, product: Product, quantity: int) -> Inventory:
        inventory = self.db.query(Inventory).filter(Inventory.product_id == product.id).first()
        if inventory is None:
            inventory = Inventory(product_id=product.id, quantity=quantity)
            self.db.add(inventory)
        else:
            inventory.quantity += quantity
        return inventory

    def process_staged_products(self, product_data: dict[Product, int]) -> UpdateInventoryResponse:
        with db_transaction(self.db):
            for product, quantity in product_data.items():
                self.add_inventory(product, quantity)
        return UpdateInventoryResponse(success=True)
    
    def add_sealed_box_to_inventory(self, product: Product, quantity: int) -> UpdateInventoryResponse:
        with db_transaction(self.db):
            inventory = self.add_inventory(product, quantity)
        return UpdateInventoryResponse(success=True)