import uuid
from typing import Sequence

from sqlalchemy import select, update as sa_update
from sqlalchemy.ext.asyncio import AsyncSession

from app.models.imp_batch import ImpBatch
from app.models.imp_product import ImpProduct
from app.schemas.imp import BatchCreate, BatchUpdate


async def create_batch(db: AsyncSession, study_id: uuid.UUID, batch_in: BatchCreate) -> ImpBatch:
    result = await db.execute(select(ImpProduct).where(ImpProduct.id == batch_in.product_id))
    product = result.scalar_one_or_none()
    if not product or product.study_id != study_id:
        raise ValueError("Product not found in study")

    batch = ImpBatch(
        study_id=study_id,
        product_id=batch_in.product_id,
        batch_no=batch_in.batch_no,
        expiry_date=batch_in.expiry_date,
        manufacture_date=batch_in.manufacture_date,
        status="ACTIVE",
    )
    db.add(batch)
    await db.commit()
    await db.refresh(batch)
    return batch


async def get_batch(db: AsyncSession, batch_id: uuid.UUID) -> ImpBatch | None:
    result = await db.execute(select(ImpBatch).where(ImpBatch.id == batch_id))
    return result.scalar_one_or_none()


async def list_batches(db: AsyncSession, study_id: uuid.UUID, product_id: uuid.UUID | None = None, status: str | None = None) -> Sequence[ImpBatch]:
    stmt = select(ImpBatch).where(ImpBatch.study_id == study_id)
    if product_id:
        stmt = stmt.where(ImpBatch.product_id == product_id)
    if status:
        stmt = stmt.where(ImpBatch.status == status)
    result = await db.execute(stmt)
    return result.scalars().all()


async def update_batch(db: AsyncSession, batch: ImpBatch, batch_in: BatchUpdate) -> ImpBatch:
    update_data = batch_in.model_dump(exclude_unset=True)
    if update_data:
        await db.execute(
            sa_update(ImpBatch)
            .where(ImpBatch.id == batch.id)
            .values(**update_data)
        )
        await db.commit()
    await db.refresh(batch)
    return batch