import uuid
from typing import Sequence

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.models.imp_inventory import ImpInventory
from app.models.imp_batch import ImpBatch


async def get_inventory_for_update(db: AsyncSession, study_id: uuid.UUID, site_id: uuid.UUID, batch_id: uuid.UUID) -> ImpInventory | None:
    result = await db.execute(
        select(ImpInventory)
        .where(
            ImpInventory.study_id == study_id,
            ImpInventory.site_id == site_id,
            ImpInventory.batch_id == batch_id,
        )
        .with_for_update()
    )
    return result.scalar_one_or_none()


async def get_or_create_inventory(db: AsyncSession, study_id: uuid.UUID, site_id: uuid.UUID, batch_id: uuid.UUID) -> ImpInventory:
    inv = await get_inventory_for_update(db, study_id, site_id, batch_id)
    if inv:
        return inv
    inv = ImpInventory(study_id=study_id, site_id=site_id, batch_id=batch_id, quantity_on_hand=0)
    db.add(inv)
    await db.flush()
    return inv


async def list_inventory(
    db: AsyncSession,
    study_id: uuid.UUID,
    site_id: uuid.UUID | None = None,
    product_id: uuid.UUID | None = None,
) -> Sequence[ImpInventory]:
    stmt = select(ImpInventory).where(ImpInventory.study_id == study_id)
    if site_id:
        stmt = stmt.where(ImpInventory.site_id == site_id)
    if product_id:
        stmt = stmt.join(ImpBatch, ImpBatch.id == ImpInventory.batch_id).where(ImpBatch.product_id == product_id)
    result = await db.execute(stmt)
    return result.scalars().all()