""" Repository for Sample entity operations. """ from datetime import UTC, datetime from uuid import UUID from sqlalchemy import ColumnElement, Numeric, distinct, func, or_, select from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import selectinload from app.models.order import Sample, SampleStatus, VisitOrder from app.models.report import Report, ReportStatus class SampleRepository: """Repository handling database operations for Sample.""" async def get_next_sample_code(self, db: AsyncSession, lab_id: UUID) -> str: """ Generate year-aware sequential sample code for a lab. Format: S{YY}{6-digit-seq} e.g. S26000001, S26000009 """ from sqlalchemy import text year_short = datetime.now(UTC).strftime("%y") prefix = f"S{year_short}" # Barcodes are globally unique, so allocation must span branches. The # transaction lock serializes concurrent allocations until sample commit. await db.execute(text("SELECT pg_advisory_xact_lock(:key)"), {"key": 730000 + int(year_short)}) numeric_suffix = func.substring(Sample.sample_code, len(prefix) + 1) stmt = select(func.max(func.cast(numeric_suffix, Numeric(50, 0)))).where( Sample.sample_code.op("~")(f"^{prefix}[0-9]+$"), ) max_seq = int((await db.execute(stmt)).scalar() or 0) return f"{prefix}{max_seq + 1:06d}" async def get_next_barcode(self, db: AsyncSession) -> str: """ Generate date-aware scanner barcode. Format: SAM-{YYYYMMDD}-{6-digit-seq} e.g. SAM-20260921-000001 """ date_str = datetime.now(UTC).strftime("%Y%m%d") prefix = f"SAM-{date_str}-" stmt = ( select(Sample.barcode) .where(Sample.barcode.like(f"{prefix}%")) .order_by(Sample.barcode.desc()) .limit(1) ) res = await db.execute(stmt) last_barcode = res.scalar_one_or_none() next_seq = 1 if last_barcode: try: parts = last_barcode.split("-") if len(parts) == 3: next_seq = int(parts[2]) + 1 except (ValueError, IndexError): next_seq = 1 return f"{prefix}{next_seq:06d}" async def get_by_id(self, db: AsyncSession, sample_id: UUID) -> Sample | None: """Fetch sample by ID with loaded order, patient, and collector relations.""" stmt = ( select(Sample) .where(Sample.id == sample_id) .options( selectinload(Sample.collected_by_user), selectinload(Sample.order).selectinload(VisitOrder.patient), selectinload(Sample.order).selectinload(VisitOrder.order_tests), ) ) res = await db.execute(stmt) return res.scalar_one_or_none() async def get_by_barcode( self, db: AsyncSession, barcode: str, lab_id: UUID | None = None ) -> Sample | None: """Fetch sample by scanner-readable barcode or sample code.""" cleaned = barcode.strip() stmt = ( select(Sample) .where(or_(Sample.barcode == cleaned, Sample.sample_code == cleaned)) .options( selectinload(Sample.collected_by_user), selectinload(Sample.order).selectinload(VisitOrder.patient), selectinload(Sample.order).selectinload(VisitOrder.order_tests), ) ) if lab_id is not None: stmt = stmt.where(Sample.lab_id == lab_id) res = await db.execute(stmt.limit(2)) matches = res.scalars().all() return matches[0] if len(matches) == 1 else None async def get_by_order_id(self, db: AsyncSession, order_id: UUID) -> list[Sample]: """Fetch all samples associated with an order.""" stmt = ( select(Sample) .where(Sample.order_id == order_id) .options( selectinload(Sample.collected_by_user), selectinload(Sample.order).selectinload(VisitOrder.patient), selectinload(Sample.order).selectinload(VisitOrder.order_tests), ) .order_by(Sample.created_at.asc()) ) res = await db.execute(stmt) return list(res.scalars().all()) async def list( self, db: AsyncSession, lab_id: UUID | None = None, status: SampleStatus | str | None = None, sample_type: str | None = None, order_id: UUID | str | None = None, collected_by: UUID | None = None, offset: int = 0, limit: int = 20, ) -> list[Sample]: """List samples matching filters with pagination.""" stmt = ( select(Sample) .options( selectinload(Sample.collected_by_user), selectinload(Sample.order).selectinload(VisitOrder.patient), selectinload(Sample.order).selectinload(VisitOrder.order_tests), selectinload(Sample.order).selectinload(VisitOrder.reports).selectinload(Report.verifier), ) .order_by(Sample.created_at.desc()) ) conditions: list[ColumnElement[bool]] = [] if lab_id: conditions.append(Sample.lab_id == lab_id) if status: st = str(status).strip().lower() if st == "rejected": latest_reports_subq = ( select(Report.order_id, func.max(Report.version).label("max_version")) .where(Report.deleted_at.is_(None)) .group_by(Report.order_id) .subquery() ) stmt = ( stmt.join(Sample.order) .outerjoin( latest_reports_subq, VisitOrder.id == latest_reports_subq.c.order_id, ) .outerjoin( Report, (Report.order_id == latest_reports_subq.c.order_id) & (Report.version == latest_reports_subq.c.max_version), ) ) conditions.append( or_( Report.status == ReportStatus.REJECTED, Sample.notes.ilike("%[REJECTED%"), ) ) else: try: parsed_status = SampleStatus(st) conditions.append(Sample.status == parsed_status) except ValueError: pass if sample_type: conditions.append(Sample.sample_type.ilike(f"%{sample_type}%")) if order_id: try: parsed_uid = UUID(str(order_id).strip()) conditions.append(Sample.order_id == parsed_uid) except (ValueError, AttributeError): if st != "rejected": stmt = stmt.join(Sample.order) conditions.append(VisitOrder.order_number.ilike(f"%{str(order_id).strip()}%")) if collected_by: conditions.append(Sample.collected_by == collected_by) if conditions: stmt = stmt.where(*conditions) stmt = stmt.offset(offset).limit(limit) res = await db.execute(stmt) return list(res.scalars().all()) async def count( self, db: AsyncSession, lab_id: UUID | None = None, status: SampleStatus | str | None = None, sample_type: str | None = None, order_id: UUID | str | None = None, collected_by: UUID | None = None, ) -> int: """Count total samples matching filters.""" stmt = select(func.count(distinct(Sample.id))).select_from(Sample) conditions: list[ColumnElement[bool]] = [] if lab_id: conditions.append(Sample.lab_id == lab_id) if status: st = str(status).strip().lower() if st == "rejected": latest_reports_subq = ( select(Report.order_id, func.max(Report.version).label("max_version")) .where(Report.deleted_at.is_(None)) .group_by(Report.order_id) .subquery() ) stmt = ( stmt.join(Sample.order) .outerjoin( latest_reports_subq, VisitOrder.id == latest_reports_subq.c.order_id, ) .outerjoin( Report, (Report.order_id == latest_reports_subq.c.order_id) & (Report.version == latest_reports_subq.c.max_version), ) ) conditions.append( or_( Report.status == ReportStatus.REJECTED, Sample.notes.ilike("%[REJECTED%"), ) ) else: try: parsed_status = SampleStatus(st) conditions.append(Sample.status == parsed_status) except ValueError: pass if sample_type: conditions.append(Sample.sample_type.ilike(f"%{sample_type}%")) if order_id: try: parsed_uid = UUID(str(order_id).strip()) conditions.append(Sample.order_id == parsed_uid) except (ValueError, AttributeError): stmt = stmt.join(Sample.order) conditions.append(VisitOrder.order_number.ilike(f"%{str(order_id).strip()}%")) if collected_by: conditions.append(Sample.collected_by == collected_by) if conditions: stmt = stmt.where(*conditions) res = await db.execute(stmt) return int(res.scalar_one() or 0) async def create(self, db: AsyncSession, sample: Sample) -> Sample: """Persist a new sample record.""" db.add(sample) await db.flush() await db.refresh(sample) return sample